1
/*
2
Package gbytes provides a buffer that supports incrementally detecting input.
3

4
You use gbytes.Buffer with the gbytes.Say matcher.  When Say finds a match, it fastforwards the buffer's read cursor to the end of that match.
5

6
Subsequent matches against the buffer will only operate against data that appears *after* the read cursor.
7

8
The read cursor is an opaque implementation detail that you cannot access.  You should use the Say matcher to sift through the buffer.  You can always
9
access the entire buffer's contents with Contents().
10
*/
11
package gbytes
12

13
import (
14
	"errors"
15
	"fmt"
16
	"io"
17
	"regexp"
18
	"sync"
19
	"time"
20
)
21

22
/*
23
gbytes.Buffer implements an io.Writer and can be used with the gbytes.Say matcher.
24

25
You should only use a gbytes.Buffer in test code.  It stores all writes in an in-memory buffer - behavior that is inappropriate for production code!
26
*/
27
type Buffer struct {
28
	contents     []byte
29
	readCursor   uint64
30
	lock         *sync.Mutex
31
	detectCloser chan any
32
	closed       bool
33
}
34

35
/*
36
NewBuffer returns a new gbytes.Buffer
37
*/
38
func NewBuffer() *Buffer {
39
	return &Buffer{
40
		lock: &sync.Mutex{},
41
	}
42
}
43

44
/*
45
BufferWithBytes returns a new gbytes.Buffer seeded with the passed in bytes
46
*/
47
func BufferWithBytes(bytes []byte) *Buffer {
48
	return &Buffer{
49
		lock:     &sync.Mutex{},
50
		contents: bytes,
51
	}
52
}
53

54
/*
55
BufferReader returns a new gbytes.Buffer that wraps a reader.  The reader's contents are read into
56
the Buffer via io.Copy
57
*/
58
func BufferReader(reader io.Reader) *Buffer {
59
	b := &Buffer{
60
		lock: &sync.Mutex{},
61
	}
62

63
	go func() {
64
		io.Copy(b, reader)
65
		b.Close()
66
	}()
67

68
	return b
69
}
70

71
/*
72
Write implements the io.Writer interface
73
*/
74
func (b *Buffer) Write(p []byte) (n int, err error) {
75
	b.lock.Lock()
76
	defer b.lock.Unlock()
77

78
	if b.closed {
79
		return 0, errors.New("attempt to write to closed buffer")
80
	}
81

82
	b.contents = append(b.contents, p...)
83
	return len(p), nil
84
}
85

86
/*
87
Read implements the io.Reader interface. It advances the
88
cursor as it reads.
89
*/
90
func (b *Buffer) Read(d []byte) (int, error) {
91
	b.lock.Lock()
92
	defer b.lock.Unlock()
93

94
	if uint64(len(b.contents)) <= b.readCursor {
95
		return 0, io.EOF
96
	}
97

98
	n := copy(d, b.contents[b.readCursor:])
99
	b.readCursor += uint64(n)
100

101
	return n, nil
102
}
103

104
/*
105
Clear clears out the buffer's contents
106
*/
107
func (b *Buffer) Clear() error {
108
	b.lock.Lock()
109
	defer b.lock.Unlock()
110

111
	if b.closed {
112
		return errors.New("attempt to clear closed buffer")
113
	}
114

115
	b.contents = []byte{}
116
	b.readCursor = 0
117
	return nil
118
}
119

120
/*
121
Close signifies that the buffer will no longer be written to
122
*/
123
func (b *Buffer) Close() error {
124
	b.lock.Lock()
125
	defer b.lock.Unlock()
126

127
	b.closed = true
128

129
	return nil
130
}
131

132
/*
133
Closed returns true if the buffer has been closed
134
*/
135
func (b *Buffer) Closed() bool {
136
	b.lock.Lock()
137
	defer b.lock.Unlock()
138

139
	return b.closed
140
}
141

142
/*
143
Contents returns all data ever written to the buffer.
144
*/
145
func (b *Buffer) Contents() []byte {
146
	b.lock.Lock()
147
	defer b.lock.Unlock()
148

149
	contents := make([]byte, len(b.contents))
150
	copy(contents, b.contents)
151
	return contents
152
}
153

154
/*
155
Detect takes a regular expression and returns a channel.
156

157
The channel will receive true the first time data matching the regular expression is written to the buffer.
158
The channel is subsequently closed and the buffer's read-cursor is fast-forwarded to just after the matching region.
159

160
You typically don't need to use Detect and should use the ghttp.Say matcher instead.  Detect is useful, however, in cases where your code must
161
be branch and handle different outputs written to the buffer.
162

163
For example, consider a buffer hooked up to the stdout of a client library.  You may (or may not, depending on state outside of your control) need to authenticate the client library.
164

165
You could do something like:
166

167
select {
168
case <-buffer.Detect("You are not logged in"):
169

170
	//log in
171

172
case <-buffer.Detect("Success"):
173

174
	//carry on
175

176
case <-time.After(time.Second):
177

178
		//welp
179
	}
180

181
buffer.CancelDetects()
182

183
You should always call CancelDetects after using Detect.  This will close any channels that have not detected and clean up the goroutines that were spawned to support them.
184

185
Finally, you can pass detect a format string followed by variadic arguments.  This will construct the regexp using fmt.Sprintf.
186
*/
187
func (b *Buffer) Detect(desired string, args ...any) chan bool {
188
	formattedRegexp := desired
189
	if len(args) > 0 {
190
		formattedRegexp = fmt.Sprintf(desired, args...)
191
	}
192
	re := regexp.MustCompile(formattedRegexp)
193

194
	b.lock.Lock()
195
	defer b.lock.Unlock()
196

197
	if b.detectCloser == nil {
198
		b.detectCloser = make(chan any)
199
	}
200

201
	closer := b.detectCloser
202
	response := make(chan bool)
203
	go func() {
204
		ticker := time.NewTicker(10 * time.Millisecond)
205
		defer ticker.Stop()
206
		defer close(response)
207
		for {
208
			select {
209
			case <-ticker.C:
210
				b.lock.Lock()
211
				data, cursor := b.contents[b.readCursor:], b.readCursor
212
				loc := re.FindIndex(data)
213
				b.lock.Unlock()
214

215
				if loc != nil {
216
					response <- true
217
					b.lock.Lock()
218
					newCursorPosition := cursor + uint64(loc[1])
219
					if newCursorPosition >= b.readCursor {
220
						b.readCursor = newCursorPosition
221
					}
222
					b.lock.Unlock()
223
					return
224
				}
225
			case <-closer:
226
				return
227
			}
228
		}
229
	}()
230

231
	return response
232
}
233

234
/*
235
CancelDetects cancels any pending detects and cleans up their goroutines.  You should always call this when you're done with a set of Detect channels.
236
*/
237
func (b *Buffer) CancelDetects() {
238
	b.lock.Lock()
239
	defer b.lock.Unlock()
240

241
	close(b.detectCloser)
242
	b.detectCloser = nil
243
}
244

245
func (b *Buffer) didSay(re *regexp.Regexp) (bool, []byte) {
246
	b.lock.Lock()
247
	defer b.lock.Unlock()
248

249
	unreadBytes := b.contents[b.readCursor:]
250
	copyOfUnreadBytes := make([]byte, len(unreadBytes))
251
	copy(copyOfUnreadBytes, unreadBytes)
252

253
	loc := re.FindIndex(unreadBytes)
254

255
	if loc != nil {
256
		b.readCursor += uint64(loc[1])
257
		return true, copyOfUnreadBytes
258
	}
259
	return false, copyOfUnreadBytes
260
}
261

262