1package wire
2
3import (
4 "bytes"
5 "errors"
6 "testing"
7
8 "google.golang.org/protobuf/encoding/protowire"
9 "google.golang.org/protobuf/proto"
10 "google.golang.org/protobuf/reflect/protoreflect"
11)
12
13// The daemonmux conversation is separate from the session and broker
14// conversations: its tags are disjoint from both, both directions are
15// closed oneof unions, the strict scanner treats it exactly like the other
16// envelopes, and the shared physical preamble pair is not an envelope.
17
18// muxClientTags is the frozen client tag inventory.
19var muxClientTags = map[string]protoreflect.FieldNumber{
20 "open": 301, "data": 302, "close": 303, "reset": 304, "window_update": 305,
21}
22
23// muxServerTags is the frozen server tag inventory.
24var muxServerTags = map[string]protoreflect.FieldNumber{
25 "opened": 401, "refused": 402, "data": 403, "close": 404, "reset": 405, "window_update": 406,
26}
27
28func muxStreamRef() *MuxStreamRef {
29 return &MuxStreamRef{
30 PhysicalStreamId: 5,
31 BrokerEpoch: 7,
32 ConnectionId: bytes.Repeat([]byte{0x11}, 16),
33 ClientStreamId: 3,
34 }
35}
36
37func muxRegistration() *RemoteRegistration {
38 return &RemoteRegistration{Endpoint: "dev@host:22", Incarnation: bytes.Repeat([]byte{0x22}, 16), Generation: 9}
39}
40
41func muxPolicy() *BrokerWirePolicy {
42 return &BrokerWirePolicy{
43 ProtocolVersion: 53, CatalogueVersion: 3, EnvironmentPolicy: 1,
44 Transport: "quic", Trust: "pinned", Launch: "agent", Isolation: "user",
45 }
46}
47
48func muxErrorDetail() *BrokerErrorDetail {
49 return &BrokerErrorDetail{Code: 2, Text: "incompatible", AdmissionCode: 1, FailureKind: 3}
50}
51
52// muxClientSamples returns one populated envelope per client variant.
53func muxClientSamples() map[string]*MuxClientEnvelope {
54 return map[string]*MuxClientEnvelope{
55 "open": {Payload: &MuxClientEnvelope_Open{Open: &MuxOpen{
56 Ref: muxStreamRef(), Purpose: 1, Local: false, Endpoint: "dev@host:22",
57 Registration: muxRegistration(),
58 Target: &ExactTarget{LifecycleId: &LifecycleID{Value: bytes.Repeat([]byte{0x55}, 16)}, SessionName: "work"},
59 Env: []string{"TERM=xterm-256color", "LANG=C.UTF-8"},
60 Policy: muxPolicy(),
61 }}},
62 "data": {Payload: &MuxClientEnvelope_Data{Data: &MuxData{PhysicalStreamId: 5, Data: bytes.Repeat([]byte{0x66}, 24)}}},
63 "close": {Payload: &MuxClientEnvelope_Close{Close: &MuxClose{PhysicalStreamId: 5}}},
64 "reset": {Payload: &MuxClientEnvelope_Reset_{Reset_: &MuxReset{PhysicalStreamId: 5, Error: muxErrorDetail()}}},
65 "window_update": {Payload: &MuxClientEnvelope_WindowUpdate{WindowUpdate: &MuxWindowUpdate{PhysicalStreamId: 5, CreditBytes: 4096}}},
66 }
67}
68
69// muxServerSamples returns one populated envelope per server variant.
70func muxServerSamples() map[string]*MuxServerEnvelope {
71 return map[string]*MuxServerEnvelope{
72 "opened": {Payload: &MuxServerEnvelope_Opened{Opened: &MuxOpened{Ref: muxStreamRef()}}},
73 "refused": {Payload: &MuxServerEnvelope_Refused{Refused: &MuxRefused{Ref: muxStreamRef(), Error: muxErrorDetail()}}},
74 "data": {Payload: &MuxServerEnvelope_Data{Data: &MuxData{PhysicalStreamId: 5, Data: bytes.Repeat([]byte{0x77}, 24)}}},
75 "close": {Payload: &MuxServerEnvelope_Close{Close: &MuxClose{PhysicalStreamId: 5}}},
76 "reset": {Payload: &MuxServerEnvelope_Reset_{Reset_: &MuxReset{PhysicalStreamId: 5}}},
77 "window_update": {Payload: &MuxServerEnvelope_WindowUpdate{WindowUpdate: &MuxWindowUpdate{PhysicalStreamId: 5, CreditBytes: 4096}}},
78 }
79}
80
81func mustMarshalMux(t *testing.T, message proto.Message) []byte {
82 t.Helper()
83 raw, err := proto.MarshalOptions{Deterministic: true}.Marshal(message)
84 if err != nil {
85 t.Fatal(err)
86 }
87 return raw
88}
89
90// TestMuxEnvelopeInventory proves both daemonmux envelopes carry exactly the
91// frozen tag inventory and share no tag with the session envelopes, the
92// broker envelopes, or each other.
93func TestMuxEnvelopeInventory(t *testing.T) {
94 client := oneofTags(t, &MuxClientEnvelope{})
95 server := oneofTags(t, &MuxServerEnvelope{})
96
97 if len(client) != len(muxClientTags) {
98 t.Fatalf("mux client envelope has %d variants, want %d (%v)", len(client), len(muxClientTags), client)
99 }
100 for name, number := range muxClientTags {
101 if got, ok := client[name]; !ok || got != number {
102 t.Fatalf("mux client variant %q = %d (present %t), want %d", name, got, ok, number)
103 }
104 }
105 if len(server) != len(muxServerTags) {
106 t.Fatalf("mux server envelope has %d variants, want %d (%v)", len(server), len(muxServerTags), server)
107 }
108 for name, number := range muxServerTags {
109 if got, ok := server[name]; !ok || got != number {
110 t.Fatalf("mux server variant %q = %d (present %t), want %d", name, got, ok, number)
111 }
112 }
113
114 // The mux tag space is disjoint from every session tag and every broker
115 // tag, in both directions.
116 otherTags := map[protoreflect.FieldNumber]string{}
117 for _, envelope := range []proto.Message{&ClientEnvelope{}, &ServerEnvelope{}} {
118 for name, number := range oneofTags(t, envelope) {
119 otherTags[number] = "session " + name
120 }
121 }
122 for _, envelope := range []proto.Message{&BrokerClientEnvelope{}, &BrokerServerEnvelope{}} {
123 for name, number := range oneofTags(t, envelope) {
124 otherTags[number] = "broker " + name
125 }
126 }
127 clientByTag := map[protoreflect.FieldNumber]string{}
128 for name, number := range muxClientTags {
129 clientByTag[number] = name
130 if other, taken := otherTags[number]; taken {
131 t.Fatalf("mux client tag %d (%s) collides with %s", number, name, other)
132 }
133 }
134 for name, number := range muxServerTags {
135 if other, taken := otherTags[number]; taken {
136 t.Fatalf("mux server tag %d (%s) collides with %s", number, name, other)
137 }
138 if peer, taken := clientByTag[number]; taken {
139 t.Fatalf("mux server tag %d (%s) collides with mux client variant %s", number, name, peer)
140 }
141 }
142
143 for _, envelope := range []proto.Message{&MuxClientEnvelope{}, &MuxServerEnvelope{}} {
144 if !isEnvelopeDescriptor(envelope.ProtoReflect().Descriptor()) {
145 t.Fatalf("%v is not recognized as a directional envelope", envelope.ProtoReflect().Descriptor().FullName())
146 }
147 }
148 // The physical preamble and the shared payloads are never envelopes.
149 for _, payload := range []proto.Message{&MuxPreambleRequest{}, &MuxPreambleResponse{}, &MuxOpen{}, &MuxData{}, &MuxClose{}, &MuxReset{}, &MuxOpened{}, &MuxRefused{}, &MuxStreamRef{}} {
150 if isEnvelopeDescriptor(payload.ProtoReflect().Descriptor()) {
151 t.Fatalf("payload %v is wrongly recognized as a directional envelope", payload.ProtoReflect().Descriptor().FullName())
152 }
153 }
154}
155
156// TestMuxEnvelopeRoundTrip proves every variant survives strict scan,
157// generated unmarshal, and equality, in both directions.
158func TestMuxEnvelopeRoundTrip(t *testing.T) {
159 for name, sample := range muxClientSamples() {
160 t.Run("client/"+name, func(t *testing.T) {
161 raw := mustMarshalMux(t, sample)
162 decoded := &MuxClientEnvelope{}
163 if err := validatedUnmarshal(decoded, raw); err != nil {
164 t.Fatalf("validatedUnmarshal() = %v", err)
165 }
166 if !proto.Equal(sample, decoded) {
167 t.Fatalf("roundtrip mismatch:\n got %v\nwant %v", decoded, sample)
168 }
169 })
170 }
171 for name, sample := range muxServerSamples() {
172 t.Run("server/"+name, func(t *testing.T) {
173 raw := mustMarshalMux(t, sample)
174 decoded := &MuxServerEnvelope{}
175 if err := validatedUnmarshal(decoded, raw); err != nil {
176 t.Fatalf("validatedUnmarshal() = %v", err)
177 }
178 if !proto.Equal(sample, decoded) {
179 t.Fatalf("roundtrip mismatch:\n got %v\nwant %v", decoded, sample)
180 }
181 })
182 }
183}
184
185// TestMuxEnvelopeRequiredVariant proves both mux envelopes demand exactly
186// one payload variant, like the session and broker envelopes.
187func TestMuxEnvelopeRequiredVariant(t *testing.T) {
188 for _, message := range []proto.Message{&MuxClientEnvelope{}, &MuxServerEnvelope{}} {
189 if err := ScanEnvelope(message, mustMarshalMux(t, message)); !errors.Is(err, ErrScanEmpty) {
190 t.Fatalf("ScanEnvelope(%T empty) = %v, want ErrScanEmpty", message, err)
191 }
192 }
193}
194
195// TestMuxEnvelopeWrongDirection proves a payload serialized in one direction
196// is rejected when scanned as the other direction, and that session and
197// broker payloads are unknown on the mux conversation.
198func TestMuxEnvelopeWrongDirection(t *testing.T) {
199 client := mustMarshalMux(t, muxClientSamples()["open"])
200 if err := ScanEnvelope(&MuxServerEnvelope{}, client); !errors.Is(err, ErrScanUnknown) {
201 t.Fatalf("client payload as server envelope = %v, want ErrScanUnknown", err)
202 }
203 server := mustMarshalMux(t, muxServerSamples()["opened"])
204 if err := ScanEnvelope(&MuxClientEnvelope{}, server); !errors.Is(err, ErrScanUnknown) {
205 t.Fatalf("server payload as client envelope = %v, want ErrScanUnknown", err)
206 }
207 session := mustMarshalMux(t, &ClientEnvelope{Payload: &ClientEnvelope_Ping{Ping: &Ping{}}})
208 if err := ScanEnvelope(&MuxClientEnvelope{}, session); !errors.Is(err, ErrScanUnknown) {
209 t.Fatalf("session payload as mux envelope = %v, want ErrScanUnknown", err)
210 }
211 broker := mustMarshalMux(t, &BrokerClientEnvelope{Payload: &BrokerClientEnvelope_Register{Register: &Register{}}})
212 if err := ScanEnvelope(&MuxClientEnvelope{}, broker); !errors.Is(err, ErrScanUnknown) {
213 t.Fatalf("broker payload as mux envelope = %v, want ErrScanUnknown", err)
214 }
215}
216
217// TestMuxEnvelopeDuplicateAlternatives proves a second oneof occurrence is
218// rejected whether it repeats the same variant or switches alternatives.
219func TestMuxEnvelopeDuplicateAlternatives(t *testing.T) {
220 closeRaw := mustMarshalMux(t, muxClientSamples()["close"])
221 dataRaw := mustMarshalMux(t, muxClientSamples()["data"])
222 if err := ScanEnvelope(&MuxClientEnvelope{}, append(append([]byte(nil), closeRaw...), dataRaw...)); !errors.Is(err, ErrScanDuplicate) {
223 t.Fatalf("alternate alternatives = %v, want ErrScanDuplicate", err)
224 }
225 if err := ScanEnvelope(&MuxClientEnvelope{}, append(append([]byte(nil), dataRaw...), dataRaw...)); !errors.Is(err, ErrScanDuplicate) {
226 t.Fatalf("duplicate variant = %v, want ErrScanDuplicate", err)
227 }
228 // The shared payload types are exclusive per direction as well.
229 serverClose := mustMarshalMux(t, muxServerSamples()["close"])
230 serverReset := mustMarshalMux(t, muxServerSamples()["reset"])
231 if err := ScanEnvelope(&MuxServerEnvelope{}, append(append([]byte(nil), serverClose...), serverReset...)); !errors.Is(err, ErrScanDuplicate) {
232 t.Fatalf("duplicate server alternatives = %v, want ErrScanDuplicate", err)
233 }
234}
235
236// TestMuxEnvelopeRejectsUnknownAndWrongTypes proves unknown and mistyped
237// fields are rejected recursively, including inside the nested stream ref
238// and the shared error detail.
239func TestMuxEnvelopeRejectsUnknownAndWrongTypes(t *testing.T) {
240 // Unknown field on the envelope itself.
241 raw := mustMarshalMux(t, muxClientSamples()["close"])
242 raw = protowire.AppendTag(raw, 399, protowire.VarintType)
243 raw = protowire.AppendVarint(raw, 1)
244 if err := ScanEnvelope(&MuxClientEnvelope{}, raw); !errors.Is(err, ErrScanUnknown) {
245 t.Fatalf("unknown envelope field = %v, want ErrScanUnknown", err)
246 }
247
248 // Unknown field two levels down: MuxClientEnvelope → MuxOpen → MuxStreamRef.
249 ref := protowire.AppendTag(nil, 1, protowire.VarintType) // MuxStreamRef.physical_stream_id
250 ref = protowire.AppendVarint(ref, 5)
251 ref = protowire.AppendTag(ref, 9, protowire.VarintType) // unknown in MuxStreamRef
252 ref = protowire.AppendVarint(ref, 1)
253 open := protowire.AppendTag(nil, 1, protowire.BytesType) // MuxOpen.ref
254 open = protowire.AppendBytes(open, ref)
255 outer := protowire.AppendTag(nil, 301, protowire.BytesType) // MuxClientEnvelope.open
256 outer = protowire.AppendBytes(outer, open)
257 if err := ScanEnvelope(&MuxClientEnvelope{}, outer); !errors.Is(err, ErrScanUnknown) {
258 t.Fatalf("nested unknown field = %v, want ErrScanUnknown", err)
259 }
260
261 // Unknown field inside the shared error detail of a reset.
262 reset := protowire.AppendTag(nil, 1, protowire.VarintType) // MuxReset.physical_stream_id
263 reset = protowire.AppendVarint(reset, 5)
264 detail := protowire.AppendTag(nil, 1, protowire.VarintType) // BrokerErrorDetail.code
265 detail = protowire.AppendVarint(detail, 2)
266 detail = protowire.AppendTag(detail, 99, protowire.VarintType)
267 detail = protowire.AppendVarint(detail, 1)
268 reset = protowire.AppendTag(reset, 2, protowire.BytesType) // MuxReset.error
269 reset = protowire.AppendBytes(reset, detail)
270 raw = protowire.AppendTag(nil, 405, protowire.BytesType) // MuxServerEnvelope.reset
271 raw = protowire.AppendBytes(raw, reset)
272 if err := ScanEnvelope(&MuxServerEnvelope{}, raw); !errors.Is(err, ErrScanUnknown) {
273 t.Fatalf("unknown nested error field = %v, want ErrScanUnknown", err)
274 }
275
276 // Wrong wire type: MuxClientEnvelope.close (bytes) as a varint.
277 wrongType := protowire.AppendTag(nil, 303, protowire.VarintType)
278 wrongType = protowire.AppendVarint(wrongType, 1)
279 if err := ScanEnvelope(&MuxClientEnvelope{}, wrongType); !errors.Is(err, ErrScanWireType) {
280 t.Fatalf("wrong wire type = %v, want ErrScanWireType", err)
281 }
282
283 // A raw varint wider than the declared uint32 is rejected before
284 // generated unmarshal can truncate its high bits.
285 open = protowire.AppendTag(nil, 2, protowire.VarintType) // MuxOpen.purpose
286 open = protowire.AppendVarint(open, 1<<33)
287 outer = protowire.AppendTag(nil, 301, protowire.BytesType)
288 outer = protowire.AppendBytes(outer, open)
289 if err := ScanEnvelope(&MuxClientEnvelope{}, outer); !errors.Is(err, ErrScanVarintRange) {
290 t.Fatalf("uint32 overflow = %v, want ErrScanVarintRange", err)
291 }
292}
293
294// TestMuxEnvelopeTruncation proves truncated tags, varints, lengths, and
295// every strict prefix of a valid envelope are rejected.
296func TestMuxEnvelopeTruncation(t *testing.T) {
297 if err := ScanEnvelope(&MuxClientEnvelope{}, []byte{0xFF}); !errors.Is(err, ErrScanTruncated) {
298 t.Fatalf("truncated tag = %v, want ErrScanTruncated", err)
299 }
300 overlong := protowire.AppendTag(nil, 301, protowire.BytesType)
301 overlong = protowire.AppendBytes(overlong, make([]byte, 8))
302 overlong = overlong[:len(overlong)-4]
303 if err := ScanEnvelope(&MuxClientEnvelope{}, overlong); !errors.Is(err, ErrScanTruncated) {
304 t.Fatalf("overlong length = %v, want ErrScanTruncated", err)
305 }
306 partialVarint := protowire.AppendTag(nil, 302, protowire.BytesType) // MuxClientEnvelope.data
307 data := protowire.AppendTag(nil, 1, protowire.VarintType) // MuxData.physical_stream_id
308 partialVarint = protowire.AppendBytes(partialVarint, append(data, 0x80))
309 if err := ScanEnvelope(&MuxClientEnvelope{}, partialVarint); !errors.Is(err, ErrScanTruncated) {
310 t.Fatalf("truncated varint = %v, want ErrScanTruncated", err)
311 }
312
313 for name, sample := range muxClientSamples() {
314 raw := mustMarshalMux(t, sample)
315 for size := range len(raw) {
316 if err := ScanEnvelope(&MuxClientEnvelope{}, raw[:size]); err == nil {
317 t.Fatalf("client prefix %s[:%d] accepted", name, size)
318 }
319 }
320 }
321 for name, sample := range muxServerSamples() {
322 raw := mustMarshalMux(t, sample)
323 for size := range len(raw) {
324 if err := ScanEnvelope(&MuxServerEnvelope{}, raw[:size]); err == nil {
325 t.Fatalf("server prefix %s[:%d] accepted", name, size)
326 }
327 }
328 }
329}
330
331// TestMuxEnvelopeTrailingGarbage proves trailing bytes and a second
332// concatenated envelope are rejected.
333func TestMuxEnvelopeTrailingGarbage(t *testing.T) {
334 raw := mustMarshalMux(t, muxServerSamples()["opened"])
335 for _, tail := range [][]byte{
336 {0x00},
337 protowire.AppendTag(nil, 405, protowire.VarintType),
338 mustMarshalMux(t, muxServerSamples()["close"]),
339 } {
340 payload := append(append([]byte(nil), raw...), tail...)
341 if err := ScanEnvelope(&MuxServerEnvelope{}, payload); err == nil {
342 t.Fatalf("trailing %x accepted", tail)
343 }
344 }
345}