github.com/coreos/goproxy@v0.0.0-20190513173959-f8dc2d7ba04e/regretable/regretreader.go (about)

     1  package regretable
     2  
     3  import (
     4  	"io"
     5  )
     6  
     7  // A RegretableReader will allow you to read from a reader, and then
     8  // to "regret" reading it, and push back everything you've read.
     9  // For example,
    10  //	rb := NewRegretableReader(bytes.NewBuffer([]byte{1,2,3}))
    11  //	var b = make([]byte,1)
    12  //	rb.Read(b) // b[0] = 1
    13  //	rb.Regret()
    14  //	ioutil.ReadAll(rb.Read) // returns []byte{1,2,3},nil
    15  type RegretableReader struct {
    16  	reader   io.Reader
    17  	overflow bool
    18  	r, w     int
    19  	buf      []byte
    20  }
    21  
    22  var defaultBufferSize = 500
    23  
    24  // Same as RegretableReader, but allows closing the underlying reader
    25  type RegretableReaderCloser struct {
    26  	RegretableReader
    27  	c      io.Closer
    28  }
    29  
    30  // Closes the underlying readCloser, you cannot regret after closing the stream
    31  func (rbc *RegretableReaderCloser) Close() error {
    32  	return rbc.c.Close()
    33  }
    34  
    35  // initialize a RegretableReaderCloser with underlying readCloser rc
    36  func NewRegretableReaderCloser(rc io.ReadCloser) *RegretableReaderCloser {
    37  	return &RegretableReaderCloser{*NewRegretableReader(rc), rc}
    38  }
    39  
    40  // initialize a RegretableReaderCloser with underlying readCloser rc
    41  func NewRegretableReaderCloserSize(rc io.ReadCloser, size int) *RegretableReaderCloser {
    42  	return &RegretableReaderCloser{*NewRegretableReaderSize(rc, size), rc}
    43  }
    44  
    45  // The next read from the RegretableReader will be as if the underlying reader
    46  // was never read (or from the last point forget is called).
    47  func (rb *RegretableReader) Regret() {
    48  	if rb.overflow {
    49  		panic("regretting after overflow makes no sense")
    50  	}
    51  	rb.r = 0
    52  }
    53  
    54  // Will "forget" everything read so far.
    55  //	rb := NewRegretableReader(bytes.NewBuffer([]byte{1,2,3}))
    56  //	var b = make([]byte,1)
    57  //	rb.Read(b) // b[0] = 1
    58  //	rb.Forget()
    59  //	rb.Read(b) // b[0] = 2
    60  //	rb.Regret()
    61  //	ioutil.ReadAll(rb.Read) // returns []byte{2,3},nil
    62  func (rb *RegretableReader) Forget() {
    63  	if rb.overflow {
    64  		panic("forgetting after overflow makes no sense")
    65  	}
    66  	rb.r = 0
    67  	rb.w = 0
    68  }
    69  
    70  // initialize a RegretableReader with underlying reader r, whose buffer is size bytes long
    71  func NewRegretableReaderSize(r io.Reader, size int) *RegretableReader {
    72  	return &RegretableReader{reader: r, buf: make([]byte, size) }
    73  }
    74  
    75  // initialize a RegretableReader with underlying reader r
    76  func NewRegretableReader(r io.Reader) *RegretableReader {
    77  	return NewRegretableReaderSize(r, defaultBufferSize)
    78  }
    79  
    80  // reads from the underlying reader. Will buffer all input until Regret is called.
    81  func (rb *RegretableReader) Read(p []byte) (n int, err error) {
    82  	if rb.overflow {
    83  		return rb.reader.Read(p)
    84  	}
    85  	if rb.r < rb.w {
    86  		n = copy(p, rb.buf[rb.r:rb.w])
    87  		rb.r += n
    88  		return
    89  	}
    90  	n, err = rb.reader.Read(p)
    91  	bn := copy(rb.buf[rb.w:], p[:n])
    92  	rb.w, rb.r = rb.w + bn, rb.w + n
    93  	if bn < n {
    94  		rb.overflow = true
    95  	}
    96  	return
    97  }