feat(execstream): add the framed protobuf stream between host and executor
- commit
- b44fbc45e98d9a8809b7043b5f8bc04e11b0b655
- author
- brice · · 2026-10-12 15:35 UTC
- parent
- c37a8f5
- signature
- GPG Verified
CD88F9E564E0102646FB6077A1A07C7C7C3E349D· arch-gpg Verified for the committer email
4 files · +915 -0
A internal/execstream/exec.proto +58 -0
| @@ -0,0 +1,58 @@ | |||
| 1 | + | // SPDX-License-Identifier: MIT | |
| 2 | + | ||
| 3 | + | // Internal wire format between the host process of ermite-runner and the | |
| 4 | + | // executor running inside a sandbox (vsock in a micro-VM, Unix socket in a | |
| 5 | + | // container). Frames are length-prefixed on the stream (see frame.go); this file | |
| 6 | + | // defines the payload. It is internal to ermite-runner: no compatibility is | |
| 7 | + | // promised between runner versions, which is safe because the host attaches its | |
| 8 | + | // own executable to every sandbox. | |
| 9 | + | ||
| 10 | + | syntax = "proto3"; | |
| 11 | + | ||
| 12 | + | package ermite.runner.exec.v1; | |
| 13 | + | ||
| 14 | + | option go_package = "bnema.dev/ermite-runner/internal/execstream/execpb"; | |
| 15 | + | ||
| 16 | + | // Frame is the only message on the stream. | |
| 17 | + | message Frame { | |
| 18 | + | oneof body { | |
| 19 | + | Hello hello = 1; // executor -> host, first frame | |
| 20 | + | Run run = 2; // host -> executor, after Hello | |
| 21 | + | Output output = 3; // executor -> host | |
| 22 | + | Exit exit = 4; // executor -> host, last frame | |
| 23 | + | Cancel cancel = 5; // host -> executor | |
| 24 | + | } | |
| 25 | + | } | |
| 26 | + | ||
| 27 | + | message Hello { | |
| 28 | + | string version = 1; // runner version of the executor | |
| 29 | + | string sandbox = 2; // "vm" or "container" | |
| 30 | + | string hostname = 3; | |
| 31 | + | } | |
| 32 | + | ||
| 33 | + | // Run asks the executor to run a script. CI-1 stub executor: the script is run | |
| 34 | + | // with `sh -c`; an empty script means the built-in diagnostics script. | |
| 35 | + | message Run { | |
| 36 | + | string script = 1; | |
| 37 | + | map<string, string> env = 2; | |
| 38 | + | } | |
| 39 | + | ||
| 40 | + | enum StreamKind { | |
| 41 | + | STREAM_KIND_UNSPECIFIED = 0; | |
| 42 | + | STREAM_KIND_STDOUT = 1; | |
| 43 | + | STREAM_KIND_STDERR = 2; | |
| 44 | + | } | |
| 45 | + | ||
| 46 | + | message Output { | |
| 47 | + | StreamKind stream = 1; | |
| 48 | + | bytes data = 2; | |
| 49 | + | int64 unix_nano = 3; | |
| 50 | + | } | |
| 51 | + | ||
| 52 | + | message Exit { | |
| 53 | + | int32 code = 1; | |
| 54 | + | string error = 2; // set when the executor failed to run the script at all | |
| 55 | + | } | |
| 56 | + | ||
| 57 | + | // Cancel asks the executor to stop the running script: SIGINT now. | |
| 58 | + | message Cancel {} | |
| @@ -0,0 +1,58 @@ | |||||
| 1 | + | // SPDX-License-Identifier: MIT | |||
| 2 | + | ||||
| 3 | + | // Internal wire format between the host process of ermite-runner and the | |||
| 4 | + | // executor running inside a sandbox (vsock in a micro-VM, Unix socket in a | |||
| 5 | + | // container). Frames are length-prefixed on the stream (see frame.go); this file | |||
| 6 | + | // defines the payload. It is internal to ermite-runner: no compatibility is | |||
| 7 | + | // promised between runner versions, which is safe because the host attaches its | |||
| 8 | + | // own executable to every sandbox. | |||
| 9 | + | ||||
| 10 | + | syntax = "proto3"; | |||
| 11 | + | ||||
| 12 | + | package ermite.runner.exec.v1; | |||
| 13 | + | ||||
| 14 | + | option go_package = "bnema.dev/ermite-runner/internal/execstream/execpb"; | |||
| 15 | + | ||||
| 16 | + | // Frame is the only message on the stream. | |||
| 17 | + | message Frame { | |||
| 18 | + | oneof body { | |||
| 19 | + | Hello hello = 1; // executor -> host, first frame | |||
| 20 | + | Run run = 2; // host -> executor, after Hello | |||
| 21 | + | Output output = 3; // executor -> host | |||
| 22 | + | Exit exit = 4; // executor -> host, last frame | |||
| 23 | + | Cancel cancel = 5; // host -> executor | |||
| 24 | + | } | |||
| 25 | + | } | |||
| 26 | + | ||||
| 27 | + | message Hello { | |||
| 28 | + | string version = 1; // runner version of the executor | |||
| 29 | + | string sandbox = 2; // "vm" or "container" | |||
| 30 | + | string hostname = 3; | |||
| 31 | + | } | |||
| 32 | + | ||||
| 33 | + | // Run asks the executor to run a script. CI-1 stub executor: the script is run | |||
| 34 | + | // with `sh -c`; an empty script means the built-in diagnostics script. | |||
| 35 | + | message Run { | |||
| 36 | + | string script = 1; | |||
| 37 | + | map<string, string> env = 2; | |||
| 38 | + | } | |||
| 39 | + | ||||
| 40 | + | enum StreamKind { | |||
| 41 | + | STREAM_KIND_UNSPECIFIED = 0; | |||
| 42 | + | STREAM_KIND_STDOUT = 1; | |||
| 43 | + | STREAM_KIND_STDERR = 2; | |||
| 44 | + | } | |||
| 45 | + | ||||
| 46 | + | message Output { | |||
| 47 | + | StreamKind stream = 1; | |||
| 48 | + | bytes data = 2; | |||
| 49 | + | int64 unix_nano = 3; | |||
| 50 | + | } | |||
| 51 | + | ||||
| 52 | + | message Exit { | |||
| 53 | + | int32 code = 1; | |||
| 54 | + | string error = 2; // set when the executor failed to run the script at all | |||
| 55 | + | } | |||
| 56 | + | ||||
| 57 | + | // Cancel asks the executor to stop the running script: SIGINT now. | |||
| 58 | + | message Cancel {} | |||
A internal/execstream/execpb/exec.pb.go +582 -0 generated file: collapsed
A internal/execstream/frame.go +136 -0
| @@ -0,0 +1,136 @@ | |||
| 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. | |
| 10 | + | package execstream | |
| 11 | + | ||
| 12 | + | //go:generate protoc --proto_path=. --go_out=execpb --go_opt=paths=source_relative exec.proto | |
| 13 | + | ||
| 14 | + | import ( | |
| 15 | + | "bufio" | |
| 16 | + | "context" | |
| 17 | + | "encoding/binary" | |
| 18 | + | "errors" | |
| 19 | + | "fmt" | |
| 20 | + | "io" | |
| 21 | + | "sync" | |
| 22 | + | ||
| 23 | + | "google.golang.org/protobuf/proto" | |
| 24 | + | ||
| 25 | + | "bnema.dev/ermite-runner/internal/execstream/execpb" | |
| 26 | + | ) | |
| 27 | + | ||
| 28 | + | // MaxFrame is the largest accepted serialized frame. | |
| 29 | + | const MaxFrame = 1 << 20 | |
| 30 | + | ||
| 31 | + | // ChunkSize is the largest output payload placed in one Output frame. | |
| 32 | + | const ChunkSize = 32 * 1024 | |
| 33 | + | ||
| 34 | + | // ErrFrameTooLarge is returned for a frame above MaxFrame, on either side. | |
| 35 | + | var ErrFrameTooLarge = errors.New("execstream: frame too large") | |
| 36 | + | ||
| 37 | + | // Conn reads and writes frames on a byte stream. Send is safe for concurrent | |
| 38 | + | // use; Recv must be called from one goroutine. | |
| 39 | + | type Conn struct { | |
| 40 | + | rw io.ReadWriter | |
| 41 | + | br *bufio.Reader | |
| 42 | + | wmu sync.Mutex | |
| 43 | + | } | |
| 44 | + | ||
| 45 | + | // NewConn wraps a byte stream. | |
| 46 | + | func NewConn(rw io.ReadWriter) *Conn { | |
| 47 | + | return &Conn{rw: rw, br: bufio.NewReaderSize(rw, 64*1024)} | |
| 48 | + | } | |
| 49 | + | ||
| 50 | + | // Send writes one frame. | |
| 51 | + | func (c *Conn) Send(f *execpb.Frame) error { | |
| 52 | + | b, err := proto.Marshal(f) | |
| 53 | + | if err != nil { | |
| 54 | + | return fmt.Errorf("execstream: marshal: %w", err) | |
| 55 | + | } | |
| 56 | + | if len(b) > MaxFrame { | |
| 57 | + | return ErrFrameTooLarge | |
| 58 | + | } | |
| 59 | + | buf := make([]byte, 4+len(b)) | |
| 60 | + | binary.BigEndian.PutUint32(buf, uint32(len(b))) | |
| 61 | + | copy(buf[4:], b) | |
| 62 | + | c.wmu.Lock() | |
| 63 | + | defer c.wmu.Unlock() | |
| 64 | + | _, err = c.rw.Write(buf) | |
| 65 | + | return err | |
| 66 | + | } | |
| 67 | + | ||
| 68 | + | // Recv reads one frame. It returns io.EOF only on a clean end of stream between | |
| 69 | + | // frames; a stream cut inside a frame yields io.ErrUnexpectedEOF. | |
| 70 | + | func (c *Conn) Recv() (*execpb.Frame, error) { | |
| 71 | + | var hdr [4]byte | |
| 72 | + | if _, err := io.ReadFull(c.br, hdr[:]); err != nil { | |
| 73 | + | return nil, err | |
| 74 | + | } | |
| 75 | + | n := binary.BigEndian.Uint32(hdr[:]) | |
| 76 | + | if n > MaxFrame { | |
| 77 | + | return nil, ErrFrameTooLarge | |
| 78 | + | } | |
| 79 | + | b := make([]byte, n) | |
| 80 | + | if _, err := io.ReadFull(c.br, b); err != nil { | |
| 81 | + | if errors.Is(err, io.EOF) { | |
| 82 | + | err = io.ErrUnexpectedEOF | |
| 83 | + | } | |
| 84 | + | return nil, err | |
| 85 | + | } | |
| 86 | + | f := &execpb.Frame{} | |
| 87 | + | if err := proto.Unmarshal(b, f); err != nil { | |
| 88 | + | return nil, fmt.Errorf("execstream: unmarshal: %w", err) | |
| 89 | + | } | |
| 90 | + | return f, nil | |
| 91 | + | } | |
| 92 | + | ||
| 93 | + | // Drive is the host side of a session: it waits for the executor's Hello, sends | |
| 94 | + | // run, passes every Output frame to onOutput, and returns the Exit frame. | |
| 95 | + | // When ctx ends first, a Cancel frame is sent once and Drive keeps waiting for | |
| 96 | + | // the executor's Exit; the caller bounds that wait by stopping the sandbox, | |
| 97 | + | // which closes the stream and makes Drive return an error. | |
| 98 | + | func Drive(ctx context.Context, rw io.ReadWriter, run *execpb.Run, onOutput func(*execpb.Output)) (*execpb.Hello, *execpb.Exit, error) { | |
| 99 | + | c := NewConn(rw) | |
| 100 | + | first, err := c.Recv() | |
| 101 | + | if err != nil { | |
| 102 | + | return nil, nil, fmt.Errorf("execstream: waiting for hello: %w", err) | |
| 103 | + | } | |
| 104 | + | hello := first.GetHello() | |
| 105 | + | if hello == nil { | |
| 106 | + | return nil, nil, fmt.Errorf("execstream: first frame is not hello") | |
| 107 | + | } | |
| 108 | + | if err := c.Send(&execpb.Frame{Body: &execpb.Frame_Run{Run: run}}); err != nil { | |
| 109 | + | return hello, nil, fmt.Errorf("execstream: sending run: %w", err) | |
| 110 | + | } | |
| 111 | + | ||
| 112 | + | done := make(chan struct{}) | |
| 113 | + | defer close(done) | |
| 114 | + | go func() { | |
| 115 | + | select { | |
| 116 | + | case <-ctx.Done(): | |
| 117 | + | _ = c.Send(&execpb.Frame{Body: &execpb.Frame_Cancel{Cancel: &execpb.Cancel{}}}) | |
| 118 | + | case <-done: | |
| 119 | + | } | |
| 120 | + | }() | |
| 121 | + | ||
| 122 | + | for { | |
| 123 | + | f, err := c.Recv() | |
| 124 | + | if err != nil { | |
| 125 | + | return hello, nil, fmt.Errorf("execstream: reading: %w", err) | |
| 126 | + | } | |
| 127 | + | switch b := f.Body.(type) { | |
| 128 | + | case *execpb.Frame_Output: | |
| 129 | + | if onOutput != nil { | |
| 130 | + | onOutput(b.Output) | |
| 131 | + | } | |
| 132 | + | case *execpb.Frame_Exit: | |
| 133 | + | return hello, b.Exit, nil | |
| 134 | + | } | |
| 135 | + | } | |
| 136 | + | } | |
| @@ -0,0 +1,136 @@ | |||||
| 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. | |||
| 10 | + | package execstream | |||
| 11 | + | ||||
| 12 | + | //go:generate protoc --proto_path=. --go_out=execpb --go_opt=paths=source_relative exec.proto | |||
| 13 | + | ||||
| 14 | + | import ( | |||
| 15 | + | "bufio" | |||
| 16 | + | "context" | |||
| 17 | + | "encoding/binary" | |||
| 18 | + | "errors" | |||
| 19 | + | "fmt" | |||
| 20 | + | "io" | |||
| 21 | + | "sync" | |||
| 22 | + | ||||
| 23 | + | "google.golang.org/protobuf/proto" | |||
| 24 | + | ||||
| 25 | + | "bnema.dev/ermite-runner/internal/execstream/execpb" | |||
| 26 | + | ) | |||
| 27 | + | ||||
| 28 | + | // MaxFrame is the largest accepted serialized frame. | |||
| 29 | + | const MaxFrame = 1 << 20 | |||
| 30 | + | ||||
| 31 | + | // ChunkSize is the largest output payload placed in one Output frame. | |||
| 32 | + | const ChunkSize = 32 * 1024 | |||
| 33 | + | ||||
| 34 | + | // ErrFrameTooLarge is returned for a frame above MaxFrame, on either side. | |||
| 35 | + | var ErrFrameTooLarge = errors.New("execstream: frame too large") | |||
| 36 | + | ||||
| 37 | + | // Conn reads and writes frames on a byte stream. Send is safe for concurrent | |||
| 38 | + | // use; Recv must be called from one goroutine. | |||
| 39 | + | type Conn struct { | |||
| 40 | + | rw io.ReadWriter | |||
| 41 | + | br *bufio.Reader | |||
| 42 | + | wmu sync.Mutex | |||
| 43 | + | } | |||
| 44 | + | ||||
| 45 | + | // NewConn wraps a byte stream. | |||
| 46 | + | func NewConn(rw io.ReadWriter) *Conn { | |||
| 47 | + | return &Conn{rw: rw, br: bufio.NewReaderSize(rw, 64*1024)} | |||
| 48 | + | } | |||
| 49 | + | ||||
| 50 | + | // Send writes one frame. | |||
| 51 | + | func (c *Conn) Send(f *execpb.Frame) error { | |||
| 52 | + | b, err := proto.Marshal(f) | |||
| 53 | + | if err != nil { | |||
| 54 | + | return fmt.Errorf("execstream: marshal: %w", err) | |||
| 55 | + | } | |||
| 56 | + | if len(b) > MaxFrame { | |||
| 57 | + | return ErrFrameTooLarge | |||
| 58 | + | } | |||
| 59 | + | buf := make([]byte, 4+len(b)) | |||
| 60 | + | binary.BigEndian.PutUint32(buf, uint32(len(b))) | |||
| 61 | + | copy(buf[4:], b) | |||
| 62 | + | c.wmu.Lock() | |||
| 63 | + | defer c.wmu.Unlock() | |||
| 64 | + | _, err = c.rw.Write(buf) | |||
| 65 | + | return err | |||
| 66 | + | } | |||
| 67 | + | ||||
| 68 | + | // Recv reads one frame. It returns io.EOF only on a clean end of stream between | |||
| 69 | + | // frames; a stream cut inside a frame yields io.ErrUnexpectedEOF. | |||
| 70 | + | func (c *Conn) Recv() (*execpb.Frame, error) { | |||
| 71 | + | var hdr [4]byte | |||
| 72 | + | if _, err := io.ReadFull(c.br, hdr[:]); err != nil { | |||
| 73 | + | return nil, err | |||
| 74 | + | } | |||
| 75 | + | n := binary.BigEndian.Uint32(hdr[:]) | |||
| 76 | + | if n > MaxFrame { | |||
| 77 | + | return nil, ErrFrameTooLarge | |||
| 78 | + | } | |||
| 79 | + | b := make([]byte, n) | |||
| 80 | + | if _, err := io.ReadFull(c.br, b); err != nil { | |||
| 81 | + | if errors.Is(err, io.EOF) { | |||
| 82 | + | err = io.ErrUnexpectedEOF | |||
| 83 | + | } | |||
| 84 | + | return nil, err | |||
| 85 | + | } | |||
| 86 | + | f := &execpb.Frame{} | |||
| 87 | + | if err := proto.Unmarshal(b, f); err != nil { | |||
| 88 | + | return nil, fmt.Errorf("execstream: unmarshal: %w", err) | |||
| 89 | + | } | |||
| 90 | + | return f, nil | |||
| 91 | + | } | |||
| 92 | + | ||||
| 93 | + | // Drive is the host side of a session: it waits for the executor's Hello, sends | |||
| 94 | + | // run, passes every Output frame to onOutput, and returns the Exit frame. | |||
| 95 | + | // When ctx ends first, a Cancel frame is sent once and Drive keeps waiting for | |||
| 96 | + | // the executor's Exit; the caller bounds that wait by stopping the sandbox, | |||
| 97 | + | // which closes the stream and makes Drive return an error. | |||
| 98 | + | func Drive(ctx context.Context, rw io.ReadWriter, run *execpb.Run, onOutput func(*execpb.Output)) (*execpb.Hello, *execpb.Exit, error) { | |||
| 99 | + | c := NewConn(rw) | |||
| 100 | + | first, err := c.Recv() | |||
| 101 | + | if err != nil { | |||
| 102 | + | return nil, nil, fmt.Errorf("execstream: waiting for hello: %w", err) | |||
| 103 | + | } | |||
| 104 | + | hello := first.GetHello() | |||
| 105 | + | if hello == nil { | |||
| 106 | + | return nil, nil, fmt.Errorf("execstream: first frame is not hello") | |||
| 107 | + | } | |||
| 108 | + | if err := c.Send(&execpb.Frame{Body: &execpb.Frame_Run{Run: run}}); err != nil { | |||
| 109 | + | return hello, nil, fmt.Errorf("execstream: sending run: %w", err) | |||
| 110 | + | } | |||
| 111 | + | ||||
| 112 | + | done := make(chan struct{}) | |||
| 113 | + | defer close(done) | |||
| 114 | + | go func() { | |||
| 115 | + | select { | |||
| 116 | + | case <-ctx.Done(): | |||
| 117 | + | _ = c.Send(&execpb.Frame{Body: &execpb.Frame_Cancel{Cancel: &execpb.Cancel{}}}) | |||
| 118 | + | case <-done: | |||
| 119 | + | } | |||
| 120 | + | }() | |||
| 121 | + | ||||
| 122 | + | for { | |||
| 123 | + | f, err := c.Recv() | |||
| 124 | + | if err != nil { | |||
| 125 | + | return hello, nil, fmt.Errorf("execstream: reading: %w", err) | |||
| 126 | + | } | |||
| 127 | + | switch b := f.Body.(type) { | |||
| 128 | + | case *execpb.Frame_Output: | |||
| 129 | + | if onOutput != nil { | |||
| 130 | + | onOutput(b.Output) | |||
| 131 | + | } | |||
| 132 | + | case *execpb.Frame_Exit: | |||
| 133 | + | return hello, b.Exit, nil | |||
| 134 | + | } | |||
| 135 | + | } | |||
| 136 | + | } | |||
A internal/execstream/frame_test.go +139 -0
| @@ -0,0 +1,139 @@ | |||
| 1 | + | // SPDX-License-Identifier: MIT | |
| 2 | + | ||
| 3 | + | package execstream | |
| 4 | + | ||
| 5 | + | import ( | |
| 6 | + | "bytes" | |
| 7 | + | "context" | |
| 8 | + | "encoding/binary" | |
| 9 | + | "errors" | |
| 10 | + | "io" | |
| 11 | + | "net" | |
| 12 | + | "testing" | |
| 13 | + | "testing/iotest" | |
| 14 | + | "time" | |
| 15 | + | ||
| 16 | + | "bnema.dev/ermite-runner/internal/execstream/execpb" | |
| 17 | + | ) | |
| 18 | + | ||
| 19 | + | func TestRoundTrip(t *testing.T) { | |
| 20 | + | var buf bytes.Buffer | |
| 21 | + | c := NewConn(&buf) | |
| 22 | + | in := []*execpb.Frame{ | |
| 23 | + | {Body: &execpb.Frame_Hello{Hello: &execpb.Hello{Version: "v", Sandbox: "vm"}}}, | |
| 24 | + | {Body: &execpb.Frame_Output{Output: &execpb.Output{Stream: execpb.StreamKind_STREAM_KIND_STDERR, Data: []byte("x\x00y")}}}, | |
| 25 | + | {Body: &execpb.Frame_Exit{Exit: &execpb.Exit{Code: 3}}}, | |
| 26 | + | } | |
| 27 | + | for _, f := range in { | |
| 28 | + | if err := c.Send(f); err != nil { | |
| 29 | + | t.Fatal(err) | |
| 30 | + | } | |
| 31 | + | } | |
| 32 | + | // One byte at a time: framing must not depend on read sizes. | |
| 33 | + | r := NewConn(struct { | |
| 34 | + | io.Reader | |
| 35 | + | io.Writer | |
| 36 | + | }{iotest.OneByteReader(&buf), io.Discard}) | |
| 37 | + | for i := range in { | |
| 38 | + | out, err := r.Recv() | |
| 39 | + | if err != nil { | |
| 40 | + | t.Fatalf("frame %d: %v", i, err) | |
| 41 | + | } | |
| 42 | + | if out.String() != in[i].String() { | |
| 43 | + | t.Fatalf("frame %d: got %v want %v", i, out, in[i]) | |
| 44 | + | } | |
| 45 | + | } | |
| 46 | + | if _, err := r.Recv(); err != io.EOF { | |
| 47 | + | t.Fatalf("want io.EOF at end, got %v", err) | |
| 48 | + | } | |
| 49 | + | } | |
| 50 | + | ||
| 51 | + | func TestRecvErrors(t *testing.T) { | |
| 52 | + | huge := make([]byte, 4) | |
| 53 | + | binary.BigEndian.PutUint32(huge, MaxFrame+1) | |
| 54 | + | if _, err := NewConn(bytes.NewBuffer(huge)).Recv(); !errors.Is(err, ErrFrameTooLarge) { | |
| 55 | + | t.Fatalf("oversize: %v", err) | |
| 56 | + | } | |
| 57 | + | trunc := []byte{0, 0, 0, 9, 1, 2} | |
| 58 | + | if _, err := NewConn(bytes.NewBuffer(trunc)).Recv(); !errors.Is(err, io.ErrUnexpectedEOF) { | |
| 59 | + | t.Fatalf("truncated body: %v", err) | |
| 60 | + | } | |
| 61 | + | if _, err := NewConn(bytes.NewBuffer([]byte{0, 0})).Recv(); !errors.Is(err, io.ErrUnexpectedEOF) { | |
| 62 | + | t.Fatalf("truncated header: %v", err) | |
| 63 | + | } | |
| 64 | + | garbage := []byte{0, 0, 0, 3, 0xff, 0xff, 0xff} | |
| 65 | + | if _, err := NewConn(bytes.NewBuffer(garbage)).Recv(); err == nil { | |
| 66 | + | t.Fatal("garbage payload accepted") | |
| 67 | + | } | |
| 68 | + | } | |
| 69 | + | ||
| 70 | + | func TestSendTooLarge(t *testing.T) { | |
| 71 | + | c := NewConn(struct { | |
| 72 | + | io.Reader | |
| 73 | + | io.Writer | |
| 74 | + | }{bytes.NewReader(nil), io.Discard}) | |
| 75 | + | err := c.Send(&execpb.Frame{Body: &execpb.Frame_Output{Output: &execpb.Output{Data: make([]byte, MaxFrame)}}}) | |
| 76 | + | if !errors.Is(err, ErrFrameTooLarge) { | |
| 77 | + | t.Fatalf("got %v", err) | |
| 78 | + | } | |
| 79 | + | } | |
| 80 | + | ||
| 81 | + | func TestDrive(t *testing.T) { | |
| 82 | + | host, guest := net.Pipe() | |
| 83 | + | defer func() { _ = host.Close() }() | |
| 84 | + | go func() { | |
| 85 | + | defer func() { _ = guest.Close() }() | |
| 86 | + | g := NewConn(guest) | |
| 87 | + | _ = g.Send(&execpb.Frame{Body: &execpb.Frame_Hello{Hello: &execpb.Hello{Sandbox: "container"}}}) | |
| 88 | + | f, err := g.Recv() | |
| 89 | + | if err != nil || f.GetRun() == nil || f.GetRun().Script != "echo hi" { | |
| 90 | + | return | |
| 91 | + | } | |
| 92 | + | _ = g.Send(&execpb.Frame{Body: &execpb.Frame_Output{Output: &execpb.Output{Data: []byte("hi\n")}}}) | |
| 93 | + | _ = g.Send(&execpb.Frame{Body: &execpb.Frame_Exit{Exit: &execpb.Exit{Code: 7}}}) | |
| 94 | + | }() | |
| 95 | + | var got []byte | |
| 96 | + | hello, exit, err := Drive(context.Background(), host, &execpb.Run{Script: "echo hi"}, func(o *execpb.Output) { got = append(got, o.Data...) }) | |
| 97 | + | if err != nil { | |
| 98 | + | t.Fatal(err) | |
| 99 | + | } | |
| 100 | + | if hello.Sandbox != "container" || exit.Code != 7 || string(got) != "hi\n" { | |
| 101 | + | t.Fatalf("hello=%v exit=%v out=%q", hello, exit, got) | |
| 102 | + | } | |
| 103 | + | } | |
| 104 | + | ||
| 105 | + | func TestDriveCancelSendsCancel(t *testing.T) { | |
| 106 | + | host, guest := net.Pipe() | |
| 107 | + | defer func() { _ = host.Close() }() | |
| 108 | + | gotCancel := make(chan bool, 1) | |
| 109 | + | go func() { | |
| 110 | + | defer func() { _ = guest.Close() }() | |
| 111 | + | g := NewConn(guest) | |
| 112 | + | _ = g.Send(&execpb.Frame{Body: &execpb.Frame_Hello{Hello: &execpb.Hello{}}}) | |
| 113 | + | _, _ = g.Recv() // run | |
| 114 | + | f, err := g.Recv() | |
| 115 | + | gotCancel <- err == nil && f.GetCancel() != nil | |
| 116 | + | _ = g.Send(&execpb.Frame{Body: &execpb.Frame_Exit{Exit: &execpb.Exit{Code: 130}}}) | |
| 117 | + | }() | |
| 118 | + | ctx, cancel := context.WithCancel(context.Background()) | |
| 119 | + | time.AfterFunc(50*time.Millisecond, cancel) | |
| 120 | + | _, exit, err := Drive(ctx, host, &execpb.Run{}, nil) | |
| 121 | + | if err != nil || exit.Code != 130 { | |
| 122 | + | t.Fatalf("exit=%v err=%v", exit, err) | |
| 123 | + | } | |
| 124 | + | if !<-gotCancel { | |
| 125 | + | t.Fatal("guest did not receive a Cancel frame") | |
| 126 | + | } | |
| 127 | + | } | |
| 128 | + | ||
| 129 | + | func TestDriveHelloMissing(t *testing.T) { | |
| 130 | + | host, guest := net.Pipe() | |
| 131 | + | defer func() { _ = host.Close() }() | |
| 132 | + | go func() { | |
| 133 | + | defer func() { _ = guest.Close() }() | |
| 134 | + | _ = NewConn(guest).Send(&execpb.Frame{Body: &execpb.Frame_Exit{Exit: &execpb.Exit{}}}) | |
| 135 | + | }() | |
| 136 | + | if _, _, err := Drive(context.Background(), host, &execpb.Run{}, nil); err == nil { | |
| 137 | + | t.Fatal("want error when the first frame is not hello") | |
| 138 | + | } | |
| 139 | + | } | |
| @@ -0,0 +1,139 @@ | |||||
| 1 | + | // SPDX-License-Identifier: MIT | |||
| 2 | + | ||||
| 3 | + | package execstream | |||
| 4 | + | ||||
| 5 | + | import ( | |||
| 6 | + | "bytes" | |||
| 7 | + | "context" | |||
| 8 | + | "encoding/binary" | |||
| 9 | + | "errors" | |||
| 10 | + | "io" | |||
| 11 | + | "net" | |||
| 12 | + | "testing" | |||
| 13 | + | "testing/iotest" | |||
| 14 | + | "time" | |||
| 15 | + | ||||
| 16 | + | "bnema.dev/ermite-runner/internal/execstream/execpb" | |||
| 17 | + | ) | |||
| 18 | + | ||||
| 19 | + | func TestRoundTrip(t *testing.T) { | |||
| 20 | + | var buf bytes.Buffer | |||
| 21 | + | c := NewConn(&buf) | |||
| 22 | + | in := []*execpb.Frame{ | |||
| 23 | + | {Body: &execpb.Frame_Hello{Hello: &execpb.Hello{Version: "v", Sandbox: "vm"}}}, | |||
| 24 | + | {Body: &execpb.Frame_Output{Output: &execpb.Output{Stream: execpb.StreamKind_STREAM_KIND_STDERR, Data: []byte("x\x00y")}}}, | |||
| 25 | + | {Body: &execpb.Frame_Exit{Exit: &execpb.Exit{Code: 3}}}, | |||
| 26 | + | } | |||
| 27 | + | for _, f := range in { | |||
| 28 | + | if err := c.Send(f); err != nil { | |||
| 29 | + | t.Fatal(err) | |||
| 30 | + | } | |||
| 31 | + | } | |||
| 32 | + | // One byte at a time: framing must not depend on read sizes. | |||
| 33 | + | r := NewConn(struct { | |||
| 34 | + | io.Reader | |||
| 35 | + | io.Writer | |||
| 36 | + | }{iotest.OneByteReader(&buf), io.Discard}) | |||
| 37 | + | for i := range in { | |||
| 38 | + | out, err := r.Recv() | |||
| 39 | + | if err != nil { | |||
| 40 | + | t.Fatalf("frame %d: %v", i, err) | |||
| 41 | + | } | |||
| 42 | + | if out.String() != in[i].String() { | |||
| 43 | + | t.Fatalf("frame %d: got %v want %v", i, out, in[i]) | |||
| 44 | + | } | |||
| 45 | + | } | |||
| 46 | + | if _, err := r.Recv(); err != io.EOF { | |||
| 47 | + | t.Fatalf("want io.EOF at end, got %v", err) | |||
| 48 | + | } | |||
| 49 | + | } | |||
| 50 | + | ||||
| 51 | + | func TestRecvErrors(t *testing.T) { | |||
| 52 | + | huge := make([]byte, 4) | |||
| 53 | + | binary.BigEndian.PutUint32(huge, MaxFrame+1) | |||
| 54 | + | if _, err := NewConn(bytes.NewBuffer(huge)).Recv(); !errors.Is(err, ErrFrameTooLarge) { | |||
| 55 | + | t.Fatalf("oversize: %v", err) | |||
| 56 | + | } | |||
| 57 | + | trunc := []byte{0, 0, 0, 9, 1, 2} | |||
| 58 | + | if _, err := NewConn(bytes.NewBuffer(trunc)).Recv(); !errors.Is(err, io.ErrUnexpectedEOF) { | |||
| 59 | + | t.Fatalf("truncated body: %v", err) | |||
| 60 | + | } | |||
| 61 | + | if _, err := NewConn(bytes.NewBuffer([]byte{0, 0})).Recv(); !errors.Is(err, io.ErrUnexpectedEOF) { | |||
| 62 | + | t.Fatalf("truncated header: %v", err) | |||
| 63 | + | } | |||
| 64 | + | garbage := []byte{0, 0, 0, 3, 0xff, 0xff, 0xff} | |||
| 65 | + | if _, err := NewConn(bytes.NewBuffer(garbage)).Recv(); err == nil { | |||
| 66 | + | t.Fatal("garbage payload accepted") | |||
| 67 | + | } | |||
| 68 | + | } | |||
| 69 | + | ||||
| 70 | + | func TestSendTooLarge(t *testing.T) { | |||
| 71 | + | c := NewConn(struct { | |||
| 72 | + | io.Reader | |||
| 73 | + | io.Writer | |||
| 74 | + | }{bytes.NewReader(nil), io.Discard}) | |||
| 75 | + | err := c.Send(&execpb.Frame{Body: &execpb.Frame_Output{Output: &execpb.Output{Data: make([]byte, MaxFrame)}}}) | |||
| 76 | + | if !errors.Is(err, ErrFrameTooLarge) { | |||
| 77 | + | t.Fatalf("got %v", err) | |||
| 78 | + | } | |||
| 79 | + | } | |||
| 80 | + | ||||
| 81 | + | func TestDrive(t *testing.T) { | |||
| 82 | + | host, guest := net.Pipe() | |||
| 83 | + | defer func() { _ = host.Close() }() | |||
| 84 | + | go func() { | |||
| 85 | + | defer func() { _ = guest.Close() }() | |||
| 86 | + | g := NewConn(guest) | |||
| 87 | + | _ = g.Send(&execpb.Frame{Body: &execpb.Frame_Hello{Hello: &execpb.Hello{Sandbox: "container"}}}) | |||
| 88 | + | f, err := g.Recv() | |||
| 89 | + | if err != nil || f.GetRun() == nil || f.GetRun().Script != "echo hi" { | |||
| 90 | + | return | |||
| 91 | + | } | |||
| 92 | + | _ = g.Send(&execpb.Frame{Body: &execpb.Frame_Output{Output: &execpb.Output{Data: []byte("hi\n")}}}) | |||
| 93 | + | _ = g.Send(&execpb.Frame{Body: &execpb.Frame_Exit{Exit: &execpb.Exit{Code: 7}}}) | |||
| 94 | + | }() | |||
| 95 | + | var got []byte | |||
| 96 | + | hello, exit, err := Drive(context.Background(), host, &execpb.Run{Script: "echo hi"}, func(o *execpb.Output) { got = append(got, o.Data...) }) | |||
| 97 | + | if err != nil { | |||
| 98 | + | t.Fatal(err) | |||
| 99 | + | } | |||
| 100 | + | if hello.Sandbox != "container" || exit.Code != 7 || string(got) != "hi\n" { | |||
| 101 | + | t.Fatalf("hello=%v exit=%v out=%q", hello, exit, got) | |||
| 102 | + | } | |||
| 103 | + | } | |||
| 104 | + | ||||
| 105 | + | func TestDriveCancelSendsCancel(t *testing.T) { | |||
| 106 | + | host, guest := net.Pipe() | |||
| 107 | + | defer func() { _ = host.Close() }() | |||
| 108 | + | gotCancel := make(chan bool, 1) | |||
| 109 | + | go func() { | |||
| 110 | + | defer func() { _ = guest.Close() }() | |||
| 111 | + | g := NewConn(guest) | |||
| 112 | + | _ = g.Send(&execpb.Frame{Body: &execpb.Frame_Hello{Hello: &execpb.Hello{}}}) | |||
| 113 | + | _, _ = g.Recv() // run | |||
| 114 | + | f, err := g.Recv() | |||
| 115 | + | gotCancel <- err == nil && f.GetCancel() != nil | |||
| 116 | + | _ = g.Send(&execpb.Frame{Body: &execpb.Frame_Exit{Exit: &execpb.Exit{Code: 130}}}) | |||
| 117 | + | }() | |||
| 118 | + | ctx, cancel := context.WithCancel(context.Background()) | |||
| 119 | + | time.AfterFunc(50*time.Millisecond, cancel) | |||
| 120 | + | _, exit, err := Drive(ctx, host, &execpb.Run{}, nil) | |||
| 121 | + | if err != nil || exit.Code != 130 { | |||
| 122 | + | t.Fatalf("exit=%v err=%v", exit, err) | |||
| 123 | + | } | |||
| 124 | + | if !<-gotCancel { | |||
| 125 | + | t.Fatal("guest did not receive a Cancel frame") | |||
| 126 | + | } | |||
| 127 | + | } | |||
| 128 | + | ||||
| 129 | + | func TestDriveHelloMissing(t *testing.T) { | |||
| 130 | + | host, guest := net.Pipe() | |||
| 131 | + | defer func() { _ = host.Close() }() | |||
| 132 | + | go func() { | |||
| 133 | + | defer func() { _ = guest.Close() }() | |||
| 134 | + | _ = NewConn(guest).Send(&execpb.Frame{Body: &execpb.Frame_Exit{Exit: &execpb.Exit{}}}) | |||
| 135 | + | }() | |||
| 136 | + | if _, _, err := Drive(context.Background(), host, &execpb.Run{}, nil); err == nil { | |||
| 137 | + | t.Fatal("want error when the first frame is not hello") | |||
| 138 | + | } | |||
| 139 | + | } | |||