1// SPDX-License-Identifier: MIT
2
3// Package execstream is the framed protobuf stream between the host process of
4// ermite-runner and the executor running inside a sandbox.
5//
6// Wire format: each frame is a 4-byte big-endian length followed by that many
7// bytes of a serialized execpb.Frame (see exec.proto). A frame is at most
8// MaxFrame bytes; longer output is split by the sender. The stream is
9// transport-agnostic: vsock in a micro-VM, a Unix socket in a container.
10package execstream
11
12//go:generate go tool buf generate ../.. --template ../../buf.gen.yaml
13
14import (
15 "bufio"
16 "context"
17 "encoding/binary"
18 "errors"
19 "fmt"
20 "io"
21 "sync"
22 "time"
23
24 "google.golang.org/protobuf/proto"
25
26 "bnema.dev/ermite-runner/internal/execstream/execpb"
27)
28
29// MaxFrame is the largest accepted serialized frame.
30// It leaves room for the largest job result: outputs of 1 MiB in total (the
31// limit of ermite) plus the overhead of the encoding.
32const MaxFrame = 2 << 20
33
34// ChunkSize is the largest output payload placed in one Output frame.
35const ChunkSize = 32 * 1024
36
37// ErrFrameTooLarge is returned for a frame above MaxFrame, on either side.
38var ErrFrameTooLarge = errors.New("execstream: frame too large")
39
40// Conn reads and writes frames on a byte stream. Send is safe for concurrent
41// use; Recv must be called from one goroutine.
42type Conn struct {
43 rw io.ReadWriter
44 br *bufio.Reader
45 wmu sync.Mutex
46}
47
48// NewConn wraps a byte stream.
49func NewConn(rw io.ReadWriter) *Conn {
50 return &Conn{rw: rw, br: bufio.NewReaderSize(rw, 64*1024)}
51}
52
53// Send writes one frame.
54func (c *Conn) Send(f *execpb.Frame) error {
55 b, err := proto.Marshal(f)
56 if err != nil {
57 return fmt.Errorf("execstream: marshal: %w", err)
58 }
59 if len(b) > MaxFrame {
60 return ErrFrameTooLarge
61 }
62 buf := make([]byte, 4+len(b))
63 binary.BigEndian.PutUint32(buf, uint32(len(b)))
64 copy(buf[4:], b)
65 c.wmu.Lock()
66 defer c.wmu.Unlock()
67 _, err = c.rw.Write(buf)
68 return err
69}
70
71// Recv reads one frame. It returns io.EOF only on a clean end of stream between
72// frames; a stream cut inside a frame yields io.ErrUnexpectedEOF.
73func (c *Conn) Recv() (*execpb.Frame, error) {
74 var hdr [4]byte
75 if _, err := io.ReadFull(c.br, hdr[:]); err != nil {
76 return nil, err
77 }
78 n := binary.BigEndian.Uint32(hdr[:])
79 if n > MaxFrame {
80 return nil, ErrFrameTooLarge
81 }
82 b := make([]byte, n)
83 if _, err := io.ReadFull(c.br, b); err != nil {
84 if errors.Is(err, io.EOF) {
85 err = io.ErrUnexpectedEOF
86 }
87 return nil, err
88 }
89 f := &execpb.Frame{}
90 if err := proto.Unmarshal(b, f); err != nil {
91 return nil, fmt.Errorf("execstream: unmarshal: %w", err)
92 }
93 return f, nil
94}
95
96// Drive is the host side of a session: it waits for the executor's Hello, sends
97// run, passes every Output frame to onOutput, and returns the Exit frame.
98//
99// cancel, when non-nil, asks the executor to stop the script: the first value
100// received is sent as one Cancel frame (the executor escalates SIGINT, SIGTERM,
101// SIGKILL by itself); closing the channel without a value sends nothing. It is
102// deliberately not a context: cancelling the session's context would also abort
103// the reads that deliver the executor's Exit. Drive keeps waiting for Exit after
104// a Cancel; the caller bounds that wait by stopping the sandbox, which closes
105// the stream and makes Drive return an error.
106func Drive(rw io.ReadWriter, run *execpb.Run, cancel <-chan *execpb.Cancel, onOutput func(*execpb.Output)) (*execpb.Hello, *execpb.Exit, error) {
107 return DriveJob(rw, run, cancel, JobHandler{Output: onOutput})
108}
109
110// resolveJoinGrace is how long DriveJob waits for the resolution worker to
111// finish once the session is over before it interrupts a write that is stuck
112// (a guest that stopped reading): the write deadline is set to now, and the
113// stream is closed if it cannot be interrupted otherwise. The session is over at
114// that point, the caller closes the stream anyway.
115var resolveJoinGrace = 2 * time.Second
116
117// joinResolver waits for the worker. A worker in a Send that cannot finish
118// because the peer does not read would block DriveJob for ever, so the write is
119// interrupted after the grace period.
120func joinResolver(rw io.ReadWriter, workers *sync.WaitGroup) {
121 finished := make(chan struct{})
122 go func() { workers.Wait(); close(finished) }()
123 select {
124 case <-finished:
125 return
126 case <-time.After(resolveJoinGrace):
127 }
128 if d, ok := rw.(interface{ SetWriteDeadline(time.Time) error }); ok {
129 _ = d.SetWriteDeadline(time.Now())
130 }
131 select {
132 case <-finished:
133 return
134 case <-time.After(resolveJoinGrace):
135 }
136 if cl, ok := rw.(io.Closer); ok {
137 _ = cl.Close()
138 <-finished
139 }
140}
141
142// resolveQueue is how many ResolveRequest frames may wait for the worker.
143const resolveQueue = 2
144
145func resolveFrame(r *execpb.ResolveResponse) *execpb.Frame {
146 return &execpb.Frame{Body: &execpb.Frame_ResolveResponse{ResolveResponse: r}}
147}
148
149// JobHandler receives the frames of a session. Nil functions ignore their
150// frames. The functions are called from the goroutine of DriveJob, in the
151// order of the stream.
152type JobHandler struct {
153 Hello func(*execpb.Hello)
154 Output func(*execpb.Output)
155 StepStarted func(*execpb.StepStarted)
156 StepCompleted func(*execpb.StepCompleted)
157 Log func(*execpb.Log)
158 JobResult func(*execpb.JobResult)
159 // Resolve answers the executor's ResolveRequest frames, one at a time and in
160 // order, from a single worker (the call goes to the server); the answer is
161 // sent as a ResolveResponse frame carrying the id of the request. ctx ends
162 // with the session: a call still running is cancelled and DriveJob waits for
163 // it before it returns. At most resolveQueue requests wait; the next ones are
164 // answered with an error at once, without a call.
165 Resolve func(ctx context.Context, r *execpb.ResolveRequest) *execpb.ResolveResponse
166 // StepAnnotations and StepSummary come before the StepCompleted of their
167 // step.
168 StepAnnotations func(*execpb.StepAnnotations)
169 StepSummary func(*execpb.StepSummary)
170}
171
172// DriveJob is Drive for a session that carries a job: besides Output frames it
173// delivers the step, log and result frames to h.
174func DriveJob(rw io.ReadWriter, run *execpb.Run, cancel <-chan *execpb.Cancel, h JobHandler) (*execpb.Hello, *execpb.Exit, error) {
175 c := NewConn(rw)
176 first, err := c.Recv()
177 if err != nil {
178 return nil, nil, fmt.Errorf("execstream: waiting for hello: %w", err)
179 }
180 hello := first.GetHello()
181 if hello == nil {
182 return nil, nil, fmt.Errorf("execstream: first frame is not hello")
183 }
184 if h.Hello != nil {
185 h.Hello(hello)
186 }
187 if err := c.Send(&execpb.Frame{Body: &execpb.Frame_Run{Run: run}}); err != nil {
188 return hello, nil, fmt.Errorf("execstream: sending run: %w", err)
189 }
190
191 done := make(chan struct{})
192 defer close(done)
193
194 // One worker answers the resolutions; it ends, and is waited for, with the
195 // session.
196 resolveCh := make(chan *execpb.ResolveRequest, resolveQueue)
197 rctx, rcancel := context.WithCancel(context.Background())
198 var workers sync.WaitGroup
199 workers.Add(1)
200 go func() {
201 defer workers.Done()
202 for {
203 select {
204 case req := <-resolveCh:
205 resp := h.Resolve(rctx, req)
206 resp.Id = req.Id
207 if rctx.Err() != nil {
208 return // the session is over: nobody reads the answer
209 }
210 _ = c.Send(resolveFrame(resp))
211 case <-rctx.Done():
212 return
213 }
214 }
215 }()
216 defer func() { rcancel(); joinResolver(rw, &workers) }()
217 if cancel != nil {
218 go func() {
219 select {
220 case cn, ok := <-cancel:
221 if ok {
222 if cn == nil {
223 cn = &execpb.Cancel{}
224 }
225 _ = c.Send(&execpb.Frame{Body: &execpb.Frame_Cancel{Cancel: cn}})
226 }
227 case <-done:
228 }
229 }()
230 }
231
232 for {
233 f, err := c.Recv()
234 if err != nil {
235 return hello, nil, fmt.Errorf("execstream: reading: %w", err)
236 }
237 switch b := f.Body.(type) {
238 case *execpb.Frame_Output:
239 if h.Output != nil {
240 h.Output(b.Output)
241 }
242 case *execpb.Frame_StepStarted:
243 if h.StepStarted != nil {
244 h.StepStarted(b.StepStarted)
245 }
246 case *execpb.Frame_StepCompleted:
247 if h.StepCompleted != nil {
248 h.StepCompleted(b.StepCompleted)
249 }
250 case *execpb.Frame_Log:
251 if h.Log != nil {
252 h.Log(b.Log)
253 }
254 case *execpb.Frame_JobResult:
255 if h.JobResult != nil {
256 h.JobResult(b.JobResult)
257 }
258 case *execpb.Frame_ResolveRequest:
259 req := b.ResolveRequest
260 if h.Resolve == nil {
261 _ = c.Send(resolveFrame(&execpb.ResolveResponse{Id: req.Id, Error: "the host cannot resolve actions"}))
262 break
263 }
264 select {
265 case resolveCh <- req:
266 default:
267 _ = c.Send(resolveFrame(&execpb.ResolveResponse{Id: req.Id, Error: "too many action resolutions in flight"}))
268 }
269 case *execpb.Frame_StepAnnotations:
270 if h.StepAnnotations != nil {
271 h.StepAnnotations(b.StepAnnotations)
272 }
273 case *execpb.Frame_StepSummary:
274 if h.StepSummary != nil {
275 h.StepSummary(b.StepSummary)
276 }
277 case *execpb.Frame_Exit:
278 return hello, b.Exit, nil
279 }
280 }
281}