zrepl/util/chunking.go

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
}

func NewUnchunker(conn io.Reader) *Unchunker {
	return &Unchunker{
		in:                  conn,
		remainingChunkBytes: 0,
	}
}

func (c *Unchunker) Read(b []byte) (n int, err error) {

	if c.remainingChunkBytes == 0 {

		var nextChunkLen uint32
		err = binary.Read(c.in, ChunkHeaderByteOrder, &nextChunkLen)
		if err != nil {
			return
		}

		// A chunk of len 0 indicates end of stream
		if nextChunkLen == 0 {
			return 0, io.EOF
		}

		c.remainingChunkBytes = nextChunkLen
		c.ChunkCount++

	}

	maxRead := min(int(c.remainingChunkBytes), len(b))
	if maxRead < 0 {
		panic("Cannot read negative amount of bytes")
	}
	if maxRead == 0 {
		return 0, nil
	}

	n, err = c.in.Read(b[0:maxRead])
	if err != nil {
		return n, err
	}
	c.remainingChunkBytes -= uint32(n)

	return

}

func min(a, b int) int {
	if a < b {
		return a
	} else {
		return b
	}
}

type Chunker struct {
	ChunkCount          int
	in                  io.Reader
	remainingChunkBytes int
	payloadBuf          []byte
	headerBuf           *bytes.Buffer
}

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,
		payloadBuf:          buf,
		headerBuf:           &bytes.Buffer{},
	}

}

func (c *Chunker) Read(b []byte) (n int, err error) {

	//fmt.Printf("chunker: c.remainingChunkBytes: %d len(b): %d\n", c.remainingChunkBytes, len(b))

	n = 0
	if c.remainingChunkBytes == 0 {

		newPayloadLen, err := c.in.Read(c.payloadBuf)

		if newPayloadLen == 0 {
			return 0, io.EOF
		} else if err != nil {
			return newPayloadLen, err
		}

		c.remainingChunkBytes = newPayloadLen

		// Write chunk header
		c.headerBuf.Reset()
		nextChunkLen := uint32(newPayloadLen)
		headerLen := binary.Size(nextChunkLen)
		err = binary.Write(c.headerBuf, ChunkHeaderByteOrder, nextChunkLen)
		if err != nil {
			return n, err
		}
		copy(b[0:headerLen], c.headerBuf.Bytes())
		n += headerLen
		c.ChunkCount++
	}

	remainingBuf := b[n:]
	n2 := copy(remainingBuf, c.payloadBuf[:c.remainingChunkBytes])
	//fmt.Printf("chunker: written: %d\n", n+int(n2))
	c.remainingChunkBytes -= n2
	return n + int(n2), err
}