From b029ebdf11d0c480aed2b0cca46d89140adf5541 Mon Sep 17 00:00:00 2001
From: Anton Evangelatov <anton.evangelatov@gmail.com>
Date: Tue, 14 Aug 2018 10:44:03 +0200
Subject: [PATCH] swarm/bmt, swarm/network, swarm/storage: update BMT hash
 initialisation

---
 swarm/bmt/bmt.go                 |  8 +++-----
 swarm/bmt/bmt_test.go            | 12 ++++++------
 swarm/network/stream/delivery.go |  3 ++-
 swarm/storage/types.go           |  4 +++-
 4 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/swarm/bmt/bmt.go b/swarm/bmt/bmt.go
index 99179cb8e5..77b9b4f99f 100644
--- a/swarm/bmt/bmt.go
+++ b/swarm/bmt/bmt.go
@@ -23,8 +23,6 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
-
-	"github.com/ethereum/go-ethereum/swarm/chunk"
 )
 
 /*
@@ -57,9 +55,9 @@ Two implementations are provided:
 */
 
 const (
-	// SegmentCount is the maximum number of segments of the underlying chunk
+	// segmentCount is the maximum number of segments of the underlying chunk
 	// Should be equal to max-chunk-data-size / hash-size
-	SegmentCount = 128
+	segmentCount = 128
 	// PoolSize is the maximum number of bmt trees used by the hashers, i.e,
 	// the maximum number of concurrent BMT hashing operations performed by the same hasher
 	PoolSize = 8
@@ -320,7 +318,7 @@ func (h *Hasher) Sum(b []byte) (s []byte) {
 // with every full segment calls writeSection in a go routine
 func (h *Hasher) Write(b []byte) (int, error) {
 	l := len(b)
-	if l == 0 || l > chunk.DefaultSize {
+	if l == 0 || l > h.pool.Size {
 		return 0, nil
 	}
 	t := h.getTree()
diff --git a/swarm/bmt/bmt_test.go b/swarm/bmt/bmt_test.go
index 891d8cbb29..8763f194ad 100644
--- a/swarm/bmt/bmt_test.go
+++ b/swarm/bmt/bmt_test.go
@@ -224,14 +224,14 @@ func TestHasherReuse(t *testing.T) {
 // tests if bmt reuse is not corrupting result
 func testHasherReuse(poolsize int, t *testing.T) {
 	hasher := sha3.NewKeccak256
-	pool := NewTreePool(hasher, SegmentCount, poolsize)
+	pool := NewTreePool(hasher, segmentCount, poolsize)
 	defer pool.Drain(0)
 	bmt := New(pool)
 
 	for i := 0; i < 100; i++ {
 		data := newData(BufferSize)
 		n := rand.Intn(bmt.Size())
-		err := testHasherCorrectness(bmt, hasher, data, n, SegmentCount)
+		err := testHasherCorrectness(bmt, hasher, data, n, segmentCount)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -241,7 +241,7 @@ func testHasherReuse(poolsize int, t *testing.T) {
 // Tests if pool can be cleanly reused even in concurrent use by several hasher
 func TestBMTConcurrentUse(t *testing.T) {
 	hasher := sha3.NewKeccak256
-	pool := NewTreePool(hasher, SegmentCount, PoolSize)
+	pool := NewTreePool(hasher, segmentCount, PoolSize)
 	defer pool.Drain(0)
 	cycles := 100
 	errc := make(chan error)
@@ -451,7 +451,7 @@ func benchmarkBMTBaseline(t *testing.B, n int) {
 func benchmarkBMT(t *testing.B, n int) {
 	data := newData(n)
 	hasher := sha3.NewKeccak256
-	pool := NewTreePool(hasher, SegmentCount, PoolSize)
+	pool := NewTreePool(hasher, segmentCount, PoolSize)
 	bmt := New(pool)
 
 	t.ReportAllocs()
@@ -465,7 +465,7 @@ func benchmarkBMT(t *testing.B, n int) {
 func benchmarkBMTAsync(t *testing.B, n int, wh whenHash, double bool) {
 	data := newData(n)
 	hasher := sha3.NewKeccak256
-	pool := NewTreePool(hasher, SegmentCount, PoolSize)
+	pool := NewTreePool(hasher, segmentCount, PoolSize)
 	bmt := New(pool).NewAsyncWriter(double)
 	idxs, segments := splitAndShuffle(bmt.SectionSize(), data)
 	shuffle(len(idxs), func(i int, j int) {
@@ -483,7 +483,7 @@ func benchmarkBMTAsync(t *testing.B, n int, wh whenHash, double bool) {
 func benchmarkPool(t *testing.B, poolsize, n int) {
 	data := newData(n)
 	hasher := sha3.NewKeccak256
-	pool := NewTreePool(hasher, SegmentCount, poolsize)
+	pool := NewTreePool(hasher, segmentCount, poolsize)
 	cycles := 100
 
 	t.ReportAllocs()
diff --git a/swarm/network/stream/delivery.go b/swarm/network/stream/delivery.go
index f6732e941f..ad74e32c34 100644
--- a/swarm/network/stream/delivery.go
+++ b/swarm/network/stream/delivery.go
@@ -24,6 +24,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/metrics"
 	"github.com/ethereum/go-ethereum/p2p/discover"
+	cp "github.com/ethereum/go-ethereum/swarm/chunk"
 	"github.com/ethereum/go-ethereum/swarm/log"
 	"github.com/ethereum/go-ethereum/swarm/network"
 	"github.com/ethereum/go-ethereum/swarm/spancontext"
@@ -244,7 +245,7 @@ R:
 			continue R
 		default:
 		}
-		if len(req.SData) > chunk.DefaultSize {
+		if len(req.SData) > int(cp.DefaultSize) {
 			log.Warn("received chunk is bigger than expected", "len", len(req.SData))
 			continue R
 		}
diff --git a/swarm/storage/types.go b/swarm/storage/types.go
index b6692d4829..b1354bfe56 100644
--- a/swarm/storage/types.go
+++ b/swarm/storage/types.go
@@ -115,7 +115,9 @@ func MakeHashFunc(hash string) SwarmHasher {
 	case "BMT":
 		return func() SwarmHash {
 			hasher := sha3.NewKeccak256
-			pool := bmt.NewTreePool(hasher, bmt.SegmentCount, bmt.PoolSize)
+			hasherSize := hasher().Size()
+			segmentCount := int(chunk.DefaultSize) / hasherSize
+			pool := bmt.NewTreePool(hasher, segmentCount, bmt.PoolSize)
 			return bmt.New(pool)
 		}
 	}