move implementation to internal/ directory (#828)

internal/util/chunking/chunking.go

@@ -0,0 +1,204 @@
+package chunking
+
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+)
+
+var ChunkBufSize uint32 = 32 * 1024
+var ChunkHeaderByteOrder = binary.LittleEndian
+
+type Unchunker struct {
+	ChunkCount          int
+	in                  io.Reader
+	remainingChunkBytes uint32
+	finishErr           error
+}
+
+func NewUnchunker(conn io.Reader) *Unchunker {
+	return &Unchunker{
+		in:                  conn,
+		remainingChunkBytes: 0,
+	}
+}
+
+func (c *Unchunker) Read(b []byte) (n int, err error) {
+
+	if c.finishErr != nil {
+		return 0, c.finishErr
+	}
+
+	if c.remainingChunkBytes == 0 {
+
+		var nextChunkLen uint32
+		err = binary.Read(c.in, ChunkHeaderByteOrder, &nextChunkLen)
+		if err != nil {
+			c.finishErr = err // can't handle this
+			return
+		}
+
+		// A chunk of len 0 indicates end of stream
+		if nextChunkLen == 0 {
+			c.finishErr = io.EOF
+			return 0, c.finishErr
+		}
+
+		c.remainingChunkBytes = nextChunkLen
+		c.ChunkCount++
+
+	}
+
+	if c.remainingChunkBytes == 0 {
+		panic("internal inconsistency: c.remainingChunkBytes must be > 0")
+	}
+	if len(b) <= 0 {
+		panic("cannot read into buffer of length 0")
+	}
+
+	maxRead := min(int(c.remainingChunkBytes), len(b))
+
+	n, err = c.in.Read(b[0:maxRead])
+	if err != nil {
+		return
+	}
+	c.remainingChunkBytes -= uint32(n)
+
+	return
+
+}
+
+func (c *Unchunker) Close() (err error) {
+
+	buf := make([]byte, 4096)
+	for err == nil {
+		_, err = c.Read(buf)
+		if err == io.EOF {
+			err = nil
+			break
+		}
+	}
+	return
+}
+
+func min(a, b int) int {
+	if a < b {
+		return a
+	} else {
+		return b
+	}
+}
+
+type Chunker struct {
+	ChunkCount          int
+	in                  io.Reader
+	inEOF               bool
+	remainingChunkBytes int
+	payloadBufLen       int
+	payloadBuf          []byte
+	headerBuf           *bytes.Buffer
+	finalHeaderBuffered bool
+}
+
+func NewChunker(conn io.Reader) Chunker {
+	return NewChunkerSized(conn, ChunkBufSize)
+}
+
+func NewChunkerSized(conn io.Reader, chunkSize uint32) Chunker {
+
+	buf := make([]byte, int(chunkSize)-binary.Size(chunkSize))
+
+	return Chunker{
+		in:                  conn,
+		remainingChunkBytes: 0,
+		payloadBufLen:       0,
+		payloadBuf:          buf,
+		headerBuf:           &bytes.Buffer{},
+	}
+
+}
+
+func (c *Chunker) Read(b []byte) (n int, err error) {
+
+	if len(b) == 0 {
+		panic("unexpected empty output buffer")
+	}
+
+	if c.inEOF && c.finalHeaderBuffered && c.headerBuf.Len() == 0 { // all bufs empty and no more bytes to expect
+		return 0, io.EOF
+	}
+
+	n = 0
+
+	if c.headerBuf.Len() > 0 { // first drain the header buf
+		nh, err := c.headerBuf.Read(b[n:])
+		if nh > 0 {
+			n += nh
+		}
+		if nh == 0 || (err != nil && err != io.EOF) {
+			panic("unexpected behavior: in-memory buffer should not throw errors")
+		}
+		if c.headerBuf.Len() != 0 {
+			return n, nil // finish writing the header before we can proceed with payload
+		}
+		if c.finalHeaderBuffered {
+			// we just wrote the final header
+			return n, io.EOF
+		}
+	}
+
+	if c.remainingChunkBytes > 0 { // then drain the payload buf
+
+		npl := copy(b[n:], c.payloadBuf[c.payloadBufLen-c.remainingChunkBytes:c.payloadBufLen])
+		c.remainingChunkBytes -= npl
+		if c.remainingChunkBytes < 0 {
+			panic("unexpected behavior, copy() should not copy more than max(cap(), len())")
+		}
+		n += npl
+	}
+
+	if c.remainingChunkBytes == 0 && !c.inEOF { // fillup bufs
+
+		newPayloadLen, err := c.in.Read(c.payloadBuf)
+
+		if newPayloadLen > 0 {
+			c.payloadBufLen = newPayloadLen
+			c.remainingChunkBytes = newPayloadLen
+		}
+		if err == io.EOF {
+			c.inEOF = true
+		} else if err != nil {
+			return n, err
+		}
+		if newPayloadLen == 0 { // apparently, this happens with some Readers
+			c.finalHeaderBuffered = true
+		}
+
+		// Fill header buf
+		{
+			c.headerBuf.Reset()
+			nextChunkLen := uint32(newPayloadLen)
+			err := binary.Write(c.headerBuf, ChunkHeaderByteOrder, nextChunkLen)
+			if err != nil {
+				panic("unexpected error, write to in-memory buffer should not throw error")
+			}
+		}
+
+		if c.headerBuf.Len() == 0 {
+			panic("unexpected empty header buf")
+		}
+
+	} else if c.remainingChunkBytes == 0 && c.inEOF && !c.finalHeaderBuffered {
+
+		c.headerBuf.Reset()
+		err := binary.Write(c.headerBuf, ChunkHeaderByteOrder, uint32(0))
+		if err != nil {
+			panic("unexpected error, write to in-memory buffer should not throw error [2]")
+		}
+		c.finalHeaderBuffered = true
+
+	}
+
+	return
+
+}

