github.com/4thel00z/pcopy@v0.0.0-20230830212547-a1758a3a86bc/util/peak.go (about) 1 package util 2 3 import ( 4 "bytes" 5 "io" 6 "strings" 7 ) 8 9 // PeakedReadCloser is a ReadCloser that allows peaking into a stream and buffering it in memory. 10 // It can be instantiated using the Peak function. After a stream has been peaked, it can still be fully 11 // read by reading the PeakedReadCloser. It first drained from the memory buffer, and then from the remaining 12 // underlying reader. 13 type PeakedReadCloser struct { 14 PeakedBytes []byte 15 LimitReached bool 16 peaked io.Reader 17 underlying io.ReadCloser 18 closed bool 19 } 20 21 // Peak reads the underlying ReadCloser into memory up until the limit and returns a PeakedReadCloser 22 func Peak(underlying io.ReadCloser, limit int) (*PeakedReadCloser, error) { 23 if underlying == nil { 24 underlying = io.NopCloser(strings.NewReader("")) 25 } 26 peaked := make([]byte, limit) 27 read, err := io.ReadFull(underlying, peaked) 28 if err != nil && err != io.ErrUnexpectedEOF && err != io.EOF { 29 return nil, err 30 } 31 return &PeakedReadCloser{ 32 PeakedBytes: peaked[:read], 33 LimitReached: read == limit, 34 underlying: underlying, 35 peaked: bytes.NewReader(peaked[:read]), 36 closed: false, 37 }, nil 38 } 39 40 // Read reads from the peaked bytes and then from the underlying stream 41 func (r *PeakedReadCloser) Read(p []byte) (n int, err error) { 42 if r.closed { 43 return 0, io.EOF 44 } 45 n, err = r.peaked.Read(p) 46 if err == io.EOF { 47 return r.underlying.Read(p) 48 } else if err != nil { 49 return 0, err 50 } 51 return 52 } 53 54 // Close closes the underlying stream 55 func (r *PeakedReadCloser) Close() error { 56 if r.closed { 57 return io.EOF 58 } 59 r.closed = true 60 return r.underlying.Close() 61 }