1package main
2
3import (
4 "bufio"
5 "context"
6 "encoding/json"
7 "io"
8 "net"
9 "os"
10 "path/filepath"
11 "testing"
12 "time"
13
14 "github.com/stretchr/testify/require"
15)
16
17func TestTCPRelayPreservesBytesAndStops(t *testing.T) {
18 ctx, cancel := context.WithCancel(t.Context())
19 defer cancel()
20 r := &relay{ctx: ctx, cancel: cancel}
21 source, writer := net.Pipe()
22 reader, destination := net.Pipe()
23 defer source.Close()
24 defer writer.Close()
25 defer reader.Close()
26 defer destination.Close()
27 done := make(chan struct{})
28 go func() { defer close(done); r.copy(destination, source) }()
29 payload := []byte("fixture input in order")
30 go func() { writer.Write(payload); writer.Close() }()
31 reader.SetReadDeadline(time.Now().Add(time.Second))
32 got := make([]byte, len(payload))
33 _, err := io.ReadFull(reader, got)
34 require.NoError(t, err)
35 require.Equal(t, payload, got)
36 select {
37 case <-done:
38 case <-time.After(time.Second):
39 t.Fatal("relay did not stop after EOF")
40 }
41 require.Equal(t, uint64(len(payload)), r.bytes.Load())
42}
43
44// controlSocketPath returns a socket path in a short private directory: the
45// macOS t.TempDir path exceeds the 104-byte AF_UNIX limit.
46func controlSocketPath(t *testing.T) string {
47 t.Helper()
48 directory, err := os.MkdirTemp("/tmp", "relay")
49 require.NoError(t, err)
50 t.Cleanup(func() { _ = os.RemoveAll(directory) })
51 require.NoError(t, os.Chmod(directory, 0700))
52 return filepath.Join(directory, "control.sock")
53}
54
55func TestRelayShutdownClosesListeners(t *testing.T) {
56 ctx, cancel := context.WithCancel(t.Context())
57 r := &relay{ctx: ctx, cancel: cancel, connections: make(map[net.Conn]struct{})}
58 path := controlSocketPath(t)
59 done := make(chan error, 1)
60 go func() { done <- r.serve("127.0.0.1:0", path) }()
61 require.Eventually(t, func() bool { _, err := os.Stat(path); return err == nil }, time.Second, time.Millisecond)
62 cancel()
63 select {
64 case err := <-done:
65 require.NoError(t, err)
66 case <-time.After(time.Second):
67 t.Fatal("relay shutdown hung")
68 }
69 _, err := os.Stat(path)
70 require.True(t, os.IsNotExist(err))
71}
72
73func TestTCPRelayBlackoutHoldsBytes(t *testing.T) {
74 ctx, cancel := context.WithCancel(t.Context())
75 defer cancel()
76 r := &relay{ctx: ctx, cancel: cancel}
77 _, err := r.apply(command{Op: "blackout", TCP: true, Active: true})
78 require.NoError(t, err)
79 source, writer := net.Pipe()
80 reader, destination := net.Pipe()
81 defer source.Close()
82 defer writer.Close()
83 defer reader.Close()
84 defer destination.Close()
85 done := make(chan struct{})
86 go func() { defer close(done); r.copy(destination, source) }()
87 payload := []byte("held")
88 wrote := make(chan error, 1)
89 go func() { _, err := writer.Write(payload); wrote <- err }()
90 select {
91 case err := <-wrote:
92 t.Fatalf("blackout read the source: %v", err)
93 case <-time.After(50 * time.Millisecond):
94 }
95 require.Zero(t, r.bytes.Load())
96 _, err = r.apply(command{Op: "blackout", TCP: true, Active: false})
97 require.NoError(t, err)
98 reader.SetReadDeadline(time.Now().Add(time.Second))
99 got := make([]byte, len(payload))
100 _, err = io.ReadFull(reader, got)
101 require.NoError(t, err)
102 require.Equal(t, payload, got)
103 require.NoError(t, <-wrote)
104 cancel()
105 writer.Close()
106 <-done
107}
108
109// TestRelayUDPProxyInheritsActiveBlackout opens a UDP route while a blackout
110// is active and proves it drops datagrams until the blackout ends.
111func TestRelayUDPProxyInheritsActiveBlackout(t *testing.T) {
112 ctx, cancel := context.WithCancel(t.Context())
113 defer cancel()
114 r := &relay{ctx: ctx, cancel: cancel, remote: "127.0.0.1"}
115 defer r.close()
116 upstream, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
117 require.NoError(t, err)
118 defer upstream.Close()
119 _, err = r.apply(command{Op: "blackout", UDP: true, Active: true})
120 require.NoError(t, err)
121 port, err := r.apply(command{Op: "udp", Port: upstream.LocalAddr().(*net.UDPAddr).Port})
122 require.NoError(t, err)
123 client, err := net.DialUDP("udp", nil, &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: port.(int)})
124 require.NoError(t, err)
125 defer client.Close()
126 _, err = client.Write([]byte("blocked"))
127 require.NoError(t, err)
128 require.Eventually(t, func() bool { return r.proxies[0].Stats().ToServer.BlackoutDrops == 1 }, time.Second, time.Millisecond)
129 _, err = r.apply(command{Op: "blackout", UDP: true, Active: false})
130 require.NoError(t, err)
131 _, err = client.Write([]byte("open"))
132 require.NoError(t, err)
133 require.NoError(t, upstream.SetReadDeadline(time.Now().Add(time.Second)))
134 data := make([]byte, 16)
135 n, _, err := upstream.ReadFromUDP(data)
136 require.NoError(t, err)
137 require.Equal(t, "open", string(data[:n]))
138}
139
140func TestRelayDisconnectClosesConnections(t *testing.T) {
141 ctx, cancel := context.WithCancel(t.Context())
142 defer cancel()
143 r := &relay{ctx: ctx, cancel: cancel, connections: make(map[net.Conn]struct{})}
144 local, peer := net.Pipe()
145 defer peer.Close()
146 r.connections[local] = struct{}{}
147 _, err := r.apply(command{Op: "disconnect"})
148 require.NoError(t, err)
149 peer.SetReadDeadline(time.Now().Add(time.Second))
150 _, err = peer.Read(make([]byte, 1))
151 require.ErrorIs(t, err, io.EOF)
152}
153
154func TestRelayControlRoundTrip(t *testing.T) {
155 ctx, cancel := context.WithCancel(t.Context())
156 defer cancel()
157 r := &relay{ctx: ctx, cancel: cancel, connections: make(map[net.Conn]struct{})}
158 path := controlSocketPath(t)
159 done := make(chan error, 1)
160 go func() { done <- r.serve("127.0.0.1:0", path) }()
161 require.Eventually(t, func() bool { _, err := os.Stat(path); return err == nil }, time.Second, time.Millisecond)
162 send := func(c command) map[string]any {
163 conn, err := net.Dial("unix", path)
164 require.NoError(t, err)
165 defer conn.Close()
166 require.NoError(t, json.NewEncoder(conn).Encode(c))
167 line, err := bufio.NewReader(conn).ReadBytes('\n')
168 require.NoError(t, err)
169 var reply map[string]any
170 require.NoError(t, json.Unmarshal(line, &reply))
171 return reply
172 }
173 reply := send(command{ID: 7, Op: "blackout", TCP: true, Active: true})
174 require.EqualValues(t, 7, reply["id"])
175 require.NotContains(t, reply, "error")
176 require.True(t, r.blackout.Load())
177 reply = send(command{ID: 8, Op: "unknown"})
178 require.Equal(t, "unknown command", reply["error"])
179 reply = send(command{ID: 9, Op: "stop"})
180 require.NotContains(t, reply, "error")
181 select {
182 case err := <-done:
183 require.NoError(t, err)
184 case <-time.After(time.Second):
185 t.Fatal("stop command did not shut the relay down")
186 }
187}
188
189func TestRelayRejectsInvalidControl(t *testing.T) {
190 ctx, cancel := context.WithCancel(t.Context())
191 defer cancel()
192 r := &relay{ctx: ctx, cancel: cancel}
193 for _, c := range []command{{Op: "unknown"}, {Op: "udp", Port: 0}, {Op: "udp", Port: 65536}} {
194 _, err := r.apply(c)
195 require.Error(t, err)
196 }
197 _, err := r.apply(command{Op: "stop"})
198 require.NoError(t, err)
199 require.ErrorIs(t, ctx.Err(), context.Canceled)
200}