internal/util/chunking/chunking_test.go

@@ -0,0 +1,74 @@
+package chunking
+
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+	"reflect"
+	"testing"
+	"testing/quick"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestUnchunker(t *testing.T) {
+
+	buf := bytes.Buffer{}
+	binary.Write(&buf, ChunkHeaderByteOrder, uint32(2))
+	buf.WriteByte(0xca)
+	buf.WriteByte(0xfe)
+	binary.Write(&buf, ChunkHeaderByteOrder, uint32(0))
+	buf.WriteByte(0xff) // sentinel, should not be read
+
+	un := NewUnchunker(&buf)
+
+	recv := bytes.Buffer{}
+	n, err := io.Copy(&recv, un)
+	assert.Nil(t, err)
+	assert.Equal(t, int64(2), n)
+	assert.Equal(t, []byte{0xca, 0xfe}, recv.Bytes())
+
+}
+
+func TestChunker(t *testing.T) {
+
+	buf := bytes.Buffer{}
+	buf.WriteByte(0xca)
+	buf.WriteByte(0xfe)
+
+	ch := NewChunker(&buf)
+
+	chunked := bytes.Buffer{}
+	n, err := io.Copy(&chunked, &ch)
+	assert.Nil(t, err)
+	assert.Equal(t, int64(4+2+4), n)
+	assert.Equal(t, []byte{0x2, 0x0, 0x0, 0x0, 0xca, 0xfe, 0x0, 0x0, 0x0, 0x0}, chunked.Bytes())
+
+}
+
+func TestUnchunkerUnchunksChunker(t *testing.T) {
+
+	f := func(b []byte) bool {
+
+		buf := bytes.NewBuffer(b)
+		ch := NewChunker(buf)
+		unch := NewUnchunker(&ch)
+		var tx bytes.Buffer
+		_, err := io.Copy(&tx, unch)
+		if err != nil {
+			return false
+		}
+
+		return reflect.DeepEqual(b, tx.Bytes())
+	}
+
+	cfg := quick.Config{
+		MaxCount:      3 * int(ChunkBufSize),
+		MaxCountScale: 2.0,
+	}
+
+	if err := quick.Check(f, &cfg); err != nil {
+		t.Error(err)
+	}
+
+}