1// Command relay is disposable acceptance infrastructure, never a production proxy.
2package main
3
4import (
5 "bufio"
6 "bytes"
7 "context"
8 "encoding/json"
9 "errors"
10 "flag"
11 "fmt"
12 "io"
13 "net"
14 "os"
15 "os/signal"
16 "path/filepath"
17 "sync"
18 "sync/atomic"
19 "syscall"
20 "time"
21
22 "github.com/bnema/vev/internal/adapters/quicnettest"
23)
24
25type command struct {
26 ID uint64 `json:"id"`
27 Op string `json:"op"`
28 Port int `json:"port,omitempty"`
29 UDP bool `json:"udp,omitempty"`
30 TCP bool `json:"tcp,omitempty"`
31 Active bool `json:"active,omitempty"`
32}
33
34type relay struct {
35 ctx context.Context
36 cancel context.CancelFunc
37 remote string
38 degraded bool
39 blackout atomic.Bool
40 bytes atomic.Uint64
41 mu sync.Mutex
42 connections map[net.Conn]struct{}
43 proxies []*quicnettest.Proxy
44 udpBlackout bool
45 wg sync.WaitGroup
46}
47
48func main() {
49 remote := flag.String("remote", "remote", "fixture upstream host")
50 control := flag.String("control", "", "required private Unix control path")
51 tcp := flag.String("tcp", ":2222", "SSH listener")
52 degraded := flag.Bool("degraded", false, "enable seeded degraded profile")
53 flag.Parse()
54 if *control == "" {
55 fmt.Fprintln(os.Stderr, "-control is required")
56 os.Exit(2)
57 }
58 ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
59 r := &relay{ctx: ctx, cancel: cancel, remote: *remote, degraded: *degraded, connections: make(map[net.Conn]struct{})}
60 if err := r.serve(*tcp, *control); err != nil {
61 fmt.Fprintln(os.Stderr, err)
62 os.Exit(1)
63 }
64}
65
66func (r *relay) serve(addr, path string) error {
67 listener, err := net.Listen("tcp", addr)
68 if err != nil {
69 return err
70 }
71 defer listener.Close()
72 parent, err := os.Stat(filepath.Dir(path))
73 if err != nil {
74 return err
75 }
76 if !parent.IsDir() || parent.Mode().Perm() != 0700 {
77 return errors.New("control socket requires a private 0700 directory")
78 }
79 owner, ok := parent.Sys().(*syscall.Stat_t)
80 if !ok || owner.Uid != uint32(os.Getuid()) {
81 return errors.New("control directory must belong to the current user")
82 }
83 control, err := net.Listen("unix", path)
84 if err != nil {
85 return err
86 }
87 defer control.Close()
88 watchDone := make(chan struct{})
89 defer close(watchDone)
90 go func() {
91 select {
92 case <-r.ctx.Done():
93 listener.Close()
94 control.Close()
95 case <-watchDone:
96 }
97 }()
98 defer os.Remove(path)
99 if err := os.Chmod(path, 0600); err != nil {
100 return err
101 }
102 defer func() { listener.Close(); r.close() }()
103 r.wg.Add(1)
104 go func() {
105 defer r.wg.Done()
106 for {
107 conn, err := listener.Accept()
108 if err != nil {
109 return
110 }
111 r.mu.Lock()
112 if len(r.connections) >= 64 {
113 r.mu.Unlock()
114 conn.Close()
115 continue
116 }
117 r.connections[conn] = struct{}{}
118 r.mu.Unlock()
119 r.wg.Add(1)
120 go func() { defer r.wg.Done(); r.forward(conn) }()
121 }
122 }()
123 for {
124 conn, err := control.Accept()
125 if err != nil {
126 if r.ctx.Err() != nil {
127 return nil
128 }
129 return err
130 }
131 // Commands are bounded and handled serially; a stalled caller gets 2s.
132 conn.SetDeadline(time.Now().Add(2 * time.Second))
133 scanner := bufio.NewScanner(conn)
134 scanner.Buffer(make([]byte, 4096), 4096)
135 if scanner.Scan() {
136 var c command
137 err := json.Unmarshal(scanner.Bytes(), &c)
138 var result any
139 if err == nil {
140 result, err = r.apply(c)
141 }
142 reply := map[string]any{"id": c.ID, "result": result}
143 if err != nil {
144 reply["error"] = err.Error()
145 }
146 json.NewEncoder(conn).Encode(reply)
147 }
148 conn.Close()
149 if r.ctx.Err() != nil {
150 listener.Close()
151 return nil
152 }
153 }
154}
155
156func (r *relay) apply(c command) (any, error) {
157 switch c.Op {
158 case "udp":
159 if c.Port <= 0 || c.Port > 65535 {
160 return nil, errors.New("invalid port")
161 }
162 upstream, err := net.ResolveUDPAddr("udp", net.JoinHostPort(r.remote, fmt.Sprint(c.Port)))
163 if err != nil {
164 return nil, err
165 }
166 link := quicnettest.LinkConfig{}
167 if r.degraded {
168 link = quicnettest.LinkConfig{BaseLatency: 150 * time.Millisecond, Jitter: 50 * time.Millisecond, LossPercent: 5}
169 }
170 r.mu.Lock()
171 defer r.mu.Unlock()
172 if len(r.proxies) >= 32 {
173 return nil, errors.New("UDP listener limit")
174 }
175 p, err := quicnettest.New(quicnettest.Config{ServerAddr: upstream, ListenAddr: &net.UDPAddr{IP: net.IPv4zero}, Seed: 317, ToServer: link, ToClient: link})
176 if err != nil {
177 return nil, err
178 }
179 p.SetBlackout(r.udpBlackout)
180 r.proxies = append(r.proxies, p)
181 return p.Addr().Port, nil
182 case "blackout":
183 if c.TCP {
184 r.blackout.Store(c.Active)
185 }
186 if c.UDP {
187 r.mu.Lock()
188 r.udpBlackout = c.Active
189 for _, p := range r.proxies {
190 p.SetBlackout(c.Active)
191 }
192 r.mu.Unlock()
193 }
194 return nil, nil
195 case "disconnect":
196 r.mu.Lock()
197 for conn := range r.connections {
198 conn.Close()
199 }
200 r.mu.Unlock()
201 return nil, nil
202 case "stats":
203 r.mu.Lock()
204 defer r.mu.Unlock()
205 stats := make([]quicnettest.Stats, 0, len(r.proxies))
206 for _, p := range r.proxies {
207 stats = append(stats, p.Stats())
208 }
209 return map[string]any{"udp": stats, "tcp_bytes": r.bytes.Load()}, nil
210 case "stop":
211 r.cancel()
212 return nil, nil
213 default:
214 return nil, errors.New("unknown command")
215 }
216}
217
218func (r *relay) forward(client net.Conn) {
219 defer func() { client.Close(); r.mu.Lock(); delete(r.connections, client); r.mu.Unlock() }()
220 server, err := (&net.Dialer{Timeout: 5 * time.Second}).DialContext(r.ctx, "tcp", net.JoinHostPort(r.remote, "22"))
221 if err != nil {
222 return
223 }
224 defer server.Close()
225 r.mu.Lock()
226 r.connections[server] = struct{}{}
227 r.mu.Unlock()
228 defer func() { r.mu.Lock(); delete(r.connections, server); r.mu.Unlock() }()
229 done := make(chan struct{})
230 go func() { r.copy(server, client); server.Close(); client.Close(); close(done) }()
231 r.copy(client, server)
232 server.Close()
233 client.Close()
234 <-done
235}
236
237func (r *relay) wait(delay time.Duration) bool {
238 timer := time.NewTimer(delay)
239 defer timer.Stop()
240 select {
241 case <-timer.C:
242 return true
243 case <-r.ctx.Done():
244 return false
245 }
246}
247
248func (r *relay) copy(dst, src net.Conn) {
249 buf := make([]byte, 32<<10)
250 for {
251 for r.blackout.Load() {
252 if !r.wait(10 * time.Millisecond) {
253 return
254 }
255 }
256 n, err := src.Read(buf)
257 if n > 0 {
258 for r.blackout.Load() {
259 if !r.wait(10 * time.Millisecond) {
260 return
261 }
262 }
263 if r.degraded && !r.wait(150*time.Millisecond+time.Duration(n)*time.Second/(64<<10)) {
264 return
265 }
266 written, writeErr := io.Copy(dst, bytes.NewReader(buf[:n]))
267 r.bytes.Add(uint64(written))
268 if writeErr != nil {
269 return
270 }
271 }
272 if err != nil {
273 return
274 }
275 }
276}
277
278func (r *relay) close() {
279 r.cancel()
280 r.mu.Lock()
281 for conn := range r.connections {
282 conn.Close()
283 }
284 for _, p := range r.proxies {
285 p.Close()
286 }
287 r.mu.Unlock()
288 r.wg.Wait()
289}