package dataext import ( "errors" "io" ) type brcMode int const ( modeSourceReading brcMode = 0 modeSourceFinished brcMode = 1 modeBufferReading brcMode = 2 modeBufferFinished brcMode = 3 ) type BufferedReadCloser interface { io.ReadCloser BufferedAll() ([]byte, error) Reset() error } type bufferedReadCloser struct { buffer []byte inner io.ReadCloser mode brcMode off int } func NewBufferedReadCloser(sub io.ReadCloser) BufferedReadCloser { return &bufferedReadCloser{ buffer: make([]byte, 0, 1024), inner: sub, mode: modeSourceReading, off: 0, } } func (b *bufferedReadCloser) Read(p []byte) (int, error) { switch b.mode { case modeSourceReading: n, err := b.inner.Read(p) if n > 0 { b.buffer = append(b.buffer, p[0:n]...) } if err == io.EOF { b.mode = modeSourceFinished } return n, err case modeSourceFinished: return 0, io.EOF case modeBufferReading: if len(b.buffer) <= b.off { b.mode = modeBufferFinished if len(p) == 0 { return 0, nil } return 0, io.EOF } n := copy(p, b.buffer[b.off:]) b.off += n return n, nil case modeBufferFinished: return 0, io.EOF default: return 0, errors.New("object in undefined status") } } func (b *bufferedReadCloser) Close() error { switch b.mode { case modeSourceReading: _, err := b.BufferedAll() if err != nil { return err } err = b.inner.Close() if err != nil { return err } b.mode = modeSourceFinished return nil case modeSourceFinished: return nil case modeBufferReading: b.mode = modeBufferFinished return nil case modeBufferFinished: return nil default: return errors.New("object in undefined status") } } func (b *bufferedReadCloser) BufferedAll() ([]byte, error) { switch b.mode { case modeSourceReading: arr := make([]byte, 1024) for b.mode == modeSourceReading { _, err := b.Read(arr) if err != nil && err != io.EOF { return nil, err } } if err := b.Reset(); err != nil { return nil, err } return b.buffer, nil case modeSourceFinished: return b.buffer, nil case modeBufferReading: return b.buffer, nil case modeBufferFinished: return b.buffer, nil default: return nil, errors.New("object in undefined status") } } // Reset resets the buffer to the beginning of the buffer. // If the original source is partially read, we will finish reading it and fill our buffer func (b *bufferedReadCloser) Reset() error { switch b.mode { case modeSourceReading: if b.off == 0 { return nil // nobody has read anything yet } err := b.Close() if err != nil { return err } b.mode = modeBufferReading b.off = 0 return nil case modeSourceFinished: err := b.Close() if err != nil { return err } b.mode = modeBufferReading b.off = 0 return nil case modeBufferReading: fallthrough case modeBufferFinished: b.mode = modeBufferReading b.off = 0 return nil default: return errors.New("object in undefined status") } }