fix: apply code-review feedback

Co-authored-by: Masih H. Derkani <m@derkani.org>
ipfs · Nov 22, 2022 · d910fbc · d910fbc
1 parent 43dba3c
commit d910fbc
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 55 deletions.
diff --git a/merkledag_test.go b/merkledag_test.go
@@ -83,54 +83,62 @@ func traverseAndCheck(t *testing.T, root ipld.Node, ds ipld.DAGService, hasF fun
 
 type brokenBuilder struct{}
 
-func (brokenBuilder) Sum(data []byte) (cid.Cid, error) { return cid.Undef, fmt.Errorf("Nope!") }
-func (brokenBuilder) GetCodec() uint64                 { return 0 }
-func (b brokenBuilder) WithCodec(uint64) cid.Builder   { return b }
+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)
-	}
 
-	// hasher we can't use, should error
-	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)
+		}
+	})
 
-	// broken builder, should error
-	err = n.SetCidBuilder(brokenBuilder{})
-	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("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) {

diff --git a/node.go b/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,21 +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
 		}
-		// ensure the hasher is a usable one
-		_, err := mhcore.GetVariableHasher(p.MhType, mhLen)
-		if err != nil {
+	case *cid.Prefix:
+		if err := checkHasher(b.MhType, b.MhLength); err != nil {
 			return err
 		}
-	} else {
-		// have to test it's a usable hasher directly
-		// this is only a basic check, there are still ways it may break
-		_, err := builder.Sum([]byte{0})
-		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
 		}
 	}
@@ -129,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