simplify chunk stream reader

pull/758/head
Darien Raymond 2017-12-02 01:19:27 +01:00
parent 2c6b46ee2d
commit 0481e6fa42
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169
1 changed files with 9 additions and 47 deletions

View File

@ -59,45 +59,24 @@ func (p *AEADChunkSizeParser) Decode(b []byte) (uint16, error) {
type ChunkStreamReader struct { type ChunkStreamReader struct {
sizeDecoder ChunkSizeDecoder sizeDecoder ChunkSizeDecoder
reader buf.Reader reader *buf.BufferedReader
buffer []byte buffer []byte
leftOver buf.MultiBuffer
leftOverSize int leftOverSize int
} }
func NewChunkStreamReader(sizeDecoder ChunkSizeDecoder, reader io.Reader) *ChunkStreamReader { func NewChunkStreamReader(sizeDecoder ChunkSizeDecoder, reader io.Reader) *ChunkStreamReader {
return &ChunkStreamReader{ return &ChunkStreamReader{
sizeDecoder: sizeDecoder, sizeDecoder: sizeDecoder,
reader: buf.NewReader(reader), reader: buf.NewBufferedReader(buf.NewReader(reader)),
buffer: make([]byte, sizeDecoder.SizeBytes()), buffer: make([]byte, sizeDecoder.SizeBytes()),
leftOver: buf.NewMultiBufferCap(16),
} }
} }
func (r *ChunkStreamReader) readAtLeast(size int) error {
mb := r.leftOver
r.leftOver = nil
for mb.Len() < size {
extra, err := r.reader.ReadMultiBuffer()
if err != nil {
mb.Release()
return err
}
mb.AppendMulti(extra)
}
r.leftOver = mb
return nil
}
func (r *ChunkStreamReader) readSize() (uint16, error) { func (r *ChunkStreamReader) readSize() (uint16, error) {
if r.sizeDecoder.SizeBytes() > r.leftOver.Len() { if _, err := io.ReadFull(r.reader, r.buffer); err != nil {
if err := r.readAtLeast(r.sizeDecoder.SizeBytes() - r.leftOver.Len()); err != nil { return 0, err
return 0, err
}
} }
common.Must2(io.ReadFull(&r.leftOver, r.buffer))
return r.sizeDecoder.Decode(r.buffer) return r.sizeDecoder.Decode(r.buffer)
} }
@ -113,31 +92,14 @@ func (r *ChunkStreamReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
} }
size = int(nextSize) size = int(nextSize)
} }
r.leftOverSize = size
if r.leftOver.IsEmpty() { mb, err := r.reader.ReadAtMost(size)
if err := r.readAtLeast(1); err != nil { if !mb.IsEmpty() {
return nil, err r.leftOverSize -= mb.Len()
}
}
leftOverLen := r.leftOver.Len()
if size >= leftOverLen {
mb := r.leftOver
r.leftOverSize = size - leftOverLen
r.leftOver = nil
return mb, nil return mb, nil
} }
return nil, err
mb := r.leftOver.SliceBySize(size)
mbLen := mb.Len()
if mbLen != size {
b := buf.New()
common.Must(b.Reset(buf.ReadFullFrom(&r.leftOver, size-mbLen)))
mb.Append(b)
}
r.leftOverSize = 0
return mb, nil
} }
type ChunkStreamWriter struct { type ChunkStreamWriter struct {