1package sessionwire
2
3import (
4 "bytes"
5 "context"
6 "encoding/binary"
7 "testing"
8 "time"
9
10 "google.golang.org/protobuf/encoding/protowire"
11 "google.golang.org/protobuf/proto"
12
13 "github.com/bnema/vev/internal/protocol/wire"
14 "github.com/stretchr/testify/require"
15)
16
17func preambleFrame(payload []byte) []byte {
18 var header [4]byte
19 binary.BigEndian.PutUint32(header[:], uint32(len(payload)))
20 return append(header[:], payload...)
21}
22
23func mustMarshalPreambleRequest(t *testing.T) []byte {
24 t.Helper()
25 raw, err := proto.Marshal(preambleRequestToWire(defaultProtoCeilings()))
26 require.NoError(t, err)
27 return raw
28}
29
30func mustMarshalPreambleResponse(t *testing.T) []byte {
31 t.Helper()
32 raw, err := proto.Marshal(preambleResponseToWire(true, defaultProtoCeilings(), 0))
33 require.NoError(t, err)
34 return raw
35}
36
37func withUnknownField(t *testing.T, payload []byte) []byte {
38 t.Helper()
39 unknown := protowire.AppendTag(append([]byte(nil), payload...), 99, protowire.VarintType)
40 return protowire.AppendVarint(unknown, 1)
41}
42
43// TestReadPreambleFrameRejectsUnscannedPayloads proves the strict envelope
44// scan runs before generated unmarshal on the preamble read path: unknown,
45// repeated, and concatenated payloads are refused.
46func TestReadPreambleFrameRejectsUnscannedPayloads(t *testing.T) {
47 request := mustMarshalPreambleRequest(t)
48 response := mustMarshalPreambleResponse(t)
49 magic := protowire.Number((&wire.PreambleRequest{}).ProtoReflect().Descriptor().Fields().ByName("magic").Number())
50
51 duplicateMagic := protowire.AppendTag(append([]byte(nil), request...), magic, protowire.VarintType)
52 duplicateMagic = protowire.AppendVarint(duplicateMagic, 7)
53
54 tests := []struct {
55 name string
56 message proto.Message
57 payload []byte
58 wantErr error
59 }{
60 {name: "valid request", message: &wire.PreambleRequest{}, payload: request},
61 {name: "valid response", message: &wire.PreambleResponse{}, payload: response},
62 {name: "unknown request field", message: &wire.PreambleRequest{}, payload: withUnknownField(t, request), wantErr: wire.ErrScanUnknown},
63 {name: "unknown response field", message: &wire.PreambleResponse{}, payload: withUnknownField(t, response), wantErr: wire.ErrScanUnknown},
64 {name: "repeated singular request field", message: &wire.PreambleRequest{}, payload: duplicateMagic, wantErr: wire.ErrScanDuplicate},
65 {name: "concatenated requests", message: &wire.PreambleRequest{}, payload: append(append([]byte(nil), request...), request...), wantErr: wire.ErrScanDuplicate},
66 {name: "concatenated responses", message: &wire.PreambleResponse{}, payload: append(append([]byte(nil), response...), response...), wantErr: wire.ErrScanDuplicate},
67 }
68 for _, tt := range tests {
69 t.Run(tt.name, func(t *testing.T) {
70 err := readPreambleFrame(bytes.NewReader(preambleFrame(tt.payload)), tt.message)
71 if tt.wantErr != nil {
72 require.ErrorIs(t, err, tt.wantErr)
73 return
74 }
75 require.NoError(t, err)
76 })
77 }
78}
79
80// TestProtoPreambleRejectsScannedMalformation proves the live preamble
81// exchange refuses unknown and concatenated payloads on both sides before
82// unmarshal, with the server sending a typed refusal.
83func TestProtoPreambleRejectsScannedMalformation(t *testing.T) {
84 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
85 defer cancel()
86
87 t.Run("server rejects unknown request field", func(t *testing.T) {
88 raw := &scriptedTransport{recv: []wire.Envelope{{Payload: withUnknownField(t, mustMarshalPreambleRequest(t))}}}
89 _, err := runProtoServerPreamble(ctx, raw, defaultProtoCeilings())
90 require.ErrorIs(t, err, wire.ErrScanUnknown)
91 require.Equal(t, 1, raw.sentLen())
92 })
93 t.Run("server rejects concatenated request", func(t *testing.T) {
94 request := mustMarshalPreambleRequest(t)
95 raw := &scriptedTransport{recv: []wire.Envelope{{Payload: append(append([]byte(nil), request...), request...)}}}
96 _, err := runProtoServerPreamble(ctx, raw, defaultProtoCeilings())
97 require.ErrorIs(t, err, wire.ErrScanDuplicate)
98 require.Equal(t, 1, raw.sentLen())
99 })
100 t.Run("client rejects unknown response field", func(t *testing.T) {
101 raw := &scriptedTransport{recv: []wire.Envelope{{Payload: withUnknownField(t, mustMarshalPreambleResponse(t))}}}
102 _, err := runProtoClientPreamble(ctx, raw, defaultProtoCeilings())
103 require.ErrorIs(t, err, wire.ErrScanUnknown)
104 })
105}