1package extensions
2
3import (
4 "context"
5 "errors"
6 "fmt"
7 "io"
8 "sync"
9 "time"
10)
11
12// orderedPipe serializes every host→extension frame through one writer
13// goroutine so a blocked subprocess stdin never pins a host goroutine for an
14// unbounded time. Frames are delivered in enqueue order. Synchronous callers
15// wait for their frame's write with a bounded timeout or context; fire-and-
16// forget callers only pay for the enqueue.
17//
18// The queue is bounded in frames and bytes. Overflow disconnects the
19// extension transport instead of growing without limit or silently dropping
20// a frame in the middle of an ordered stream.
21type orderedPipe struct {
22 pipe io.WriteCloser
23 mu sync.Mutex
24 queue []*outboundFrame
25 bytes int
26 closed bool
27 wake chan struct{}
28 done chan struct{}
29 once sync.Once
30}
31
32type outboundFrame struct {
33 data []byte
34 ack chan error
35 // writing is set once the writer goroutine has taken ownership of the
36 // frame. Until then the frame may still be withdrawn safely.
37 writing bool
38}
39
40const (
41 outboundBytes = 16 * 1024 * 1024
42 outboundFrames = 256
43
44 // transportWriteTimeout bounds how long a synchronous host frame may wait
45 // for the extension to drain its stdin. It is independent of any reply
46 // deadline: an interactive tool may wait forever for a reply, but the
47 // request itself must be accepted by the subprocess promptly.
48 transportWriteTimeout = 5 * time.Second
49
50 // abandonGrace is how long a cancelled caller waits for a frame that is
51 // already mid-write to finish before the transport is disconnected. A
52 // torn frame followed by unrelated bytes would corrupt the stream, so
53 // closing is the only safe fallback.
54 abandonGrace = 250 * time.Millisecond
55)
56
57var errOutboundOverflow = errors.New("extension outbound queue exceeded its limit")
58
59func newOrderedPipe(pipe io.WriteCloser) *orderedPipe {
60 p := &orderedPipe{pipe: pipe, wake: make(chan struct{}, 1), done: make(chan struct{})}
61 go p.run()
62 return p
63}
64
65// enqueue appends a frame and returns a handle for withdrawal. ack, when not
66// nil, receives exactly one write outcome unless the frame is withdrawn.
67func (p *orderedPipe) enqueue(data []byte, ack chan error) (*outboundFrame, error) {
68 p.mu.Lock()
69 if p.closed {
70 p.mu.Unlock()
71 return nil, io.ErrClosedPipe
72 }
73 if len(p.queue) >= outboundFrames || p.bytes+len(data) > outboundBytes {
74 p.mu.Unlock()
75 p.Close()
76 return nil, errOutboundOverflow
77 }
78 frame := &outboundFrame{data: append([]byte(nil), data...), ack: ack}
79 p.queue = append(p.queue, frame)
80 p.bytes += len(data)
81 p.mu.Unlock()
82 select {
83 case p.wake <- struct{}{}:
84 default:
85 }
86 return frame, nil
87}
88
89// withdraw removes a frame that has not yet started writing. It reports
90// false when the writer already owns the frame; in that case the bytes may be
91// partially on the wire and cannot be taken back.
92func (p *orderedPipe) withdraw(frame *outboundFrame) bool {
93 p.mu.Lock()
94 defer p.mu.Unlock()
95 if frame.writing {
96 return false
97 }
98 for i, queued := range p.queue {
99 if queued == frame {
100 p.queue = append(p.queue[:i], p.queue[i+1:]...)
101 p.bytes -= len(frame.data)
102 return true
103 }
104 }
105 // Not queued and not writing: already written or the pipe was closed.
106 return true
107}
108
109// Write enqueues and waits for the write with the standard transport bound.
110func (p *orderedPipe) Write(data []byte) (int, error) {
111 return p.writeContext(context.Background(), data, transportWriteTimeout)
112}
113
114// writeContext enqueues data and waits until it is written, the timeout
115// elapses, ctx ends, or the pipe closes. A frame still waiting in the queue
116// is withdrawn on cancellation without disturbing the transport; a frame
117// already being written must finish within the cancellation grace or the
118// transport is closed to avoid a torn frame. A write timeout disconnects the
119// transport whether the frame is queued or in progress.
120func (p *orderedPipe) writeContext(ctx context.Context, data []byte, timeout time.Duration) (int, error) {
121 if err := ctx.Err(); err != nil {
122 return 0, err
123 }
124 if timeout <= 0 {
125 return 0, fmt.Errorf("extension outbound write timed out: %w", context.DeadlineExceeded)
126 }
127 ack := make(chan error, 1)
128 frame, err := p.enqueue(data, ack)
129 if err != nil {
130 return 0, err
131 }
132 timer := time.NewTimer(timeout)
133 defer timer.Stop()
134 abandon := func(reason error) (int, error) {
135 if p.withdraw(frame) {
136 return 0, reason
137 }
138 // Mid-write: wait briefly for the writer to finish this frame so a
139 // complete frame lands or the pipe error surfaces; otherwise close.
140 select {
141 case err := <-ack:
142 if err != nil {
143 return 0, err
144 }
145 return len(data), nil
146 case <-p.done:
147 return 0, io.ErrClosedPipe
148 case <-time.After(abandonGrace):
149 p.Close()
150 return 0, reason
151 }
152 }
153 select {
154 case err := <-ack:
155 if err != nil {
156 return 0, err
157 }
158 return len(data), nil
159 case <-p.done:
160 // The frame may have completed just before the disconnect. If the
161 // writer owns it, its outcome arrives promptly because the closed
162 // pipe fails any in-progress write; report that outcome accurately so
163 // callers distinguish "request delivered, peer gone" from "never sent".
164 if !p.withdraw(frame) {
165 if err := <-ack; err != nil {
166 return 0, err
167 }
168 return len(data), nil
169 }
170 select {
171 case err := <-ack:
172 if err == nil {
173 return len(data), nil
174 }
175 return 0, err
176 default:
177 }
178 return 0, io.ErrClosedPipe
179 case <-ctx.Done():
180 return abandon(ctx.Err())
181 case <-timer.C:
182 p.Close()
183 return 0, fmt.Errorf("extension outbound write timed out: %w", context.DeadlineExceeded)
184 }
185}
186
187// Close disconnects the transport. Queued frames are discarded and every
188// waiter observes done. Safe to call more than once.
189func (p *orderedPipe) Close() error {
190 p.once.Do(func() {
191 p.mu.Lock()
192 p.closed = true
193 p.queue = nil
194 p.bytes = 0
195 p.mu.Unlock()
196 close(p.done)
197 _ = p.pipe.Close()
198 })
199 return nil
200}
201
202// Done is closed when the transport has been disconnected.
203func (p *orderedPipe) Done() <-chan struct{} { return p.done }
204
205func (p *orderedPipe) run() {
206 for {
207 select {
208 case <-p.done:
209 return
210 case <-p.wake:
211 }
212 for {
213 p.mu.Lock()
214 if p.closed || len(p.queue) == 0 {
215 p.mu.Unlock()
216 break
217 }
218 frame := p.queue[0]
219 frame.writing = true
220 p.queue[0] = nil
221 p.queue = p.queue[1:]
222 p.bytes -= len(frame.data)
223 p.mu.Unlock()
224 n, err := p.pipe.Write(frame.data)
225 if err == nil && n != len(frame.data) {
226 err = io.ErrShortWrite
227 }
228 if frame.ack != nil {
229 frame.ack <- err
230 }
231 if err != nil {
232 p.Close()
233 return
234 }
235 }
236 }
237}