1// Package streamframe carries length-delimited Protobuf envelopes over one
2// reliable ordered byte stream. It owns no application semantics: IPC, SSH
3// stdio, and QUIC adapters compose it for framing and share its bounds,
4// serialized writes, and close behavior.
5package streamframe
6
7import (
8 "encoding/binary"
9 "errors"
10 "io"
11 "sync"
12)
13
14// Framing bounds. The absolute ceiling never moves; per-connection
15// negotiated ceilings only lower it.
16const (
17 // AbsoluteLimit caps one serialized envelope on any stream.
18 AbsoluteLimit = 16 << 20
19 // DefaultQueuedMessages bounds frames accepted while a peer is slow.
20 DefaultQueuedMessages = 8
21 // DefaultQueuedBytes bounds queued payload bytes per connection.
22 DefaultQueuedBytes = 32 << 20
23)
24
25var (
26 ErrClosed = errors.New("streamframe: closed")
27 ErrBackpressure = errors.New("streamframe: queue full")
28 ErrTooLarge = errors.New("streamframe: envelope exceeds maximum length")
29 ErrZeroLength = errors.New("streamframe: zero-length envelope")
30)
31
32// Option tunes a Framer.
33type Option func(*Framer)
34
35// WithMaxEnvelope sets the negotiated per-connection ceiling. Values above
36// AbsoluteLimit are clamped; values below 1 disable the framer.
37func WithMaxEnvelope(limit uint64) Option {
38 return func(f *Framer) {
39 switch {
40 case limit == 0:
41 f.maxEnvelope = 0
42 case limit > AbsoluteLimit:
43 f.maxEnvelope = AbsoluteLimit
44 default:
45 f.maxEnvelope = int(limit)
46 }
47 }
48}
49
50// WithQueueBounds sets the bounded async admission budget.
51func WithQueueBounds(messages int, bytes uint64) Option {
52 return func(f *Framer) {
53 if messages > 0 {
54 f.queueMessages = messages
55 }
56 if bytes > 0 {
57 f.queueBytes = bytes
58 }
59 }
60}
61
62// Framer frames complete envelopes over r/w: 4-byte big-endian length then
63// payload, no type byte. One writer is serialized internally; exactly one
64// Recv caller at a time. Close is concurrent-safe and unblocks in-flight
65// Send and Recv, then invokes the injected closer exactly once so blocked
66// I/O on the underlying stream releases.
67type Framer struct {
68 r io.Reader
69 w io.Writer
70 closer func() error
71
72 maxEnvelope int
73 queueMessages int
74 queueBytes uint64
75
76 // readMu serializes Recv so concurrent receivers each get one
77 // complete envelope instead of interleaving header and body reads.
78 readMu sync.Mutex
79 // header is the length-prefix buffer, guarded by readMu.
80 header [4]byte
81
82 egressMu sync.Mutex
83 egressClosing bool
84 // egressSenders counts synchronous senders between admission and queue
85 // hand-off; the writer drains the queue only after they all leave.
86 egressSenders sync.WaitGroup
87 egress chan sendRequest
88 queuedBytes uint64
89 done chan struct{}
90 writerDone chan struct{}
91 closeOnce sync.Once
92 closeErr error
93}
94
95type sendRequest struct {
96 data []byte
97 done chan error // nil for async admission
98 // counted is the byte-budget charge admitted for this request.
99 // Synchronous sends rendezvous with the writer instead of
100 // charging the async byte budget, so only async requests carry
101 // a charge the write loop releases on dequeue.
102 counted uint64
103}
104
105// NewFramer builds a Framer over r/w. closer releases the underlying stream
106// and is invoked exactly once by Close; it may be nil. Blocked Recv calls
107// observe Close promptly only when closer releases the underlying read
108// (concrete adapters close their connection/stream there); otherwise Recv
109// returns when that source itself errors or closes.
110func NewFramer(r io.Reader, w io.Writer, closer func() error, opts ...Option) *Framer {
111 f := &Framer{
112 r: r,
113 w: w,
114 closer: closer,
115 maxEnvelope: AbsoluteLimit,
116 queueMessages: DefaultQueuedMessages,
117 queueBytes: DefaultQueuedBytes,
118 egress: make(chan sendRequest, DefaultQueuedMessages),
119 done: make(chan struct{}),
120 writerDone: make(chan struct{}),
121 }
122 for _, opt := range opts {
123 if opt != nil {
124 opt(f)
125 }
126 }
127 if cap(f.egress) < f.queueMessages {
128 f.egress = make(chan sendRequest, f.queueMessages)
129 }
130 go f.writeLoop()
131 return f
132}
133
134// Send copies payload into the bounded queue and waits for the wire attempt.
135// It reports the write result, ErrClosed, or ErrTooLarge. The synchronous
136// path rendezvous with the writer instead of charging the async byte
137// budget, so it carries no counted charge to release on dequeue.
138func (f *Framer) Send(payload []byte) error {
139 data, err := f.frame(payload)
140 if err != nil {
141 return err
142 }
143 result, _ := resultPool.Get().(chan error)
144 if result == nil {
145 result = make(chan error, 1)
146 }
147 if err := f.enqueueWait(sendRequest{data: data, done: result}); err != nil {
148 resultPool.Put(result)
149 return err
150 }
151 select {
152 case err := <-result:
153 resultPool.Put(result)
154 return err
155 case <-f.done:
156 // The writer still owns result and may complete it later, so it is
157 // not returned to the pool.
158 return ErrClosed
159 }
160}
161
162// resultPool recycles the one-slot completion channels of synchronous Send.
163// A channel is pooled only after its single result was received.
164var resultPool sync.Pool
165
166// SendAsync admits payload for ordered background transmission without
167// waiting for the write. It returns ErrBackpressure when the byte or message
168// budget is exhausted; ownership includes a payload copy.
169func (f *Framer) SendAsync(payload []byte) error {
170 data, err := f.frame(payload)
171 if err != nil {
172 return err
173 }
174 return f.enqueueAsync(sendRequest{data: data})
175}
176
177// Recv reads one complete envelope and returns a freshly allocated copy the
178// caller owns. Concurrent Recv calls are serialized; each receives one
179// complete envelope. It blocks until a full envelope arrives, the framer
180// closes (io.EOF), or an error occurs.
181func (f *Framer) Recv() ([]byte, error) {
182 f.readMu.Lock()
183 defer f.readMu.Unlock()
184 return f.recvLocked(f.negotiatedBound())
185}
186
187// RecvBounded reads one complete envelope like Recv, but validates the four-
188// byte length prefix against limit before allocating or reading the body, so a
189// peer cannot force an allocation larger than the caller's bound by
190// advertising a huge frame. The effective bound is the smaller of limit and
191// the framer's negotiated ceiling, so a per-call limit can only narrow the
192// framer; a zero limit refuses every non-empty frame. Concurrent calls are
193// serialized with Recv. The caller owns the returned copy.
194func (f *Framer) RecvBounded(limit uint64) ([]byte, error) {
195 f.readMu.Lock()
196 defer f.readMu.Unlock()
197 return f.recvLocked(min(f.negotiatedBound(), limit))
198}
199
200// negotiatedBound is the framer's configured per-connection ceiling, or 0
201// when the framer is disabled (which refuses every non-empty frame).
202func (f *Framer) negotiatedBound() uint64 {
203 if f.maxEnvelope <= 0 {
204 return 0
205 }
206 return uint64(f.maxEnvelope)
207}
208
209// recvLocked reads and validates one frame under readMu. The length prefix is
210// validated against bound before any body buffer exists, and the returned slice
211// is freshly allocated and owned by the caller.
212func (f *Framer) recvLocked(bound uint64) ([]byte, error) {
213 header := f.header[:]
214 if _, err := io.ReadFull(f.r, header); err != nil {
215 if f.isClosed() {
216 return nil, io.EOF
217 }
218 return nil, err
219 }
220 length := binary.BigEndian.Uint32(header)
221 if length == 0 {
222 return nil, ErrZeroLength
223 }
224 if uint64(length) > bound {
225 return nil, ErrTooLarge
226 }
227 payload := make([]byte, length)
228 if _, err := io.ReadFull(f.r, payload); err != nil {
229 // Best-effort erase: a truncated frame may hold a partial secret
230 // (for example a bounded bootstrap auth record).
231 clear(payload)
232 if f.isClosed() {
233 return nil, io.EOF
234 }
235 return nil, err
236 }
237 return payload, nil
238}
239
240// Close unblocks in-flight Send and Recv, fails queued payloads, waits for
241// the sole writer, then invokes the injected closer exactly once.
242func (f *Framer) Close() error {
243 f.closeOnce.Do(func() {
244 f.egressMu.Lock()
245 f.egressClosing = true
246 close(f.done)
247 f.egressMu.Unlock()
248 if f.closer != nil {
249 f.closeErr = f.closer()
250 }
251 <-f.writerDone
252 })
253 return f.closeErr
254}
255
256func (f *Framer) isClosed() bool {
257 select {
258 case <-f.done:
259 return true
260 default:
261 return false
262 }
263}
264
265func (f *Framer) frame(payload []byte) ([]byte, error) {
266 if f.maxEnvelope <= 0 {
267 return nil, ErrClosed
268 }
269 if len(payload) == 0 {
270 return nil, ErrZeroLength
271 }
272 if len(payload) > f.maxEnvelope {
273 return nil, ErrTooLarge
274 }
275 buf := make([]byte, 4+len(payload))
276 binary.BigEndian.PutUint32(buf[:4], uint32(len(payload)))
277 copy(buf[4:], payload)
278 return buf, nil
279}
280
281func (f *Framer) enqueueWait(req sendRequest) error {
282 f.egressMu.Lock()
283 if f.egressClosing {
284 f.egressMu.Unlock()
285 return ErrClosed
286 }
287 // Add happens under egressMu while egressClosing is false, so it always
288 // precedes the writer's Wait, which starts only after Close set it.
289 f.egressSenders.Add(1)
290 f.egressMu.Unlock()
291 defer f.egressSenders.Done()
292
293 select {
294 case <-f.done:
295 return ErrClosed
296 case f.egress <- req:
297 return nil
298 }
299}
300
301// withCountedCharge stamps the async byte-budget charge onto the request
302// before it enters the queue: channel sends copy the struct, so the charge
303// must be set prior to admission for the write loop to release it later.
304func withCountedCharge(req sendRequest) sendRequest {
305 req.counted = uint64(len(req.data))
306 return req
307}
308
309func (f *Framer) enqueueAsync(req sendRequest) error {
310 f.egressMu.Lock()
311 defer f.egressMu.Unlock()
312 if f.egressClosing {
313 return ErrClosed
314 }
315 if uint64(len(req.data)) > f.queueBytes-f.queuedBytes {
316 return ErrBackpressure
317 }
318 select {
319 case f.egress <- withCountedCharge(req):
320 f.queuedBytes += uint64(len(req.data))
321 return nil
322 default:
323 return ErrBackpressure
324 }
325}
326
327func (f *Framer) writeLoop() {
328 defer close(f.writerDone)
329 for {
330 select {
331 case req := <-f.egress:
332 f.egressMu.Lock()
333 f.queuedBytes -= min(f.queuedBytes, req.counted)
334 f.egressMu.Unlock()
335 f.completeWrite(req, writeAll(f.w, req.data))
336 case <-f.done:
337 f.egressSenders.Wait()
338 f.egressMu.Lock()
339 for {
340 select {
341 case req := <-f.egress:
342 f.completeWrite(req, ErrClosed)
343 default:
344 f.egressMu.Unlock()
345 return
346 }
347 }
348 }
349 }
350}
351
352func (f *Framer) completeWrite(req sendRequest, err error) {
353 if req.done != nil {
354 req.done <- err
355 }
356}
357
358func writeAll(w io.Writer, data []byte) error {
359 for len(data) > 0 {
360 n, err := w.Write(data)
361 if n > 0 {
362 data = data[n:]
363 }
364 if err != nil {
365 return err
366 }
367 if n == 0 {
368 return io.ErrShortWrite
369 }
370 }
371 return nil
372}