Merge pull request #3108 from homotopycolimit/FixChunkerWithBrokenReader

swarm/storage: fixes for tree chunker in the context of a broken reader

1 files changed, 21 insertions, 21 deletions

diff --git a/swarm/storage/common_test.go b/swarm/storage/common_test.go
index e81a82b7b..889b28a70 100644
--- a/swarm/storage/common_test.go
+++ b/swarm/storage/common_test.go
@@ -19,6 +19,7 @@ package storage
 import (
 	"bytes"
 	"crypto/rand"
+	"fmt"
 	"io"
 	"sync"
 	"testing"
@@ -27,32 +28,31 @@ import (
 	"github.com/ethereum/go-ethereum/logger/glog"
 )
 
-type limitedReader struct {
-	r    io.Reader
-	off  int64
-	size int64
+type brokenLimitedReader struct {
+	lr    io.Reader
+	errAt int
+	off   int
+	size  int
 }
 
-func limitReader(r io.Reader, size int) *limitedReader {
-	return &limitedReader{r, 0, int64(size)}
-}
-
-func (self *limitedReader) Read(buf []byte) (int, error) {
-	limit := int64(len(buf))
-	left := self.size - self.off
-	if limit >= left {
-		limit = left
-	}
-	n, err := self.r.Read(buf[:limit])
-	if err == nil && limit == left {
-		err = io.EOF
+func brokenLimitReader(data io.Reader, size int, errAt int) *brokenLimitedReader {
+	return &brokenLimitedReader{
+		lr:    data,
+		errAt: errAt,
+		size:  size,
 	}
-	self.off += int64(n)
-	return n, err
 }
 
 func testDataReader(l int) (r io.Reader) {
-	return limitReader(rand.Reader, l)
+	return io.LimitReader(rand.Reader, int64(l))
+}
+
+func (self *brokenLimitedReader) Read(buf []byte) (int, error) {
+	if self.off+len(buf) > self.errAt {
+		return 0, fmt.Errorf("Broken reader")
+	}
+	self.off += len(buf)
+	return self.lr.Read(buf)
 }
 
 func testDataReaderAndSlice(l int) (r io.Reader, slice []byte) {
@@ -60,7 +60,7 @@ func testDataReaderAndSlice(l int) (r io.Reader, slice []byte) {
 	if _, err := rand.Read(slice); err != nil {
 		panic("rand error")
 	}
-	r = limitReader(bytes.NewReader(slice), l)
+	r = io.LimitReader(bytes.NewReader(slice), int64(l))
 	return
 }