| 1 | package tunnel |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "fmt" |
| 6 | "io" |
| 7 | "net/http" |
| 8 | "net/http/httptest" |
| 9 | "net/url" |
| 10 | "strings" |
| 11 | "sync" |
| 12 | "testing" |
| 13 | "time" |
| 14 | |
| 15 | "github.com/coder/websocket" |
| 16 | ) |
| 17 | |
| 18 | // The whole loop, no router required: a local "instance" (plain HTTP + |
| 19 | // a WebSocket echo), a gateway, and a client tunnelling between them. Verifies |
| 20 | // ordinary requests, header passthrough, and — the part that matters for the |
| 21 | // chat — a live WebSocket upgraded end-to-end through the tunnel. |
| 22 | func TestTunnelEndToEnd(t *testing.T) { |
| 23 | backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 24 | switch r.URL.Path { |
| 25 | case "/api/version": |
| 26 | w.Header().Set("Content-Type", "application/json") |
| 27 | io.WriteString(w, `{"build":"test","version":"test-1"}`) |
| 28 | case "/ws": |
| 29 | c, err := websocket.Accept(w, r, nil) |
| 30 | if err != nil { |
| 31 | return |
| 32 | } |
| 33 | defer c.CloseNow() |
| 34 | ctx := r.Context() |
| 35 | for { |
| 36 | typ, data, err := c.Read(ctx) |
| 37 | if err != nil { |
| 38 | return |
| 39 | } |
| 40 | if err := c.Write(ctx, typ, append([]byte("echo:"), data...)); err != nil { |
| 41 | return |
| 42 | } |
| 43 | } |
| 44 | default: |
| 45 | w.Header().Set("X-Seen-Host", r.Host) |
| 46 | io.WriteString(w, "hello through the tunnel") |
| 47 | } |
| 48 | })) |
| 49 | defer backend.Close() |
| 50 | |
| 51 | gw := &Gateway{Domain: "test.local", DialTimeout: 5 * time.Second} |
| 52 | registered := make(chan string, 1) |
| 53 | gw.OnRegister = func(host string) error { registered <- host; return nil } |
| 54 | closed := make(chan string, 1) |
| 55 | gw.OnClose = func(host string) { closed <- host } |
| 56 | |
| 57 | gwSrv := httptest.NewServer(gw) |
| 58 | defer gwSrv.Close() |
| 59 | u, _ := url.Parse(gwSrv.URL) |
| 60 | gw.ControlHost = u.Hostname() // ServeHTTP compares hosts port-stripped |
| 61 | |
| 62 | ctx, cancel := context.WithCancel(context.Background()) |
| 63 | defer cancel() |
| 64 | cl := &Client{Gateway: gwSrv.URL, LocalAddr: strings.TrimPrefix(backend.URL, "http://")} |
| 65 | ready := make(chan string, 1) |
| 66 | clientDone := make(chan error, 1) |
| 67 | go func() { clientDone <- cl.Run(ctx, func(h string) { ready <- h }) }() |
| 68 | |
| 69 | var host string |
| 70 | select { |
| 71 | case host = <-ready: |
| 72 | case err := <-clientDone: |
| 73 | t.Fatalf("client exited early: %v", err) |
| 74 | case <-time.After(5 * time.Second): |
| 75 | t.Fatal("no ready") |
| 76 | } |
| 77 | if !strings.HasPrefix(host, "try-") || !strings.HasSuffix(host, ".test.local") { |
| 78 | t.Fatalf("unexpected host %q", host) |
| 79 | } |
| 80 | if got := <-registered; got != host { |
| 81 | t.Fatalf("OnRegister got %q, ready got %q", got, host) |
| 82 | } |
| 83 | |
| 84 | // A plain request, addressed to the slug host through the gateway. |
| 85 | hostRewrite := func(rt http.RoundTripper) http.RoundTripper { |
| 86 | return roundTripperFunc(func(r *http.Request) (*http.Response, error) { |
| 87 | r.Host = host |
| 88 | return rt.RoundTrip(r) |
| 89 | }) |
| 90 | } |
| 91 | httpc := &http.Client{Transport: hostRewrite(http.DefaultTransport.(*http.Transport).Clone())} |
| 92 | resp, err := httpc.Get(gwSrv.URL + "/anything") |
| 93 | if err != nil { |
| 94 | t.Fatalf("GET: %v", err) |
| 95 | } |
| 96 | body, _ := io.ReadAll(resp.Body) |
| 97 | resp.Body.Close() |
| 98 | if string(body) != "hello through the tunnel" { |
| 99 | t.Fatalf("body = %q", body) |
| 100 | } |
| 101 | if got := resp.Header.Get("X-Seen-Host"); got != host { |
| 102 | t.Fatalf("backend saw Host %q, want %q (public hostname must reach the instance)", got, host) |
| 103 | } |
| 104 | |
| 105 | // A WebSocket, upgraded end-to-end through the tunnel. |
| 106 | wsCtx, wsCancel := context.WithTimeout(context.Background(), 5*time.Second) |
| 107 | defer wsCancel() |
| 108 | c, _, err := websocket.Dial(wsCtx, "ws://"+u.Host+"/ws", &websocket.DialOptions{ |
| 109 | HTTPClient: &http.Client{Transport: hostRewrite(http.DefaultTransport.(*http.Transport).Clone())}, |
| 110 | }) |
| 111 | if err != nil { |
| 112 | t.Fatalf("ws dial: %v", err) |
| 113 | } |
| 114 | defer c.CloseNow() |
| 115 | if err := c.Write(wsCtx, websocket.MessageText, []byte("ping")); err != nil { |
| 116 | t.Fatalf("ws write: %v", err) |
| 117 | } |
| 118 | _, data, err := c.Read(wsCtx) |
| 119 | if err != nil { |
| 120 | t.Fatalf("ws read: %v", err) |
| 121 | } |
| 122 | if string(data) != "echo:ping" { |
| 123 | t.Fatalf("ws echo = %q", data) |
| 124 | } |
| 125 | |
| 126 | // Tearing the client down retires the route. |
| 127 | cancel() |
| 128 | select { |
| 129 | case got := <-closed: |
| 130 | if got != host { |
| 131 | t.Fatalf("OnClose got %q", got) |
| 132 | } |
| 133 | case <-time.After(5 * time.Second): |
| 134 | t.Fatal("route never closed") |
| 135 | } |
| 136 | } |
| 137 | |
| 138 | // The per-IP limiter refuses the sixth tunnel in an hour from one address. |
| 139 | func TestTunnelRateLimit(t *testing.T) { |
| 140 | gw := &Gateway{Domain: "test.local"} |
| 141 | for range 5 { |
| 142 | if err := gw.admit("203.0.113.9"); err != nil { |
| 143 | t.Fatalf("admit: %v", err) |
| 144 | } |
| 145 | gw.byIP["203.0.113.9"] = append(gw.byIP["203.0.113.9"], time.Now()) |
| 146 | } |
| 147 | if err := gw.admit("203.0.113.9"); err == nil { |
| 148 | t.Fatal("sixth tunnel admitted") |
| 149 | } |
| 150 | if err := gw.admit("203.0.113.10"); err != nil { |
| 151 | t.Fatalf("other address blocked: %v", err) |
| 152 | } |
| 153 | } |
| 154 | |
| 155 | // A backend that is not an Eleven instance is refused before its hostname is |
| 156 | // ever published. |
| 157 | func TestTunnelRejectsNonEleven(t *testing.T) { |
| 158 | backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 159 | io.WriteString(w, "just a web server") |
| 160 | })) |
| 161 | defer backend.Close() |
| 162 | |
| 163 | gw := &Gateway{Domain: "test.local", DialTimeout: 5 * time.Second} |
| 164 | registered := false |
| 165 | gw.OnRegister = func(string) error { registered = true; return nil } |
| 166 | gwSrv := httptest.NewServer(gw) |
| 167 | defer gwSrv.Close() |
| 168 | u, _ := url.Parse(gwSrv.URL) |
| 169 | gw.ControlHost = u.Hostname() |
| 170 | |
| 171 | cl := &Client{Gateway: gwSrv.URL, LocalAddr: strings.TrimPrefix(backend.URL, "http://")} |
| 172 | err := cl.Run(context.Background(), func(string) {}) |
| 173 | if err == nil || !strings.Contains(err.Error(), "only Eleven instances") { |
| 174 | t.Fatalf("expected probe refusal, got %v", err) |
| 175 | } |
| 176 | if registered { |
| 177 | t.Fatal("route was published for a non-Eleven backend") |
| 178 | } |
| 179 | } |
| 180 | |
| 181 | // The per-tunnel connection cap refuses dial-backs once live + in-flight |
| 182 | // connections hit the limit. |
| 183 | func TestTunnelConnCap(t *testing.T) { |
| 184 | g := &Gateway{MaxConns: 2} |
| 185 | tun := &liveTunnel{closed: make(chan struct{})} |
| 186 | tun.conns.Store(2) |
| 187 | if _, err := g.dialBack(context.Background(), tun); err == nil || |
| 188 | !strings.Contains(err.Error(), "connection cap") { |
| 189 | t.Fatalf("expected cap refusal, got %v", err) |
| 190 | } |
| 191 | if got := tun.conns.Load(); got != 2 { |
| 192 | t.Fatalf("refused dial leaked a slot: %d", got) |
| 193 | } |
| 194 | } |
| 195 | |
| 196 | // startTunnelPair boots a backend serving handler, the given gateway, and a |
| 197 | // client tunnelling between them, returning the gateway server and the public |
| 198 | // slug host once it's live. Everything tears down with the test. |
| 199 | func startTunnelPair(t *testing.T, gw *Gateway, handler http.HandlerFunc) (gwSrv *httptest.Server, host string) { |
| 200 | t.Helper() |
| 201 | backend := httptest.NewServer(handler) |
| 202 | t.Cleanup(backend.Close) |
| 203 | gwSrv = httptest.NewServer(gw) |
| 204 | t.Cleanup(gwSrv.Close) |
| 205 | u, _ := url.Parse(gwSrv.URL) |
| 206 | gw.ControlHost = u.Hostname() |
| 207 | |
| 208 | ctx, cancel := context.WithCancel(context.Background()) |
| 209 | t.Cleanup(cancel) |
| 210 | cl := &Client{Gateway: gwSrv.URL, LocalAddr: strings.TrimPrefix(backend.URL, "http://")} |
| 211 | ready := make(chan string, 1) |
| 212 | clientDone := make(chan error, 1) |
| 213 | go func() { clientDone <- cl.Run(ctx, func(h string) { ready <- h }) }() |
| 214 | select { |
| 215 | case host = <-ready: |
| 216 | case err := <-clientDone: |
| 217 | t.Fatalf("client exited early: %v", err) |
| 218 | case <-time.After(5 * time.Second): |
| 219 | t.Fatal("no ready") |
| 220 | } |
| 221 | return gwSrv, host |
| 222 | } |
| 223 | |
| 224 | // hostRewriteTransport addresses requests to the slug host through the |
| 225 | // gateway's test listener. |
| 226 | func hostRewriteTransport(host string) http.RoundTripper { |
| 227 | rt := http.DefaultTransport.(*http.Transport).Clone() |
| 228 | return roundTripperFunc(func(r *http.Request) (*http.Response, error) { |
| 229 | r.Host = host |
| 230 | return rt.RoundTrip(r) |
| 231 | }) |
| 232 | } |
| 233 | |
| 234 | // versionThen serves the probe's /api/version and hands everything else to |
| 235 | // next — the minimum for a backend to pass the Eleven door check. |
| 236 | func versionThen(next http.HandlerFunc) http.HandlerFunc { |
| 237 | return func(w http.ResponseWriter, r *http.Request) { |
| 238 | if r.URL.Path == "/api/version" { |
| 239 | w.Header().Set("Content-Type", "application/json") |
| 240 | io.WriteString(w, `{"build":"test","version":"test-1"}`) |
| 241 | return |
| 242 | } |
| 243 | next(w, r) |
| 244 | } |
| 245 | } |
| 246 | |
| 247 | // A first-load burst wider than the per-tunnel connection cap queues for a |
| 248 | // free connection instead of failing. The no-build shell fetches 40+ modules |
| 249 | // in one burst over HTTP/2, and a single 502 permanently fails the page's ES |
| 250 | // module graph — this is what made `lchat try` unusable in a browser. |
| 251 | func TestTunnelBurstQueuesBeyondConnCap(t *testing.T) { |
| 252 | gw := &Gateway{Domain: "test.local", DialTimeout: 5 * time.Second, MaxConns: 4} |
| 253 | gwSrv, host := startTunnelPair(t, gw, versionThen(func(w http.ResponseWriter, r *http.Request) { |
| 254 | time.Sleep(20 * time.Millisecond) // hold connections open so the burst overlaps |
| 255 | io.WriteString(w, "module") |
| 256 | })) |
| 257 | |
| 258 | httpc := &http.Client{Transport: hostRewriteTransport(host)} |
| 259 | var wg sync.WaitGroup |
| 260 | codes := make([]int, 24) |
| 261 | for i := range codes { |
| 262 | wg.Add(1) |
| 263 | go func() { |
| 264 | defer wg.Done() |
| 265 | resp, err := httpc.Get(fmt.Sprintf("%s/static/m%d.js", gwSrv.URL, i)) |
| 266 | if err != nil { |
| 267 | codes[i] = -1 |
| 268 | return |
| 269 | } |
| 270 | io.Copy(io.Discard, resp.Body) |
| 271 | resp.Body.Close() |
| 272 | codes[i] = resp.StatusCode |
| 273 | }() |
| 274 | } |
| 275 | wg.Wait() |
| 276 | for i, c := range codes { |
| 277 | if c != http.StatusOK { |
| 278 | t.Fatalf("request %d: status %d — a burst above the cap must queue, not fail", i, c) |
| 279 | } |
| 280 | } |
| 281 | } |
| 282 | |
| 283 | // A visitor abandoning a request mid-dial (reload, closed tab) fails that one |
| 284 | // dial only — it must never drop the tunnel for everyone else. |
| 285 | func TestTunnelSurvivesAbandonedRequest(t *testing.T) { |
| 286 | closed := make(chan string, 1) |
| 287 | gw := &Gateway{Domain: "test.local", DialTimeout: 5 * time.Second} |
| 288 | gw.OnClose = func(host string) { closed <- host } |
| 289 | gwSrv, host := startTunnelPair(t, gw, versionThen(func(w http.ResponseWriter, r *http.Request) { |
| 290 | io.WriteString(w, "ok") |
| 291 | })) |
| 292 | |
| 293 | gw.mu.Lock() |
| 294 | tun := gw.tunnels[host] |
| 295 | gw.mu.Unlock() |
| 296 | if tun == nil { |
| 297 | t.Fatal("no live tunnel") |
| 298 | } |
| 299 | |
| 300 | abandoned, cancel := context.WithCancel(context.Background()) |
| 301 | cancel() |
| 302 | if conn, err := gw.dialBack(abandoned, tun); err == nil { |
| 303 | conn.Close() // the dial-back can win the race; the tunnel must survive either way |
| 304 | } |
| 305 | |
| 306 | select { |
| 307 | case h := <-closed: |
| 308 | t.Fatalf("abandoned request dropped the whole tunnel (%s)", h) |
| 309 | default: |
| 310 | } |
| 311 | httpc := &http.Client{Transport: hostRewriteTransport(host)} |
| 312 | resp, err := httpc.Get(gwSrv.URL + "/after") |
| 313 | if err != nil { |
| 314 | t.Fatalf("tunnel unusable after an abandoned request: %v", err) |
| 315 | } |
| 316 | defer resp.Body.Close() |
| 317 | if resp.StatusCode != http.StatusOK { |
| 318 | t.Fatalf("tunnel unusable after an abandoned request: HTTP %d", resp.StatusCode) |
| 319 | } |
| 320 | } |
| 321 | |
| 322 | type roundTripperFunc func(*http.Request) (*http.Response, error) |
| 323 | |
| 324 | func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } |
| 325 | |
amadan