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  }