feat: improve broken cid.Builder testing for CidBuilder

If you don't use something derived from a cid.Prefix then it'll test
execute your hasher to make sure it doesn't error.

The reason for this is to avoid more cases where a panic could occur
when encoding a ProtoNode.

fix: apply code-review feedback

Co-authored-by: Masih H. Derkani <[email protected]>

Author: rvagg

merkledag_test.go

@@ -81,38 +81,64 @@ func traverseAndCheck(t *testing.T, root ipld.Node, ds ipld.DAGService, hasF fun
 	}
 }
 
+type brokenBuilder struct{}
+
+func (brokenBuilder) Sum([]byte) (cid.Cid, error)    { return cid.Undef, errors.New("Nope!") }
+func (brokenBuilder) GetCodec() uint64               { return 0 }
+func (b brokenBuilder) WithCodec(uint64) cid.Builder { return b }
+
 func TestBadBuilderEncode(t *testing.T) {
 	n := NodeWithData([]byte("boop"))
-	_, err := n.EncodeProtobuf(false)
-	if err != nil {
-		t.Fatal(err)
-	}
-	err = n.SetCidBuilder(
-		&cid.Prefix{
-			MhType:   mh.SHA2_256,
-			MhLength: -1,
-			Version:  1,
-			Codec:    cid.DagProtobuf,
-		},
-	)
-	if err != nil {
-		t.Fatal(err)
-	}
-	err = n.SetCidBuilder(
-		&cid.Prefix{
-			MhType:   mh.SHA2_256_TRUNC254_PADDED,
-			MhLength: 256,
-			Version:  1,
-			Codec:    cid.DagProtobuf,
-		},
-	)
-	if err == nil {
-		t.Fatal("expected SetCidBuilder to error on unusable hasher")
-	}
-	_, err = n.EncodeProtobuf(false)
-	if err != nil {
-		t.Fatalf("expected EncodeProtobuf to use safe CidBuilder: %v", err)
-	}
+
+	t.Run("good builder sanity check", func(t *testing.T) {
+		if _, err := n.EncodeProtobuf(false); err != nil {
+			t.Fatal(err)
+		}
+		if err := n.SetCidBuilder(
+			&cid.Prefix{
+				MhType:   mh.SHA2_256,
+				MhLength: -1,
+				Version:  1,
+				Codec:    cid.DagProtobuf,
+			},
+		); err != nil {
+			t.Fatal(err)
+		}
+	})
+
+	t.Run("hasher we can't use, should error", func(t *testing.T) {
+		if err := n.SetCidBuilder(
+			&cid.Prefix{
+				MhType:   mh.SHA2_256_TRUNC254_PADDED,
+				MhLength: 256,
+				Version:  1,
+				Codec:    cid.DagProtobuf,
+			},
+		); err == nil {
+			t.Fatal("expected SetCidBuilder to error on unusable hasher")
+		}
+		if _, err := n.EncodeProtobuf(false); err != nil {
+			t.Fatalf("expected EncodeProtobuf to use safe CidBuilder: %v", err)
+		}
+	})
+
+	t.Run("broken custom builder, should error", func(t *testing.T) {
+		if err := n.SetCidBuilder(brokenBuilder{}); err == nil {
+			t.Fatal("expected SetCidBuilder to error on unusable hasher")
+		}
+		if _, err := n.EncodeProtobuf(false); err != nil {
+			t.Fatalf("expected EncodeProtobuf to use safe CidBuilder: %v", err)
+		}
+	})
+
+	t.Run("broken custom builder as pointer, should error", func(t *testing.T) {
+		if err := n.SetCidBuilder(&brokenBuilder{}); err == nil {
+			t.Fatal("expected SetCidBuilder to error on unusable hasher")
+		}
+		if _, err := n.EncodeProtobuf(false); err != nil {
+			t.Fatalf("expected EncodeProtobuf to use safe CidBuilder: %v", err)
+		}
+	})
 }
 
 func TestLinkChecking(t *testing.T) {

node.go

@@ -25,6 +25,9 @@ var (
 	ErrLinkNotFound = fmt.Errorf("no link by that name")
 )
 
+// for testing custom CidBuilders
+var zeros [256]byte
+
 type immutableProtoNode struct {
 	encoded []byte
 	dagpb.PBNode
@@ -106,13 +109,20 @@ func (n *ProtoNode) SetCidBuilder(builder cid.Builder) error {
 		n.builder = v0CidPrefix
 		return nil
 	}
-	if p, ok := builder.(*cid.Prefix); ok {
-		mhLen := p.MhLength
-		if mhLen <= 0 {
-			mhLen = -1
+	switch b := builder.(type) {
+	case cid.Prefix:
+		if err := checkHasher(b.MhType, b.MhLength); err != nil {
+			return err
+		}
+	case *cid.Prefix:
+		if err := checkHasher(b.MhType, b.MhLength); err != nil {
+			return err
 		}
-		_, err := mhcore.GetVariableHasher(p.MhType, mhLen)
-		if err != nil {
+	default:
+		// We have to test it's a usable hasher by invoking it and checking it
+		// doesn't error. This is only a basic check, there are still ways it may
+		// break
+		if _, err := builder.Sum(zeros[:]); err != nil {
 			return err
 		}
 	}
@@ -121,6 +131,16 @@ func (n *ProtoNode) SetCidBuilder(builder cid.Builder) error {
 	return nil
 }
 
+// check whether the hasher is likely to be a usable one
+func checkHasher(indicator uint64, sizeHint int) error {
+	mhLen := sizeHint
+	if mhLen <= 0 {
+		mhLen = -1
+	}
+	_, err := mhcore.GetVariableHasher(indicator, mhLen)
+	return err
+}
+
 // LinkSlice is a slice of format.Links
 type LinkSlice []*format.Link