This repository has no description
1// package millproto carries the mill<->executor session protocol. a new message
2// vocabulary over the same length-prefixed protobuf framing the spindle already
3// uses to talk to the microVM guest (see spindle/agentproto). only the framing
4// pattern is shared. the messages are entirely separate
5package millproto
6
7import (
8 "encoding/binary"
9 "fmt"
10 "io"
11 "sync"
12
13 "buf.build/go/protovalidate"
14 "google.golang.org/protobuf/proto"
15
16 millv1 "tangled.org/core/spindle/mill/proto/gen"
17)
18
19const (
20 ProtocolVersion = 1
21 // generous vs agentproto's 1 MiB. a ReserveSeat carries the raw pipeline and
22 // workflow JSON, and streamed log lines can be chunky
23 MaxMessageBytes = 8 * 1024 * 1024
24)
25
26type Message = millv1.Message
27
28var validator protovalidate.Validator
29
30func init() {
31 var err error
32 validator, err = protovalidate.New()
33 if err != nil {
34 panic(fmt.Errorf("failed to initialize protovalidate validator: %w", err))
35 }
36}
37
38type Encoder struct {
39 mu sync.Mutex
40 w io.Writer
41}
42
43func NewEncoder(w io.Writer) *Encoder {
44 return &Encoder{w: w}
45}
46
47func (e *Encoder) Encode(msg *Message) error {
48 if err := validator.Validate(msg); err != nil {
49 return fmt.Errorf("validate fleet message: %w", err)
50 }
51
52 data, err := proto.Marshal(msg)
53 if err != nil {
54 return fmt.Errorf("marshal fleet message: %w", err)
55 }
56 if len(data) > MaxMessageBytes {
57 return fmt.Errorf("fleet message exceeded %d bytes", MaxMessageBytes)
58 }
59
60 // single write of header and payload maps to exactly one websocket binary
61 // frame when the writer is a ws stream
62 frame := make([]byte, 4+len(data))
63 binary.BigEndian.PutUint32(frame[:4], uint32(len(data)))
64 copy(frame[4:], data)
65
66 e.mu.Lock()
67 defer e.mu.Unlock()
68 _, err = e.w.Write(frame)
69 return err
70}
71
72type Decoder struct {
73 r io.Reader
74}
75
76func NewDecoder(r io.Reader) *Decoder {
77 return &Decoder{r: r}
78}
79
80func (d *Decoder) Decode() (*Message, error) {
81 msg := &Message{}
82 var header [4]byte
83 if _, err := io.ReadFull(d.r, header[:]); err != nil {
84 return msg, err
85 }
86
87 size := binary.BigEndian.Uint32(header[:])
88 if size > MaxMessageBytes {
89 return msg, fmt.Errorf("fleet message exceeded %d bytes", MaxMessageBytes)
90 }
91
92 data := make([]byte, size)
93 if _, err := io.ReadFull(d.r, data); err != nil {
94 return msg, err
95 }
96 if err := proto.Unmarshal(data, msg); err != nil {
97 return msg, fmt.Errorf("parse fleet message: %w", err)
98 }
99 if err := validator.Validate(msg); err != nil {
100 return msg, fmt.Errorf("validate fleet message: %w", err)
101 }
102 return msg, nil
103}