From 62bb7c966df68cb8e9058b04d1a32ab47499f31c Mon Sep 17 00:00:00 2001
From: Jeromy <why@ipfs.io>
Date: Tue, 22 Nov 2016 12:31:10 -0800
Subject: [PATCH] chekc that size input to newHamtShard is a power of two

License: MIT
Signed-off-by: Jeromy <why@ipfs.io>
---
 unixfs/hamt/hamt.go             | 33 ++++++++++++++++++++++++---------
 unixfs/hamt/hamt_stress_test.go | 18 +++++++++++++++---
 unixfs/hamt/hamt_test.go        | 21 ++++++++++++++-------
 unixfs/io/dirbuilder.go         |  7 ++++++-
 4 files changed, 59 insertions(+), 20 deletions(-)

diff --git a/unixfs/hamt/hamt.go b/unixfs/hamt/hamt.go
index a3b0c24d9..c9d6eb9dc 100644
--- a/unixfs/hamt/hamt.go
+++ b/unixfs/hamt/hamt.go
@@ -64,23 +64,31 @@ type child interface {
 	Label() string
 }
 
-func NewHamtShard(dserv dag.DAGService, size int) *HamtShard {
-	ds := makeHamtShard(dserv, size)
+func NewHamtShard(dserv dag.DAGService, size int) (*HamtShard, error) {
+	ds, err := makeHamtShard(dserv, size)
+	if err != nil {
+		return nil, err
+	}
+
 	ds.bitfield = big.NewInt(0)
 	ds.nd = new(dag.ProtoNode)
 	ds.hashFunc = HashMurmur3
-	return ds
+	return ds, nil
 }
 
-func makeHamtShard(ds dag.DAGService, size int) *HamtShard {
+func makeHamtShard(ds dag.DAGService, size int) (*HamtShard, error) {
+	lg2s := int(math.Log2(float64(size)))
+	if 1<<uint(lg2s) != size {
+		return nil, fmt.Errorf("hamt size should be a power of two")
+	}
 	maxpadding := fmt.Sprintf("%X", size-1)
 	return &HamtShard{
-		tableSizeLg2: int(math.Log2(float64(size))),
+		tableSizeLg2: lg2s,
 		prefixPadStr: fmt.Sprintf("%%0%dX", len(maxpadding)),
 		maxpadlen:    len(maxpadding),
 		tableSize:    size,
 		dserv:        ds,
-	}
+	}, nil
 }
 
 func NewHamtFromDag(dserv dag.DAGService, nd node.Node) (*HamtShard, error) {
@@ -102,7 +110,11 @@ func NewHamtFromDag(dserv dag.DAGService, nd node.Node) (*HamtShard, error) {
 		return nil, fmt.Errorf("only murmur3 supported as hash function")
 	}
 
-	ds := makeHamtShard(dserv, int(pbd.GetFanout()))
+	ds, err := makeHamtShard(dserv, int(pbd.GetFanout()))
+	if err != nil {
+		return nil, err
+	}
+
 	ds.nd = pbnd.Copy().(*dag.ProtoNode)
 	ds.children = make([]child, len(pbnd.Links()))
 	ds.bitfield = new(big.Int).SetBytes(pbd.GetData())
@@ -446,13 +458,16 @@ func (ds *HamtShard) modifyValue(ctx context.Context, hv *hashBits, key string,
 			return nil
 
 		default: // replace value with another shard, one level deeper
-			ns := NewHamtShard(ds.dserv, ds.tableSize)
+			ns, err := NewHamtShard(ds.dserv, ds.tableSize)
+			if err != nil {
+				return err
+			}
 			chhv := &hashBits{
 				b:        hash([]byte(child.key)),
 				consumed: hv.consumed,
 			}
 
-			err := ns.modifyValue(ctx, hv, key, val)
+			err = ns.modifyValue(ctx, hv, key, val)
 			if err != nil {
 				return err
 			}
diff --git a/unixfs/hamt/hamt_stress_test.go b/unixfs/hamt/hamt_stress_test.go
index 83ad25597..76357b23d 100644
--- a/unixfs/hamt/hamt_stress_test.go
+++ b/unixfs/hamt/hamt_stress_test.go
@@ -117,7 +117,11 @@ func validateOpSetCompletion(t *testing.T, s *HamtShard, keep, temp []string) er
 
 func executeOpSet(t *testing.T, ds dag.DAGService, width int, ops []testOp) (*HamtShard, error) {
 	ctx := context.TODO()
-	s := NewHamtShard(ds, width)
+	s, err := NewHamtShard(ds, width)
+	if err != nil {
+		return nil, err
+	}
+
 	e := ft.EmptyDirNode()
 	ds.Add(e)
 
@@ -188,7 +192,11 @@ func genOpSet(seed int64, keep, temp []string) []testOp {
 
 // executes the given op set with a repl to allow easier debugging
 func debugExecuteOpSet(ds dag.DAGService, width int, ops []testOp) (*HamtShard, error) {
-	s := NewHamtShard(ds, width)
+	s, err := NewHamtShard(ds, width)
+	if err != nil {
+		return nil, err
+	}
+
 	e := ft.EmptyDirNode()
 	ds.Add(e)
 	ctx := context.TODO()
@@ -236,7 +244,11 @@ mainloop:
 					}
 				}
 			case "restart":
-				s = NewHamtShard(ds, width)
+				var err error
+				s, err = NewHamtShard(ds, width)
+				if err != nil {
+					panic(err)
+				}
 				i = -1
 				continue mainloop
 			case "print":
diff --git a/unixfs/hamt/hamt_test.go b/unixfs/hamt/hamt_test.go
index cafa1ce8c..25b7125c4 100644
--- a/unixfs/hamt/hamt_test.go
+++ b/unixfs/hamt/hamt_test.go
@@ -30,7 +30,7 @@ func makeDir(ds dag.DAGService, size int) ([]string, *HamtShard, error) {
 }
 
 func makeDirWidth(ds dag.DAGService, size, width int) ([]string, *HamtShard, error) {
-	s := NewHamtShard(ds, width)
+	s, _ := NewHamtShard(ds, width)
 
 	var dirs []string
 	for i := 0; i < size; i++ {
@@ -136,7 +136,7 @@ func TestBasicSet(t *testing.T) {
 
 func TestDirBuilding(t *testing.T) {
 	ds := mdtest.Mock()
-	s := NewHamtShard(ds, 256)
+	s, _ := NewHamtShard(ds, 256)
 
 	_, s, err := makeDir(ds, 200)
 	if err != nil {
@@ -159,7 +159,7 @@ func TestDirBuilding(t *testing.T) {
 
 func TestShardReload(t *testing.T) {
 	ds := mdtest.Mock()
-	s := NewHamtShard(ds, 256)
+	s, _ := NewHamtShard(ds, 256)
 	ctx := context.Background()
 
 	_, s, err := makeDir(ds, 200)
@@ -287,7 +287,7 @@ func TestSetAfterMarshal(t *testing.T) {
 
 func TestDuplicateAddShard(t *testing.T) {
 	ds := mdtest.Mock()
-	dir := NewHamtShard(ds, 256)
+	dir, _ := NewHamtShard(ds, 256)
 	nd := new(dag.ProtoNode)
 	ctx := context.Background()
 
@@ -410,7 +410,7 @@ func TestRemoveElemsAfterMarshal(t *testing.T) {
 
 func TestBitfieldIndexing(t *testing.T) {
 	ds := mdtest.Mock()
-	s := NewHamtShard(ds, 256)
+	s, _ := NewHamtShard(ds, 256)
 
 	set := func(i int) {
 		s.bitfield.SetBit(s.bitfield, i, 1)
@@ -444,7 +444,7 @@ func TestBitfieldIndexing(t *testing.T) {
 // itself.
 func TestSetHamtChild(t *testing.T) {
 	ds := mdtest.Mock()
-	s := NewHamtShard(ds, 256)
+	s, _ := NewHamtShard(ds, 256)
 	ctx := context.Background()
 
 	e := ft.EmptyDirNode()
@@ -519,7 +519,7 @@ func printDiff(ds dag.DAGService, a, b *dag.ProtoNode) {
 
 func BenchmarkHAMTSet(b *testing.B) {
 	ds := mdtest.Mock()
-	sh := NewHamtShard(ds, 256)
+	sh, _ := NewHamtShard(ds, 256)
 	nd, err := sh.Node()
 	if err != nil {
 		b.Fatal(err)
@@ -550,3 +550,10 @@ func BenchmarkHAMTSet(b *testing.B) {
 		nd = out
 	}
 }
+
+func TestHamtBadSize(t *testing.T) {
+	_, err := NewHamtShard(nil, 7)
+	if err == nil {
+		t.Fatal("should have failed to construct hamt with bad size")
+	}
+}
diff --git a/unixfs/io/dirbuilder.go b/unixfs/io/dirbuilder.go
index 5f0fc2242..964eac582 100644
--- a/unixfs/io/dirbuilder.go
+++ b/unixfs/io/dirbuilder.go
@@ -85,7 +85,12 @@ func (d *Directory) AddChild(ctx context.Context, name string, nd node.Node) err
 }
 
 func (d *Directory) switchToSharding(ctx context.Context) error {
-	d.shard = hamt.NewHamtShard(d.dserv, DefaultShardWidth)
+	s, err := hamt.NewHamtShard(d.dserv, DefaultShardWidth)
+	if err != nil {
+		return err
+	}
+
+	d.shard = s
 	for _, lnk := range d.dirnode.Links() {
 		cnd, err := d.dserv.Get(ctx, lnk.Cid)
 		if err != nil {