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 "encoding/binary"
17 "errors"
18 "fmt"
19 "io"
20 "sync"
21
22 "google.golang.org/protobuf/proto"
23
24 "bnema.dev/ermite-runner/internal/execstream/execpb"
25)
26
27// MaxFrame is the largest accepted serialized frame.
28// It leaves room for the largest job result: outputs of 1 MiB in total (the
29// limit of ermite) plus the overhead of the encoding.
30const MaxFrame = 2 << 20
31
32// ChunkSize is the largest output payload placed in one Output frame.
33const ChunkSize = 32 * 1024
34
35// ErrFrameTooLarge is returned for a frame above MaxFrame, on either side.
36var ErrFrameTooLarge = errors.New("execstream: frame too large")
37
38// Conn reads and writes frames on a byte stream. Send is safe for concurrent
39// use; Recv must be called from one goroutine.
40type Conn struct {
41 rw io.ReadWriter
42 br *bufio.Reader
43 wmu sync.Mutex
44}
45
46// NewConn wraps a byte stream.
47func NewConn(rw io.ReadWriter) *Conn {
48 return &Conn{rw: rw, br: bufio.NewReaderSize(rw, 64*1024)}
49}
50
51// Send writes one frame.
52func (c *Conn) Send(f *execpb.Frame) error {
53 b, err := proto.Marshal(f)
54 if err != nil {
55 return fmt.Errorf("execstream: marshal: %w", err)
56 }
57 if len(b) > MaxFrame {
58 return ErrFrameTooLarge
59 }
60 buf := make([]byte, 4+len(b))
61 binary.BigEndian.PutUint32(buf, uint32(len(b)))
62 copy(buf[4:], b)
63 c.wmu.Lock()
64 defer c.wmu.Unlock()
65 _, err = c.rw.Write(buf)
66 return err
67}
68
69// Recv reads one frame. It returns io.EOF only on a clean end of stream between
70// frames; a stream cut inside a frame yields io.ErrUnexpectedEOF.
71func (c *Conn) Recv() (*execpb.Frame, error) {
72 var hdr [4]byte
73 if _, err := io.ReadFull(c.br, hdr[:]); err != nil {
74 return nil, err
75 }
76 n := binary.BigEndian.Uint32(hdr[:])
77 if n > MaxFrame {
78 return nil, ErrFrameTooLarge
79 }
80 b := make([]byte, n)
81 if _, err := io.ReadFull(c.br, b); err != nil {
82 if errors.Is(err, io.EOF) {
83 err = io.ErrUnexpectedEOF
84 }
85 return nil, err
86 }
87 f := &execpb.Frame{}
88 if err := proto.Unmarshal(b, f); err != nil {
89 return nil, fmt.Errorf("execstream: unmarshal: %w", err)
90 }
91 return f, nil
92}
93
94// Drive is the host side of a session: it waits for the executor's Hello, sends
95// run, passes every Output frame to onOutput, and returns the Exit frame.
96//
97// cancel, when non-nil, asks the executor to stop the script: the first value
98// received is sent as one Cancel frame (the executor escalates SIGINT, SIGTERM,
99// SIGKILL by itself); closing the channel without a value sends nothing. It is
100// deliberately not a context: cancelling the session's context would also abort
101// the reads that deliver the executor's Exit. Drive keeps waiting for Exit after
102// a Cancel; the caller bounds that wait by stopping the sandbox, which closes
103// the stream and makes Drive return an error.
104func Drive(rw io.ReadWriter, run *execpb.Run, cancel <-chan *execpb.Cancel, onOutput func(*execpb.Output)) (*execpb.Hello, *execpb.Exit, error) {
105 return DriveJob(rw, run, cancel, JobHandler{Output: onOutput})
106}
107
108// JobHandler receives the frames of a session. Nil functions ignore their
109// frames. The functions are called from the goroutine of DriveJob, in the
110// order of the stream.
111type JobHandler struct {
112 Hello func(*execpb.Hello)
113 Output func(*execpb.Output)
114 StepStarted func(*execpb.StepStarted)
115 StepCompleted func(*execpb.StepCompleted)
116 Log func(*execpb.Log)
117 JobResult func(*execpb.JobResult)
118}
119
120// DriveJob is Drive for a session that carries a job: besides Output frames it
121// delivers the step, log and result frames to h.
122func DriveJob(rw io.ReadWriter, run *execpb.Run, cancel <-chan *execpb.Cancel, h JobHandler) (*execpb.Hello, *execpb.Exit, error) {
123 c := NewConn(rw)
124 first, err := c.Recv()
125 if err != nil {
126 return nil, nil, fmt.Errorf("execstream: waiting for hello: %w", err)
127 }
128 hello := first.GetHello()
129 if hello == nil {
130 return nil, nil, fmt.Errorf("execstream: first frame is not hello")
131 }
132 if h.Hello != nil {
133 h.Hello(hello)
134 }
135 if err := c.Send(&execpb.Frame{Body: &execpb.Frame_Run{Run: run}}); err != nil {
136 return hello, nil, fmt.Errorf("execstream: sending run: %w", err)
137 }
138
139 done := make(chan struct{})
140 defer close(done)
141 if cancel != nil {
142 go func() {
143 select {
144 case cn, ok := <-cancel:
145 if ok {
146 if cn == nil {
147 cn = &execpb.Cancel{}
148 }
149 _ = c.Send(&execpb.Frame{Body: &execpb.Frame_Cancel{Cancel: cn}})
150 }
151 case <-done:
152 }
153 }()
154 }
155
156 for {
157 f, err := c.Recv()
158 if err != nil {
159 return hello, nil, fmt.Errorf("execstream: reading: %w", err)
160 }
161 switch b := f.Body.(type) {
162 case *execpb.Frame_Output:
163 if h.Output != nil {
164 h.Output(b.Output)
165 }
166 case *execpb.Frame_StepStarted:
167 if h.StepStarted != nil {
168 h.StepStarted(b.StepStarted)
169 }
170 case *execpb.Frame_StepCompleted:
171 if h.StepCompleted != nil {
172 h.StepCompleted(b.StepCompleted)
173 }
174 case *execpb.Frame_Log:
175 if h.Log != nil {
176 h.Log(b.Log)
177 }
178 case *execpb.Frame_JobResult:
179 if h.JobResult != nil {
180 h.JobResult(b.JobResult)
181 }
182 case *execpb.Frame_Exit:
183 return hello, b.Exit, nil
184 }
185 }
186}