1package streamframe
2
3import (
4 "bytes"
5 "encoding/binary"
6 "errors"
7 "io"
8 "sync"
9 "testing"
10 "time"
11
12 "github.com/stretchr/testify/require"
13)
14
15// TestFramerRoundTrip proves framed envelopes survive partial IO: the
16// writer emits 4-byte length + payload, the reader reassembles.
17func TestFramerRoundTrip(t *testing.T) {
18 reader, writer := io.Pipe()
19 defer func() {
20 _ = reader.Close()
21 _ = writer.Close()
22 }()
23 sender := NewFramer(nil, writer, nil)
24 receiver := NewFramer(reader, nil, nil)
25 defer func() {
26 _ = sender.Close()
27 _ = receiver.Close()
28 }()
29 payloads := [][]byte{[]byte("hello"), bytes.Repeat([]byte("x"), 70000), []byte("bye")}
30 done := make(chan error, 1)
31 go func() {
32 for _, payload := range payloads {
33 if err := sender.Send(payload); err != nil {
34 done <- err
35 return
36 }
37 }
38 done <- nil
39 }()
40 for _, want := range payloads {
41 got, err := receiver.Recv()
42 require.NoError(t, err)
43 require.Equal(t, want, got)
44 }
45 require.NoError(t, <-done)
46}
47
48// TestFramerPartialReads feeds the length header and body byte-by-byte.
49func TestFramerPartialReads(t *testing.T) {
50 payload := []byte("partial payload")
51 var encoded bytes.Buffer
52 var header [4]byte
53 binary.BigEndian.PutUint32(header[:], uint32(len(payload)))
54 encoded.Write(header[:])
55 encoded.Write(payload)
56 raw := encoded.Bytes()
57
58 pr, pw := io.Pipe()
59 defer func() {
60 _ = pr.Close()
61 _ = pw.Close()
62 }()
63 receiver := NewFramer(pr, nil, nil)
64 defer func() { _ = receiver.Close() }()
65 go func() {
66 for _, b := range raw {
67 _, _ = pw.Write([]byte{b})
68 }
69 }()
70 got, err := receiver.Recv()
71 require.NoError(t, err)
72 require.Equal(t, payload, got)
73}
74
75// TestFramerRejectsZeroAndOversize proves bounds are enforced before
76// allocation: zero lengths fail, lengths above the negotiated ceiling fail,
77// lengths above the absolute ceiling fail even with a raised negotiation.
78func TestFramerRejectsZeroAndOversize(t *testing.T) {
79 t.Run("zero length", func(t *testing.T) {
80 pr, pw := io.Pipe()
81 defer func() {
82 _ = pr.Close()
83 _ = pw.Close()
84 }()
85 receiver := NewFramer(pr, nil, nil)
86 defer func() { _ = receiver.Close() }()
87 go func() { _, _ = pw.Write([]byte{0, 0, 0, 0}) }()
88 _, err := receiver.Recv()
89 require.ErrorIs(t, err, ErrZeroLength)
90 })
91 t.Run("above negotiated", func(t *testing.T) {
92 pr, pw := io.Pipe()
93 defer func() {
94 _ = pr.Close()
95 _ = pw.Close()
96 }()
97 receiver := NewFramer(pr, nil, nil, WithMaxEnvelope(16))
98 defer func() { _ = receiver.Close() }()
99 go func() {
100 var header [4]byte
101 binary.BigEndian.PutUint32(header[:], 17)
102 _, _ = pw.Write(header[:])
103 }()
104 _, err := receiver.Recv()
105 require.ErrorIs(t, err, ErrTooLarge)
106 })
107 t.Run("send above negotiated", func(t *testing.T) {
108 sender := NewFramer(nil, io.Discard, nil, WithMaxEnvelope(4))
109 defer func() { _ = sender.Close() }()
110 require.ErrorIs(t, sender.Send([]byte("12345")), ErrTooLarge)
111 require.ErrorIs(t, sender.Send(nil), ErrZeroLength)
112 })
113}
114
115// TestFramerTruncation proves a short body fails instead of blocking forever.
116func TestFramerTruncation(t *testing.T) {
117 pr, pw := io.Pipe()
118 receiver := NewFramer(pr, nil, nil)
119 defer func() { _ = receiver.Close() }()
120 go func() {
121 var header [4]byte
122 binary.BigEndian.PutUint32(header[:], 100)
123 _, _ = pw.Write(header[:])
124 _, _ = pw.Write([]byte("short"))
125 _ = pw.Close()
126 }()
127 _, err := receiver.Recv()
128 require.Error(t, err)
129}
130
131// TestFramerTrailingBoundary proves two back-to-back envelopes split exactly.
132func TestFramerTrailingBoundary(t *testing.T) {
133 var encoded bytes.Buffer
134 for _, payload := range [][]byte{[]byte("one"), []byte("two")} {
135 var header [4]byte
136 binary.BigEndian.PutUint32(header[:], uint32(len(payload)))
137 encoded.Write(header[:])
138 encoded.Write(payload)
139 }
140 receiver := NewFramer(bytes.NewReader(encoded.Bytes()), nil, nil)
141 defer func() { _ = receiver.Close() }()
142 first, err := receiver.Recv()
143 require.NoError(t, err)
144 require.Equal(t, []byte("one"), first)
145 second, err := receiver.Recv()
146 require.NoError(t, err)
147 require.Equal(t, []byte("two"), second)
148}
149
150// TestFramerCloseUnblocksSendRecv proves Close interrupts both directions
151// and runs the injected closer exactly once.
152func TestFramerCloseUnblocksSendRecv(t *testing.T) {
153 pr, pw := io.Pipe()
154 closes := 0
155 // The injected closer releases blocked I/O, like the concrete
156 // adapters will in P3.3; the Framer itself never owns the stream.
157 framer := NewFramer(pr, pw, func() error {
158 closes++
159 _ = pr.Close()
160 _ = pw.Close()
161 return nil
162 })
163 recvErr := make(chan error, 1)
164 go func() {
165 _, err := framer.Recv()
166 recvErr <- err
167 }()
168 time.Sleep(50 * time.Millisecond)
169 require.NoError(t, framer.Close())
170 require.NoError(t, framer.Close())
171 select {
172 case err := <-recvErr:
173 require.True(t, errors.Is(err, io.EOF) || errors.Is(err, io.ErrClosedPipe) || errors.Is(err, ErrClosed))
174 case <-time.After(5 * time.Second):
175 t.Fatal("Recv was not unblocked by Close")
176 }
177 require.Equal(t, 1, closes)
178}
179
180// TestFramerByteBudgetExhaustion proves async admission stops at the byte
181// budget and resumes after the writer drains.
182func TestFramerByteBudgetExhaustion(t *testing.T) {
183 writer := &gateWriter{gate: make(chan struct{}, 2)}
184 sender := NewFramer(nil, writer, nil, WithQueueBounds(8, 64))
185 defer func() {
186 writer.gate <- struct{}{}
187 writer.gate <- struct{}{}
188 _ = sender.Close()
189 }()
190
191 payload := bytes.Repeat([]byte("a"), 32)
192 require.NoError(t, sender.SendAsync(payload))
193 // Wait until the sole writer has dequeued and is blocked writing the first
194 // frame. Dequeued writes no longer consume the queue's admission budget.
195 require.Eventually(t, func() bool {
196 sender.egressMu.Lock()
197 defer sender.egressMu.Unlock()
198 return len(sender.egress) == 0 && sender.queuedBytes == 0
199 }, 5*time.Second, time.Millisecond)
200 require.NoError(t, sender.SendAsync(payload))
201 require.ErrorIs(t, sender.SendAsync(payload), ErrBackpressure)
202
203 writer.gate <- struct{}{}
204 require.Eventually(t, func() bool {
205 sender.egressMu.Lock()
206 defer sender.egressMu.Unlock()
207 return len(sender.egress) == 0 && sender.queuedBytes == 0
208 }, 5*time.Second, time.Millisecond)
209 require.NoError(t, sender.SendAsync(payload))
210}
211
212// TestFramerSyncBypassKeepsByteBudget proves synchronous sends do not
213// corrupt async byte accounting: a dequeued sync frame releases no
214// charge, so later async admission cannot exceed the byte budget. Each
215// writer Write consumes one gate token, pinning the dequeue order to
216// async A, sync S, async B: after A is written the used budget must
217// still be B's 64 bytes, so a further 64-byte frame is backpressure.
218// (The old code released the dequeued sync frame's 36 bytes and wrongly
219// admitted it.)
220func TestFramerSyncBypassKeepsByteBudget(t *testing.T) {
221 gate := &gateWriter{gate: make(chan struct{}, 8)}
222 sender := NewFramer(nil, gate, nil, WithQueueBounds(8, 100))
223 defer func() { _ = sender.Close() }()
224
225 asyncA := bytes.Repeat([]byte("a"), 32) // framed 36, counted
226 syncS := bytes.Repeat([]byte("s"), 32) // framed 36, uncounted
227 asyncB := bytes.Repeat([]byte("b"), 60) // framed 64, counted
228 require.NoError(t, sender.SendAsync(asyncA))
229
230 // Wait until the writer holds A (channel empty), so the sync Send
231 // below can only land behind A and the length poll below cannot
232 // mistake A for S.
233 queueLen := func() int {
234 sender.egressMu.Lock()
235 defer sender.egressMu.Unlock()
236 return len(sender.egress)
237 }
238 queuedBytes := func() uint64 {
239 sender.egressMu.Lock()
240 defer sender.egressMu.Unlock()
241 return sender.queuedBytes
242 }
243 require.Eventually(t, func() bool {
244 return queueLen() == 0
245 }, 5*time.Second, time.Millisecond)
246
247 syncDone := make(chan error, 1)
248 go func() { syncDone <- sender.Send(syncS) }()
249 require.Eventually(t, func() bool {
250 return queueLen() == 1
251 }, 5*time.Second, time.Millisecond)
252 require.NoError(t, sender.SendAsync(asyncB))
253
254 // Finish A only, then wait until S is dequeued: S's dequeue must
255 // release no charge, leaving B's 64 bytes charged (36 free). A
256 // further 64-byte frame no longer fits.
257 gate.gate <- struct{}{}
258 require.Eventually(t, func() bool {
259 return gate.writtenForTest() == 1 && queueLen() == 1
260 }, 5*time.Second, time.Millisecond)
261 require.Equal(t, uint64(64), queuedBytes(), "sync dequeue must release no byte charge")
262
263 // A 64-byte frame no longer fits the remaining 36-byte budget.
264 probe := bytes.Repeat([]byte("c"), 60) // framed 64
265 require.ErrorIs(t, sender.SendAsync(probe), ErrBackpressure)
266
267 // Drain the rest in queue order.
268 gate.gate <- struct{}{}
269 gate.gate <- struct{}{}
270 require.NoError(t, <-syncDone)
271 require.Eventually(t, func() bool {
272 return gate.writtenForTest() == 3
273 }, 5*time.Second, time.Millisecond)
274 for _, want := range [][]byte{asyncA, syncS, asyncB} {
275 got, err := gate.nextFrame()
276 require.NoError(t, err)
277 require.Equal(t, want, got[4:])
278 }
279}
280
281// TestFramerRetainedPayloadSafety proves Send copies: mutating the caller's
282// buffer afterwards does not corrupt the queued envelope.
283func TestFramerRetainedPayloadSafety(t *testing.T) {
284 pr, pw := io.Pipe()
285 defer func() {
286 _ = pr.Close()
287 _ = pw.Close()
288 }()
289 sender := NewFramer(nil, pw, nil)
290 defer func() { _ = sender.Close() }()
291 receiver := NewFramer(pr, nil, nil)
292 defer func() { _ = receiver.Close() }()
293 payload := []byte("mutable")
294 require.NoError(t, sender.SendAsync(payload))
295 for i := range payload {
296 payload[i] = 'X'
297 }
298 got, err := receiver.Recv()
299 require.NoError(t, err)
300 require.Equal(t, []byte("mutable"), got)
301}
302
303// TestFramerConcurrentCloseSendRecv hammers all three operations at once
304// under -race: no panic, no deadlock, every call returns.
305func TestFramerConcurrentCloseSendRecv(t *testing.T) {
306 pr, pw := io.Pipe()
307 framer := NewFramer(pr, pw, func() error {
308 _ = pr.Close()
309 _ = pw.Close()
310 return nil
311 })
312 var wg sync.WaitGroup
313 for range 4 {
314 wg.Add(1)
315 go func() {
316 defer wg.Done()
317 _ = framer.Send([]byte("x"))
318 }()
319 wg.Add(1)
320 go func() {
321 defer wg.Done()
322 _, _ = framer.Recv()
323 }()
324 }
325 time.Sleep(20 * time.Millisecond)
326 require.NoError(t, framer.Close())
327 wg.Wait()
328}
329
330type blockingWriter struct {
331 release chan struct{}
332}
333
334func (w *blockingWriter) Write(p []byte) (int, error) {
335 <-w.release
336 return len(p), nil
337}
338
339// gateWriter releases one framed write per gate token and records
340// every write for ordered inspection. It pins the writer-side
341// interleaving a test needs without touching Framer internals.
342type gateWriter struct {
343 gate chan struct{}
344
345 mu sync.Mutex
346 frames [][]byte
347}
348
349func (w *gateWriter) Write(p []byte) (int, error) {
350 <-w.gate
351 w.mu.Lock()
352 defer w.mu.Unlock()
353 w.frames = append(w.frames, append([]byte(nil), p...))
354 return len(p), nil
355}
356
357func (w *gateWriter) writtenForTest() int {
358 w.mu.Lock()
359 defer w.mu.Unlock()
360 return len(w.frames)
361}
362
363func (w *gateWriter) nextFrame() ([]byte, error) {
364 deadline := time.Now().Add(5 * time.Second)
365 for {
366 w.mu.Lock()
367 if len(w.frames) > 0 {
368 frame := w.frames[0]
369 w.frames = w.frames[1:]
370 w.mu.Unlock()
371 return frame, nil
372 }
373 w.mu.Unlock()
374 if time.Now().After(deadline) {
375 return nil, errors.New("gateWriter: timed out waiting for frame")
376 }
377 time.Sleep(time.Millisecond)
378 }
379}
380
381// TestFramerRecvBoundedAppliesCallerLimitBeforeBody proves a per-call limit
382// narrower than the framer ceiling is enforced on the four-byte prefix: the
383// over-limit body is left unread, and a matching frame still round-trips.
384func TestFramerRecvBoundedAppliesCallerLimitBeforeBody(t *testing.T) {
385 body := []byte("01234567")
386 encode := func() []byte {
387 var encoded bytes.Buffer
388 var header [4]byte
389 binary.BigEndian.PutUint32(header[:], uint32(len(body)))
390 encoded.Write(header[:])
391 encoded.Write(body)
392 return encoded.Bytes()
393 }
394
395 src := bytes.NewReader(encode())
396 receiver := NewFramer(src, nil, nil)
397 defer func() { _ = receiver.Close() }()
398 _, err := receiver.RecvBounded(uint64(len(body) - 1))
399 require.ErrorIs(t, err, ErrTooLarge)
400 require.Equal(t, len(body), src.Len(), "over-limit body must not be read")
401
402 // A fresh stream at the matching limit round-trips the same frame. (A
403 // rejected prefix desynchronizes the stream, exactly as an oversize
404 // frame always has, so it is never re-read.)
405 src = bytes.NewReader(encode())
406 receiver = NewFramer(src, nil, nil)
407 defer func() { _ = receiver.Close() }()
408 got, err := receiver.RecvBounded(uint64(len(body)))
409 require.NoError(t, err)
410 require.Equal(t, body, got)
411}
412
413// TestFramerRecvBoundedPartialPrefixAndBody feeds the prefix and body one byte
414// at a time: the bounded path reassembles exactly like Recv.
415func TestFramerRecvBoundedPartialPrefixAndBody(t *testing.T) {
416 payload := []byte("bounded partial payload")
417 var encoded bytes.Buffer
418 var header [4]byte
419 binary.BigEndian.PutUint32(header[:], uint32(len(payload)))
420 encoded.Write(header[:])
421 encoded.Write(payload)
422
423 pr, pw := io.Pipe()
424 defer func() {
425 _ = pr.Close()
426 _ = pw.Close()
427 }()
428 receiver := NewFramer(pr, nil, nil)
429 defer func() { _ = receiver.Close() }()
430 go func() {
431 for _, b := range encoded.Bytes() {
432 if _, err := pw.Write([]byte{b}); err != nil {
433 return
434 }
435 }
436 }()
437 got, err := receiver.RecvBounded(uint64(len(payload)))
438 require.NoError(t, err)
439 require.Equal(t, payload, got)
440}
441
442// TestFramerRecvBoundedZeroLengthAndZeroLimit proves the bounded path keeps
443// Recv's zero-length refusal and treats a zero caller limit as refusing every
444// non-empty frame before its body is read.
445func TestFramerRecvBoundedZeroLengthAndZeroLimit(t *testing.T) {
446 t.Run("zero length", func(t *testing.T) {
447 pr, pw := io.Pipe()
448 defer func() {
449 _ = pr.Close()
450 _ = pw.Close()
451 }()
452 receiver := NewFramer(pr, nil, nil)
453 defer func() { _ = receiver.Close() }()
454 go func() { _, _ = pw.Write([]byte{0, 0, 0, 0}) }()
455 _, err := receiver.RecvBounded(64)
456 require.ErrorIs(t, err, ErrZeroLength)
457 })
458 t.Run("zero limit refuses body", func(t *testing.T) {
459 src := bytes.NewReader([]byte{0, 0, 0, 4, 'a', 'b', 'c', 'd'})
460 receiver := NewFramer(src, nil, nil)
461 defer func() { _ = receiver.Close() }()
462 _, err := receiver.RecvBounded(0)
463 require.ErrorIs(t, err, ErrTooLarge)
464 require.Equal(t, 4, src.Len(), "body must not be read at a zero limit")
465 })
466}
467
468// TestFramerRecvBoundedNeverWidensTheFramerCeiling proves a caller limit above
469// the framer's negotiated ceiling cannot raise it.
470func TestFramerRecvBoundedNeverWidensTheFramerCeiling(t *testing.T) {
471 var encoded bytes.Buffer
472 var header [4]byte
473 binary.BigEndian.PutUint32(header[:], 17)
474 encoded.Write(header[:])
475 encoded.Write(bytes.Repeat([]byte("z"), 17))
476
477 src := bytes.NewReader(encoded.Bytes())
478 receiver := NewFramer(src, nil, nil, WithMaxEnvelope(16))
479 defer func() { _ = receiver.Close() }()
480 _, err := receiver.RecvBounded(1 << 20)
481 require.ErrorIs(t, err, ErrTooLarge)
482 require.Equal(t, 17, src.Len(), "body must not be read above the framer ceiling")
483}
484
485// TestFramerRecvBoundedMatchesRecvByteForByte proves the bounded path returns
486// exactly what Recv returns for the same wire bytes.
487func TestFramerRecvBoundedMatchesRecvByteForByte(t *testing.T) {
488 payloads := [][]byte{[]byte("one"), bytes.Repeat([]byte("x"), 70000), []byte("three")}
489 var encoded bytes.Buffer
490 for _, payload := range payloads {
491 var header [4]byte
492 binary.BigEndian.PutUint32(header[:], uint32(len(payload)))
493 encoded.Write(header[:])
494 encoded.Write(payload)
495 }
496 plain := NewFramer(bytes.NewReader(encoded.Bytes()), nil, nil)
497 bounded := NewFramer(bytes.NewReader(encoded.Bytes()), nil, nil)
498 defer func() {
499 _ = plain.Close()
500 _ = bounded.Close()
501 }()
502 for _, want := range payloads {
503 got, err := plain.Recv()
504 require.NoError(t, err)
505 require.Equal(t, want, got)
506 got, err = bounded.RecvBounded(AbsoluteLimit)
507 require.NoError(t, err)
508 require.Equal(t, want, got)
509 }
510}
511
512// TestFramerRecvBoundedCloseUnblocks proves Close releases a blocked bounded
513// receive through the injected closer.
514func TestFramerRecvBoundedCloseUnblocks(t *testing.T) {
515 pr, pw := io.Pipe()
516 framer := NewFramer(pr, pw, func() error {
517 _ = pr.Close()
518 _ = pw.Close()
519 return nil
520 })
521 recvErr := make(chan error, 1)
522 go func() {
523 _, err := framer.RecvBounded(64)
524 recvErr <- err
525 }()
526 time.Sleep(50 * time.Millisecond)
527 require.NoError(t, framer.Close())
528 select {
529 case err := <-recvErr:
530 require.True(t, errors.Is(err, io.EOF) || errors.Is(err, io.ErrClosedPipe) || errors.Is(err, ErrClosed))
531 case <-time.After(5 * time.Second):
532 t.Fatal("RecvBounded was not unblocked by Close")
533 }
534}
535
536// TestFramerConcurrentRecvBoundedAndClose hammers bounded receives against a
537// concurrent Close under -race: every call returns and no receive is stranded.
538func TestFramerConcurrentRecvBoundedAndClose(t *testing.T) {
539 pr, pw := io.Pipe()
540 framer := NewFramer(pr, pw, func() error {
541 _ = pr.Close()
542 _ = pw.Close()
543 return nil
544 })
545 var wg sync.WaitGroup
546 for range 4 {
547 wg.Add(1)
548 go func() {
549 defer wg.Done()
550 _, _ = framer.RecvBounded(64)
551 }()
552 }
553 time.Sleep(20 * time.Millisecond)
554 require.NoError(t, framer.Close())
555 wg.Wait()
556}