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 racing on the reused read buffer.
78 readMu sync.Mutex
79
80 egressMu sync.Mutex
81 egressClosing bool
82 egressSenders int
83 egressSendersDone chan struct{}
84 egress chan sendRequest
85 queuedBytes uint64
86 done chan struct{}
87 writerDone chan struct{}
88 closeOnce sync.Once
89 closeErr error
90
91 readBuf []byte
92}
93
94type sendRequest struct {
95 data []byte
96 done chan error // nil for async admission
97 // counted is the byte-budget charge admitted for this request.
98 // Synchronous sends rendezvous with the writer instead of
99 // charging the async byte budget, so only async requests carry
100 // a charge the write loop releases on dequeue.
101 counted uint64
102}
103
104// NewFramer builds a Framer over r/w. closer releases the underlying stream
105// and is invoked exactly once by Close; it may be nil. Blocked Recv calls
106// observe Close promptly only when closer releases the underlying read
107// (concrete adapters close their connection/stream there); otherwise Recv
108// returns when that source itself errors or closes.
109func NewFramer(r io.Reader, w io.Writer, closer func() error, opts ...Option) *Framer {
110 f := &Framer{
111 r: r,
112 w: w,
113 closer: closer,
114 maxEnvelope: AbsoluteLimit,
115 queueMessages: DefaultQueuedMessages,
116 queueBytes: DefaultQueuedBytes,
117 egress: make(chan sendRequest, DefaultQueuedMessages),
118 done: make(chan struct{}),
119 writerDone: make(chan struct{}),
120 }
121 for _, opt := range opts {
122 if opt != nil {
123 opt(f)
124 }
125 }
126 if cap(f.egress) < f.queueMessages {
127 f.egress = make(chan sendRequest, f.queueMessages)
128 }
129 go f.writeLoop()
130 return f
131}
132
133// Send copies payload into the bounded queue and waits for the wire attempt.
134// It reports the write result, ErrClosed, or ErrTooLarge. The synchronous
135// path rendezvous with the writer instead of charging the async byte
136// budget, so it carries no counted charge to release on dequeue.
137func (f *Framer) Send(payload []byte) error {
138 data, err := f.frame(payload)
139 if err != nil {
140 return err
141 }
142 result := make(chan error, 1)
143 if err := f.enqueueWait(sendRequest{data: data, done: result}); err != nil {
144 return err
145 }
146 select {
147 case err := <-result:
148 return err
149 case <-f.done:
150 return ErrClosed
151 }
152}
153
154// SendAsync admits payload for ordered background transmission without
155// waiting for the write. It returns ErrBackpressure when the byte or message
156// budget is exhausted; ownership includes a payload copy.
157func (f *Framer) SendAsync(payload []byte) error {
158 data, err := f.frame(payload)
159 if err != nil {
160 return err
161 }
162 return f.enqueueAsync(sendRequest{data: data})
163}
164
165// Recv reads one complete envelope and returns a freshly allocated copy the
166// caller owns. Concurrent Recv calls are serialized; each receives one
167// complete envelope. It blocks until a full envelope arrives, the framer
168// closes (io.EOF), or an error occurs.
169func (f *Framer) Recv() ([]byte, error) {
170 f.readMu.Lock()
171 defer f.readMu.Unlock()
172 return f.recvLocked(f.negotiatedBound())
173}
174
175// RecvBounded reads one complete envelope like Recv, but validates the four-
176// byte length prefix against limit before allocating or reading the body, so a
177// peer cannot force an allocation larger than the caller's bound by
178// advertising a huge frame. The effective bound is the smaller of limit and
179// the framer's negotiated ceiling, so a per-call limit can only narrow the
180// framer; a zero limit refuses every non-empty frame. Concurrent calls are
181// serialized with Recv. The caller owns the returned copy.
182func (f *Framer) RecvBounded(limit uint64) ([]byte, error) {
183 f.readMu.Lock()
184 defer f.readMu.Unlock()
185 return f.recvLocked(min(f.negotiatedBound(), limit))
186}
187
188// negotiatedBound is the framer's configured per-connection ceiling, or 0
189// when the framer is disabled (which refuses every non-empty frame).
190func (f *Framer) negotiatedBound() uint64 {
191 if f.maxEnvelope <= 0 {
192 return 0
193 }
194 return uint64(f.maxEnvelope)
195}
196
197// recvLocked reads and validates one frame under readMu. The length prefix is
198// validated against bound before any body buffer exists, and the returned slice
199// is a fresh copy the caller owns.
200func (f *Framer) recvLocked(bound uint64) ([]byte, error) {
201 var header [4]byte
202 if _, err := io.ReadFull(f.r, header[:]); err != nil {
203 if f.isClosed() {
204 return nil, io.EOF
205 }
206 return nil, err
207 }
208 length := binary.BigEndian.Uint32(header[:])
209 if length == 0 {
210 return nil, ErrZeroLength
211 }
212 if uint64(length) > bound {
213 return nil, ErrTooLarge
214 }
215 if cap(f.readBuf) < int(length) {
216 f.readBuf = make([]byte, length)
217 } else {
218 f.readBuf = f.readBuf[:length]
219 }
220 payload := f.readBuf
221 if _, err := io.ReadFull(f.r, payload); err != nil {
222 // Best-effort erase: a truncated frame may hold a partial secret
223 // (for example a bounded bootstrap auth record).
224 clear(payload)
225 if f.isClosed() {
226 return nil, io.EOF
227 }
228 return nil, err
229 }
230 return append([]byte(nil), payload...), nil
231}
232
233// Close unblocks in-flight Send and Recv, fails queued payloads, waits for
234// the sole writer, then invokes the injected closer exactly once.
235func (f *Framer) Close() error {
236 f.closeOnce.Do(func() {
237 f.egressMu.Lock()
238 f.egressClosing = true
239 close(f.done)
240 f.egressMu.Unlock()
241 if f.closer != nil {
242 f.closeErr = f.closer()
243 }
244 <-f.writerDone
245 })
246 return f.closeErr
247}
248
249func (f *Framer) isClosed() bool {
250 select {
251 case <-f.done:
252 return true
253 default:
254 return false
255 }
256}
257
258func (f *Framer) frame(payload []byte) ([]byte, error) {
259 if f.maxEnvelope <= 0 {
260 return nil, ErrClosed
261 }
262 if len(payload) == 0 {
263 return nil, ErrZeroLength
264 }
265 if len(payload) > f.maxEnvelope {
266 return nil, ErrTooLarge
267 }
268 buf := make([]byte, 4+len(payload))
269 binary.BigEndian.PutUint32(buf[:4], uint32(len(payload)))
270 copy(buf[4:], payload)
271 return buf, nil
272}
273
274func (f *Framer) enqueueWait(req sendRequest) error {
275 f.egressMu.Lock()
276 if f.egressClosing {
277 f.egressMu.Unlock()
278 return ErrClosed
279 }
280 if f.egressSenders == 0 {
281 f.egressSendersDone = make(chan struct{})
282 }
283 f.egressSenders++
284 f.egressMu.Unlock()
285
286 select {
287 case <-f.done:
288 f.finishEgressSender()
289 return ErrClosed
290 case f.egress <- req:
291 f.finishEgressSender()
292 return nil
293 }
294}
295
296// withCountedCharge stamps the async byte-budget charge onto the request
297// before it enters the queue: channel sends copy the struct, so the charge
298// must be set prior to admission for the write loop to release it later.
299func withCountedCharge(req sendRequest) sendRequest {
300 req.counted = uint64(len(req.data))
301 return req
302}
303
304func (f *Framer) enqueueAsync(req sendRequest) error {
305 f.egressMu.Lock()
306 defer f.egressMu.Unlock()
307 if f.egressClosing {
308 return ErrClosed
309 }
310 if uint64(len(req.data)) > f.queueBytes-f.queuedBytes {
311 return ErrBackpressure
312 }
313 select {
314 case f.egress <- withCountedCharge(req):
315 f.queuedBytes += uint64(len(req.data))
316 return nil
317 default:
318 return ErrBackpressure
319 }
320}
321
322func (f *Framer) finishEgressSender() {
323 f.egressMu.Lock()
324 defer f.egressMu.Unlock()
325 f.egressSenders--
326 if f.egressSenders == 0 {
327 close(f.egressSendersDone)
328 }
329}
330
331func (f *Framer) waitForEgressSenders() {
332 f.egressMu.Lock()
333 done := f.egressSendersDone
334 f.egressMu.Unlock()
335 if done != nil {
336 <-done
337 }
338}
339
340func (f *Framer) writeLoop() {
341 defer close(f.writerDone)
342 for {
343 select {
344 case req := <-f.egress:
345 f.egressMu.Lock()
346 f.queuedBytes -= min(f.queuedBytes, req.counted)
347 f.egressMu.Unlock()
348 f.completeWrite(req, writeAll(f.w, req.data))
349 case <-f.done:
350 f.waitForEgressSenders()
351 f.egressMu.Lock()
352 for {
353 select {
354 case req := <-f.egress:
355 f.completeWrite(req, ErrClosed)
356 default:
357 f.egressMu.Unlock()
358 return
359 }
360 }
361 }
362 }
363}
364
365func (f *Framer) completeWrite(req sendRequest, err error) {
366 if req.done != nil {
367 req.done <- err
368 }
369}
370
371func writeAll(w io.Writer, data []byte) error {
372 for len(data) > 0 {
373 n, err := w.Write(data)
374 if n > 0 {
375 data = data[n:]
376 }
377 if err != nil {
378 return err
379 }
380 if n == 0 {
381 return io.ErrShortWrite
382 }
383 }
384 return nil
385}