1package traffic
2
3import (
4 "errors"
5 "io"
6 "net"
7 "strings"
8 "testing"
9 "time"
10
11 "github.com/stretchr/testify/assert"
12 "github.com/stretchr/testify/require"
13)
14
15func TestSniffSmartTCPClassifiesHTTP(t *testing.T) {
16 for _, tc := range []struct {
17 name string
18 data string
19 }{
20 {name: "get", data: "GET / HTTP/1.1\r\n"},
21 {name: "propfind", data: "PROPFIND / HTTP/1.1\r\n"},
22 {name: "mkcol", data: "MKCOL / HTTP/1.1\r\n"},
23 {name: "unknown_method", data: "FOO / HTTP/1.1\r\n"},
24 {name: "extension_method", data: "BAR / HTTP/1.1\r\n"},
25 } {
26 t.Run(tc.name, func(t *testing.T) {
27 result := sniffBytes(t, []byte(tc.data), time.Second)
28 assert.Equal(t, dispatchHTTP, result.kind)
29 assertReplayPrefix(t, result.conn, []byte(tc.data))
30 })
31 }
32}
33
34func TestSniffSmartTCPClassifiesH2C(t *testing.T) {
35 preface := []byte("PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n")
36 result := sniffBytes(t, preface, time.Second)
37 require.Equal(t, dispatchH2C, result.kind)
38 assertReplayPrefix(t, result.conn, preface)
39}
40
41func TestSniffSmartTCPClassifiesTLSAndAllowsClientHelloParsing(t *testing.T) {
42 hello := clientHelloBytes(t, "sniff.example.com")
43 result := sniffBytes(t, hello, time.Second)
44 require.Equal(t, dispatchTLS, result.kind)
45
46 sni, replayed, err := peekClientHelloSNI(result.conn)
47 require.NoError(t, err)
48 assert.Equal(t, "sniff.example.com", sni)
49 assertReplayPrefix(t, replayed, hello)
50}
51
52func TestSniffSmartTCPRejectsMalformedTLSLookingZeroLengthRecord(t *testing.T) {
53 result := sniffBytes(t, []byte{22, 3, 3, 0, 0}, time.Second)
54 assert.Equal(t, dispatchReject, result.kind)
55 assertReplayPrefix(t, result.conn, []byte{22, 3, 3, 0, 0})
56}
57
58func TestSniffSmartTCPClassifiesSSHBannerAsUnknown(t *testing.T) {
59 banner := []byte("SSH-2.0-client\r\n")
60 kind, done := classifySmartTCPPrefix(banner)
61 require.True(t, done)
62 assert.Equal(t, dispatchUnknown, kind)
63
64 result := sniffBytes(t, banner, time.Second)
65 assert.Equal(t, dispatchUnknown, result.kind)
66 assertReplayPrefix(t, result.conn, banner)
67}
68
69func TestSniffSmartTCPRejectsMalformedHTTPLookingRequestLine(t *testing.T) {
70 data := []byte("GET /broken\r\n")
71 kind, done := classifySmartTCPPrefix(data)
72 require.True(t, done)
73 assert.Equal(t, dispatchReject, kind)
74
75 result := sniffBytes(t, data, time.Second)
76 assert.Equal(t, dispatchReject, result.kind)
77 assertReplayPrefix(t, result.conn, data)
78}
79
80func TestSniffSmartTCPClassifiesTokenSpacePlaintextAsUnknown(t *testing.T) {
81 for _, data := range [][]byte{
82 []byte("USER anonymous\r\n"),
83 []byte("MAIL FROM:<a@example.com>\r\n"),
84 } {
85 result := sniffBytes(t, data, time.Second)
86 assert.Equal(t, dispatchUnknown, result.kind)
87 assertReplayPrefix(t, result.conn, data)
88 }
89}
90
91func TestSniffSmartTCPRejectsPROXY(t *testing.T) {
92 for _, tc := range []struct {
93 name string
94 data []byte
95 }{
96 {name: "v1", data: []byte("PROXY TCP4 192.0.2.1 192.0.2.2 12345 443\r\n")},
97 {name: "v2", data: append([]byte{0x0d, 0x0a, 0x0d, 0x0a, 0x00, 0x0d, 0x0a, 0x51, 0x55, 0x49, 0x54, 0x0a}, []byte{0x21, 0x11, 0x00, 0x0c}...)},
98 } {
99 t.Run(tc.name, func(t *testing.T) {
100 result := sniffBytes(t, tc.data, time.Second)
101 assert.Equal(t, dispatchRejectPROXY, result.kind)
102 assertReplayPrefix(t, result.conn, tc.data)
103 })
104 }
105}
106
107func TestSniffSmartTCPUnknownIsImmediate(t *testing.T) {
108 client, server := net.Pipe()
109 defer client.Close()
110
111 start := time.Now()
112 done := make(chan sniffOutcome, 1)
113 go func() {
114 result, err := sniffSmartTCP(server, time.Second)
115 done <- sniffOutcome{result: result, err: err}
116 }()
117
118 _, err := client.Write([]byte{0})
119 require.NoError(t, err)
120 select {
121 case outcome := <-done:
122 require.NoError(t, outcome.err)
123 assert.Equal(t, dispatchUnknown, outcome.result.kind)
124 assert.Less(t, time.Since(start), 500*time.Millisecond)
125 case <-time.After(500 * time.Millisecond):
126 t.Fatal("sniff did not classify unknown bytes immediately")
127 }
128}
129
130func TestSniffSmartTCPAmbiguousPartialsTimeOut(t *testing.T) {
131 for _, tc := range []struct {
132 name string
133 data []byte
134 }{
135 {name: "g", data: []byte("G")},
136 {name: "b", data: []byte("B")},
137 {name: "bar", data: []byte("BAR")},
138 {name: "pr", data: []byte("PR")},
139 {name: "tls_content_type", data: []byte{22}},
140 {name: "tls_version_major", data: []byte{22, 3}},
141 {name: "tls_version_minor", data: []byte{22, 3, 3}},
142 {name: "tls_record_length_partial", data: []byte{22, 3, 3, 0}},
143 } {
144 t.Run(tc.name, func(t *testing.T) {
145 result := sniffBytes(t, tc.data, 20*time.Millisecond)
146 assert.Equal(t, dispatchSniffTimeout, result.kind)
147 assertReplayPrefix(t, result.conn, tc.data)
148 })
149 }
150}
151
152func TestSniffSmartTCPReplayReturnsSniffedBytesBeforeUnderlyingBytes(t *testing.T) {
153 client, server := net.Pipe()
154 defer client.Close()
155
156 done := make(chan sniffOutcome, 1)
157 go func() {
158 result, err := sniffSmartTCP(server, time.Second)
159 done <- sniffOutcome{result: result, err: err}
160 }()
161 _, err := client.Write([]byte("GET / HTTP/1.1\r\n"))
162 require.NoError(t, err)
163 outcome := <-done
164 require.NoError(t, outcome.err)
165 result := outcome.result
166 writeDone := make(chan error, 1)
167 go func() {
168 _, err := client.Write([]byte("Host: example.com\r\n\r\n"))
169 writeDone <- err
170 }()
171
172 buf := make([]byte, len("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n"))
173 _, err = io.ReadFull(result.conn, buf)
174 require.NoError(t, err)
175 assert.Equal(t, "GET / HTTP/1.1\r\nHost: example.com\r\n\r\n", string(buf))
176 require.NoError(t, <-writeDone)
177}
178
179func TestSniffSmartTCPHandlesLongHTTPRequestLine(t *testing.T) {
180 client, server := net.Pipe()
181 defer client.Close()
182
183 longRequest := "GET /" + strings.Repeat("a", 512) + " HTTP/1.1\r\n"
184 writeDone := make(chan error, 1)
185 go func() {
186 _, err := client.Write([]byte(longRequest))
187 writeDone <- err
188 }()
189 result, err := sniffSmartTCP(server, time.Second)
190 require.NoError(t, err)
191 assert.Equal(t, dispatchHTTP, result.kind)
192 assertReplayPrefix(t, result.conn, []byte(longRequest))
193 require.NoError(t, <-writeDone)
194}
195
196func TestSniffSmartTCPClearsReadDeadlineBeforeHandoff(t *testing.T) {
197 conn := &deadlineConn{Conn: pipeConnWithData(t, []byte("GET / HTTP/1.1\r\n"))}
198 result, err := sniffSmartTCP(conn, time.Second)
199 require.NoError(t, err)
200 assert.Equal(t, dispatchHTTP, result.kind)
201 assert.True(t, conn.sawNonZeroReadDeadline)
202 assert.True(t, conn.sawClearedReadDeadline)
203}
204
205func sniffBytes(t *testing.T, data []byte, timeout time.Duration) smartTCPSniffResult {
206 t.Helper()
207 conn := pipeConnWithData(t, data)
208 result, err := sniffSmartTCP(conn, timeout)
209 require.NoError(t, err)
210 return result
211}
212
213func pipeConnWithData(t *testing.T, data []byte) net.Conn {
214 t.Helper()
215 client, server := net.Pipe()
216 done := make(chan struct{})
217 go func() {
218 defer close(done)
219 _, _ = client.Write(data)
220 }()
221 t.Cleanup(func() {
222 _ = client.Close()
223 _ = server.Close()
224 <-done
225 })
226 return server
227}
228
229func assertReplayPrefix(t *testing.T, conn net.Conn, want []byte) {
230 t.Helper()
231 buf := make([]byte, len(want))
232 _, err := io.ReadFull(conn, buf)
233 require.NoError(t, err)
234 assert.Equal(t, want, buf)
235}
236
237type sniffOutcome struct {
238 result smartTCPSniffResult
239 err error
240}
241
242type deadlineConn struct {
243 net.Conn
244 sawNonZeroReadDeadline bool
245 sawClearedReadDeadline bool
246}
247
248func (c *deadlineConn) SetReadDeadline(t time.Time) error {
249 if t.IsZero() {
250 c.sawClearedReadDeadline = true
251 } else {
252 c.sawNonZeroReadDeadline = true
253 }
254 return c.Conn.SetReadDeadline(t)
255}
256
257func TestSniffSmartTCPDoesNotTreatClosedPipeAsSuccess(t *testing.T) {
258 client, server := net.Pipe()
259 _ = client.Close()
260 _, err := sniffSmartTCP(server, time.Second)
261 require.Error(t, err)
262 assert.True(t, errors.Is(err, net.ErrClosed) || errors.Is(err, io.ErrClosedPipe), "expected closed connection error, got %v", err)
263}