tinkerdown/proxy_ws_e2e_test.go at main · livetemplate/tinkerdown · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
//go:build !ci

package tinkerdown_test

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/chromedp/cdproto/network"
	"github.com/chromedp/cdproto/runtime"
	"github.com/chromedp/chromedp"
	"github.com/gorilla/websocket"
)

// TestProxyRoute_WebSocketUpgradeE2E is the real-browser regression for #257.
//
// It stands up a strict same-origin WebSocket upstream, fronts it with a
// tinkerdown `routes: type: proxy` route, drives a real browser to the
// proxied page, and asserts via CDP network events that the WS handshake
// returns 101 (not 403) and that a frame actually flows. The acceptance
// signal is deliberately the handshake status + frame delivery — NOT
// liveTemplateClient.isReady(), which returns true under the silent HTTP
// fallback that masked this bug in production.
//
// Per project convention this captures all four diagnostic channels:
// (1) browser console logs, (2) server logs (both the upstream's view of
// the forwarded request and tinkerdown's own log output), (3) WebSocket
// handshake status + frames via CDP, (4) rendered HTML.
func TestProxyRoute_WebSocketUpgradeE2E(t *testing.T) {
	const wsMarker = "WS_PUSH_MARKER_257"

	// --- Capture tinkerdown's package-level log output (server logs, hop 1).
	// This mutates the process-global logger, so it is only safe while this
	// test runs serially. SERIAL — never add t.Parallel() to this test (or
	// run it with a parallel sibling that also redirects log output); the
	// global redirect would race. Restored on cleanup.
	var tdLog syncBuffer
	prevOut := log.Writer()
	prevFlags := log.Flags()
	prevPrefix := log.Prefix()
	log.SetOutput(&tdLog)
	t.Cleanup(func() { log.SetOutput(prevOut); log.SetFlags(prevFlags); log.SetPrefix(prevPrefix) })

	// --- Upstream: strict same-origin WS check (gorilla default / livetemplate
	// prod default). Serves an HTML page that opens a WS back to the same host,
	// and on the WS path upgrades + pushes one frame. Logs every request so the
	// upstream's view of the forwarded Origin/Upgrade headers is visible (server
	// logs, hop 2 — the decisive diagnostic for this bug).
	upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool {
		return r.Header.Get("Origin") == "http://"+r.Host
	}}
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Logf("[upstream] %s %s Host=%q Origin=%q Upgrade=%q", r.Method, r.URL.Path, r.Host, r.Header.Get("Origin"), r.Header.Get("Upgrade"))
		if strings.EqualFold(r.Header.Get("Upgrade"), "websocket") {
			conn, err := upgrader.Upgrade(w, r, nil)
			if err != nil {
				t.Logf("[upstream] WS upgrade rejected: %v", err)
				return // upgrader already wrote the 403
			}
			defer conn.Close()
			if err := conn.WriteMessage(websocket.TextMessage, []byte(wsMarker)); err != nil {
				t.Logf("[upstream] WS write: %v", err)
				return
			}
			conn.SetReadDeadline(time.Now().Add(3 * time.Second))
			_, _, _ = conn.ReadMessage() // drain the client's "hi" / wait for close
			return
		}
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		_, _ = fmt.Fprint(w, `<!doctype html><html><body>
<div id="status">connecting</div>
<script>
  var proto = location.protocol === 'https:' ? 'wss://' : 'ws://';
  var ws = new WebSocket(proto + location.host + '/proxy/ws');
  ws.onopen = function () { console.log('ws open'); ws.send('hi'); };
  ws.onmessage = function (e) { document.getElementById('status').textContent = 'WS_FRAME:' + e.data; };
  ws.onerror = function () { document.getElementById('status').textContent = 'WS_ERROR'; console.log('ws error'); };
  ws.onclose = function (e) { console.log('ws close code=' + e.code); };
</script>
</body></html>`)
	}))
	t.Cleanup(upstream.Close)

	// --- tinkerdown fixture: a proxy route fronting the upstream.
	tdDir := authorProxyFixture(t, upstream.URL)
	_, tdURL := startTinkerdown(t, tdDir)

	// --- Drive a real browser, capturing all four channels.
	chromeCtx, cleanup := SetupDockerChrome(t, 60*time.Second)
	t.Cleanup(cleanup)
	ctx := chromeCtx.Context

	var (
		mu              sync.Mutex
		consoleLogs     []string
		handshakeStatus []int64
		framesReceived  []string
		frameErrors     []string
	)
	chromedp.ListenTarget(ctx, func(ev any) {
		mu.Lock()
		defer mu.Unlock()
		switch e := ev.(type) {
		case *runtime.EventConsoleAPICalled:
			for _, arg := range e.Args {
				consoleLogs = append(consoleLogs, string(arg.Value))
			}
		case *network.EventWebSocketHandshakeResponseReceived:
			if e.Response != nil {
				handshakeStatus = append(handshakeStatus, e.Response.Status)
			}
		case *network.EventWebSocketFrameReceived:
			if e.Response != nil {
				framesReceived = append(framesReceived, e.Response.PayloadData)
			}
		case *network.EventWebSocketFrameError:
			frameErrors = append(frameErrors, e.ErrorMessage)
		}
	})

	url := ConvertURLForDockerChrome(tdURL)
	t.Logf("Test server URL: %s (Docker: %s)", tdURL, url)

	var status, htmlContent string
	runErr := chromedp.Run(ctx,
		network.Enable(),
		chromedp.Navigate(url+"/proxy/page"),
		// Wait for the WS frame to land and update the DOM.
		chromedp.WaitVisible(`#status`, chromedp.ByID),
		waitForStatus(`#status`, "WS_FRAME:", 5*time.Second),
		chromedp.Text(`#status`, &status, chromedp.ByID),
		chromedp.OuterHTML("html", &htmlContent),
	)

	// Hold the lock across the diagnostic dump and the assertions below —
	// the CDP listener goroutine may still append events (e.g. ws close).
	mu.Lock()
	defer mu.Unlock()

	// Dump all four diagnostic channels unconditionally — cheap, and
	// invaluable on any failure below.
	t.Logf("[console] %v", consoleLogs)
	t.Logf("[handshake statuses] %v", handshakeStatus)
	t.Logf("[frames] %v", framesReceived)
	t.Logf("[frame errors] %v", frameErrors)
	t.Logf("[tinkerdown log]\n%s", tdLog.String())
	t.Logf("[html] %s", firstChars(htmlContent, 2000))
	if runErr != nil {
		t.Fatalf("chromedp.Run: %v", runErr)
	}

	// (1) No 403, and a 101 Switching Protocols was observed.
	saw101 := false
	for _, s := range handshakeStatus {
		if s == 403 {
			t.Errorf("WS handshake returned 403 — proxy did not rewrite Origin (the #257 bug)")
		}
		if s == 101 {
			saw101 = true
		}
	}
	if !saw101 {
		t.Errorf("expected a 101 Switching Protocols handshake, got statuses %v", handshakeStatus)
	}

	// (2) A WS frame carrying the upstream marker actually flowed.
	frameSeen := false
	for _, f := range framesReceived {
		if strings.Contains(f, wsMarker) {
			frameSeen = true
		}
	}
	if !frameSeen {
		t.Errorf("expected a WS frame containing %q, got frames %v", wsMarker, framesReceived)
	}

	// (3) Rendered HTML reflects the pushed frame (not the HTTP fallback).
	if !strings.Contains(status, wsMarker) {
		t.Errorf("#status = %q, want it to contain the pushed marker %q", status, wsMarker)
	}
}

