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