File: regretreader.go

package info (click to toggle)
golang-github-elazarl-goproxy 1.1%2Bgit20240726.8b0c205%2Bdfsg-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 612 kB
  • sloc: sh: 36; makefile: 8
file content (97 lines) | stat: -rw-r--r-- 2,714 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
package regretable

import (
	"io"
)

// A RegretableReader will allow you to read from a reader, and then
// to "regret" reading it, and push back everything you've read.
// For example,
//	rb := NewRegretableReader(bytes.NewBuffer([]byte{1,2,3}))
//	var b = make([]byte,1)
//	rb.Read(b) // b[0] = 1
//	rb.Regret()
//	ioutil.ReadAll(rb.Read) // returns []byte{1,2,3},nil
type RegretableReader struct {
	reader   io.Reader
	overflow bool
	r, w     int
	buf      []byte
}

var defaultBufferSize = 500

// Same as RegretableReader, but allows closing the underlying reader
type RegretableReaderCloser struct {
	RegretableReader
	c io.Closer
}

// Closes the underlying readCloser, you cannot regret after closing the stream
func (rbc *RegretableReaderCloser) Close() error {
	return rbc.c.Close()
}

// initialize a RegretableReaderCloser with underlying readCloser rc
func NewRegretableReaderCloser(rc io.ReadCloser) *RegretableReaderCloser {
	return &RegretableReaderCloser{*NewRegretableReader(rc), rc}
}

// initialize a RegretableReaderCloser with underlying readCloser rc
func NewRegretableReaderCloserSize(rc io.ReadCloser, size int) *RegretableReaderCloser {
	return &RegretableReaderCloser{*NewRegretableReaderSize(rc, size), rc}
}

// The next read from the RegretableReader will be as if the underlying reader
// was never read (or from the last point forget is called).
func (rb *RegretableReader) Regret() {
	if rb.overflow {
		panic("regretting after overflow makes no sense")
	}
	rb.r = 0
}

// Will "forget" everything read so far.
//	rb := NewRegretableReader(bytes.NewBuffer([]byte{1,2,3}))
//	var b = make([]byte,1)
//	rb.Read(b) // b[0] = 1
//	rb.Forget()
//	rb.Read(b) // b[0] = 2
//	rb.Regret()
//	ioutil.ReadAll(rb.Read) // returns []byte{2,3},nil
func (rb *RegretableReader) Forget() {
	if rb.overflow {
		panic("forgetting after overflow makes no sense")
	}
	rb.r = 0
	rb.w = 0
}

// initialize a RegretableReader with underlying reader r, whose buffer is size bytes long
func NewRegretableReaderSize(r io.Reader, size int) *RegretableReader {
	return &RegretableReader{reader: r, buf: make([]byte, size)}
}

// initialize a RegretableReader with underlying reader r
func NewRegretableReader(r io.Reader) *RegretableReader {
	return NewRegretableReaderSize(r, defaultBufferSize)
}

// reads from the underlying reader. Will buffer all input until Regret is called.
func (rb *RegretableReader) Read(p []byte) (n int, err error) {
	if rb.overflow {
		return rb.reader.Read(p)
	}
	if rb.r < rb.w {
		n = copy(p, rb.buf[rb.r:rb.w])
		rb.r += n
		return
	}
	n, err = rb.reader.Read(p)
	bn := copy(rb.buf[rb.w:], p[:n])
	rb.w, rb.r = rb.w+bn, rb.w+n
	if bn < n {
		rb.overflow = true
	}
	return
}