// authorProxyFixture writes a tinkerdown site with a `routes: type: proxy`
// route at /proxy/ pointing at the given upstream, plus a home page so
// Discover succeeds. Returns the temp dir tinkerdown should serve from.
func authorProxyFixture(t *testing.T, upstreamURL string) string {
	t.Helper()
	dir := t.TempDir()
	// %q quotes the upstream URL so the helper stays safe if a future test
	// reuses it with a URL containing YAML-significant characters.
	cfg := fmt.Sprintf("title: \"Proxy WS E2E\"\nroutes:\n  - pattern: \"/proxy/\"\n    type: proxy\n    upstream: %q\n", upstreamURL)
	if err := os.WriteFile(filepath.Join(dir, "tinkerdown.yaml"), []byte(cfg), 0o644); err != nil {
		t.Fatal(err)
	}
	body := "---\ntitle: \"Home\"\n---\n\n# Home\n"
	if err := os.WriteFile(filepath.Join(dir, "index.md"), []byte(body), 0o644); err != nil {
		t.Fatal(err)
	}
	return dir
}

// waitForStatus polls an element's textContent until it has the wanted
// prefix or the timeout elapses. Avoids a fixed Sleep race on the WS frame.
func waitForStatus(sel, wantPrefix string, timeout time.Duration) chromedp.Action {
	return chromedp.ActionFunc(func(ctx context.Context) error {
		ctx, cancel := context.WithTimeout(ctx, timeout)
		defer cancel()
		ticker := time.NewTicker(100 * time.Millisecond)
		defer ticker.Stop()
		for {
			var txt string
			if err := chromedp.Text(sel, &txt, chromedp.ByID).Do(ctx); err == nil {
				// WS_ERROR is a settled state — stop polling and let the
				// assertions report it. Only a real timeout is an error.
				if strings.HasPrefix(txt, wantPrefix) || strings.HasPrefix(txt, "WS_ERROR") {
					return nil
				}
			}
			select {
			case <-ctx.Done():
				// Surface the timeout/cancellation in the chromedp error (the
				// caller t.Fatalf's with it after dumping diagnostics);
				// otherwise the only signal would be a confusing
				// "#status = connecting" assertion failure.
				return fmt.Errorf("gave up after %v waiting for %q prefix on %s: %w", timeout, wantPrefix, sel, ctx.Err())
			case <-ticker.C:
			}
		}
	})
}

// syncBuffer is a goroutine-safe bytes buffer for capturing log output that
// the http server writes from its handler goroutines.
type syncBuffer struct {
	mu  sync.Mutex
	buf strings.Builder
}

func (b *syncBuffer) Write(p []byte) (int, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.buf.Write(p)
}

func (b *syncBuffer) String() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.buf.String()
}