amadan

carlos / platform Public

edge: never wrap Connection: Upgrade requests in recordingResponseWriter

Fixes a regression from 22c2f62: recordingResponseWriter unconditionally
wrapped the ResponseWriter on hibernating routes, but it didn't implement
http.Hijacker. The stdlib reverse proxy's WebSocket-upgrade path needs to
hijack the raw connection, and an open WebSocket is also what keeps a
hibernating instance's idle clock from firing (internal/activator's
inflight counting) — so this would have 500'd every upgrade to a
hibernating route once caching was enabled.

Fix is two-layered per review: ServeHTTP now detects Connection: Upgrade
requests up front and routes them straight to the activator/proxy with
the real ResponseWriter, skipping recordingResponseWriter entirely (byte-
for-byte pre-cache behavior). recordingResponseWriter also now forwards
http.Hijacker and http.Flusher to its underlying ResponseWriter, closing
the same latent gap for any other caller and for streaming (SSE)
responses that do get recorded.
Paul Campbell 9ead6e1d59c80d21919a5f5c0014d84d9c1ce05f parent 22c2f62
2 files changed, +115 −0
  • internal/edge/edge.go +59 −0
  • internal/edge/edge_test.go +56 −0
diff --git a/internal/edge/edge.go b/internal/edge/edge.go
index 7d5479529..e382c0fff 100644
--- a/internal/edge/edge.go
+++ b/internal/edge/edge.go
@@ -16,6 +16,7 @@
package edge
import (
+ "bufio"
"context"
"crypto/ed25519"
"errors"
@@ -255,6 +256,26 @@ func (e *Edge) ServeHTTP(w http.ResponseWriter, r *http.Request) {
return
}
+ // A Connection: Upgrade request (a WebSocket handshake, typically GET)
+ // must reach the activator/proxy with the real ResponseWriter, completely
+ // untouched by caching: the stdlib reverse proxy's upgrade path type-
+ // asserts the ResponseWriter to http.Hijacker to take over the raw
+ // connection, and once upgraded, the open WebSocket is also what keeps a
+ // hibernating instance's idle clock from firing (the activator's inflight
+ // counting — see internal/activator/activator.go's actEntry.inflight doc
+ // comment). So this path skips caching entirely and stays byte-for-byte
+ // identical to pre-cache behavior, rather than relying on
+ // recordingResponseWriter's Hijack forwarding below.
+ if isUpgradeRequest(r) {
+ if up.hibernating {
+ if e.act.ServeHTTP(w, r, up.route, up.proxy) {
+ return
+ }
+ }
+ up.proxy.ServeHTTP(w, r)
+ return
+ }
+
// Cache check happens before the hibernation branch below — a hit must
// never reach the activator, since avoiding pointless wakes for
// bot/crawler/repeat traffic is the entire point of this cache.
@@ -294,6 +315,19 @@ func (e *Edge) ServeHTTP(w http.ResponseWriter, r *http.Request) {
up.proxy.ServeHTTP(w, r)
}
+// isUpgradeRequest reports whether r is a protocol-upgrade request (e.g. a
+// WebSocket handshake) per the Connection header's token list (RFC 7230
+// §6.7: comma-separated, case-insensitive — "Upgrade", "keep-alive, Upgrade",
+// etc. all qualify).
+func isUpgradeRequest(r *http.Request) bool {
+ for _, token := range strings.Split(r.Header.Get("Connection"), ",") {
+ if strings.EqualFold(strings.TrimSpace(token), "Upgrade") {
+ return true
+ }
+ }
+ return false
+}
+
// cacheMaxBody reports the body-size ceiling recordingResponseWriter should
// enforce, or 0 when caching is disabled (recordingResponseWriter is never
// constructed with cacheable=true in that case, so the value is unused).
@@ -422,6 +456,31 @@ func (r *recordingResponseWriter) Write(p []byte) (int, error) {
return r.ResponseWriter.Write(p)
}
+// Hijack forwards to the underlying ResponseWriter's http.Hijacker, when it
+// implements one — the same optional-interface pattern the stdlib itself
+// uses. ServeHTTP special-cases upgrade requests before ever constructing a
+// recordingResponseWriter (see isUpgradeRequest), so this method is not
+// expected to be called on the hot path today; it exists so the wrapper type
+// is never itself the reason a hijack fails, matching http.ResponseWriter
+// wrapper conventions generally.
+func (r *recordingResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+ hj, ok := r.ResponseWriter.(http.Hijacker)
+ if !ok {
+ return nil, nil, fmt.Errorf("edge: underlying ResponseWriter does not support http.Hijacker")
+ }
+ return hj.Hijack()
+}
+
+// Flush forwards to the underlying ResponseWriter's http.Flusher, when it
+// implements one — needed so a recorded (non-upgrade) response can still
+// stream incrementally to the real client, e.g. Server-Sent Events, instead
+// of buffering until the handler returns.
+func (r *recordingResponseWriter) Flush() {
+ if f, ok := r.ResponseWriter.(http.Flusher); ok {
+ f.Flush()
+ }
+}
+
// unknownHost answers a request for a host we don't serve: a redirect to the
// fallback (302 — a typo'd host must not be cached as permanently moved), or
// the classic 404 when no fallback is configured.
diff --git a/internal/edge/edge_test.go b/internal/edge/edge_test.go
index f82339d90..591e6379f 100644
--- a/internal/edge/edge_test.go
+++ b/internal/edge/edge_test.go
@@ -1,6 +1,7 @@
package edge
import (
+ "bufio"
"context"
"errors"
"io"
@@ -635,3 +636,58 @@ func TestServeHTTPRefusesVaryCookie(t *testing.T) {
t.Fatal("a response Vary-ing on Cookie must never be cached — it would defeat per-visitor correctness")
}
}
+
+// hijackableRecorder embeds httptest.ResponseRecorder (which deliberately
+// does NOT implement http.Hijacker) and adds a no-op Hijack, standing in for
+// a real connection's ResponseWriter — the minimum shape needed to prove
+// ServeHTTP passes an Upgrade request's ResponseWriter through untouched
+// instead of wrapping it in recordingResponseWriter (which pre-fix had no
+// Hijack, breaking the stdlib reverse proxy's WebSocket-upgrade path).
+type hijackableRecorder struct {
+ *httptest.ResponseRecorder
+}
+
+func (h *hijackableRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+ return nil, nil, nil
+}
+
+// hijackCheckingActivator fails the test if the ResponseWriter it's handed
+// doesn't satisfy http.Hijacker — proving ServeHTTP special-cased the
+// Upgrade request instead of wrapping it in recordingResponseWriter.
+type hijackCheckingActivator struct{ t *testing.T }
+
+func (h hijackCheckingActivator) ServeHTTP(w http.ResponseWriter, r *http.Request, rt registry.Route, proxy http.Handler) bool {
+ if _, ok := w.(http.Hijacker); !ok {
+ h.t.Fatal("ResponseWriter passed to activator for an Upgrade request does not implement http.Hijacker — recordingResponseWriter must not wrap upgrade requests")
+ }
+ return true
+}
+
+func (h hijackCheckingActivator) RetireExcept(keep map[string]bool) {}
+
+// TestServeHTTPUpgradeRequestBypassesRecording proves a Connection: Upgrade
+// request (a WebSocket handshake) to a hibernating route reaches the
+// activator with a still-Hijacker-capable ResponseWriter — i.e. ServeHTTP
+// never wraps it in recordingResponseWriter, which would break the stdlib
+// reverse proxy's hijack-based upgrade path and, with it, the open WebSocket
+// the activator relies on to keep the instance's idle clock from firing.
+func TestServeHTTPUpgradeRequestBypassesRecording(t *testing.T) {
+ reg := openReg(t)
+ // Addr is a bogus, never-dialed address — this test only needs to reach
+ // the activator, never the real backend.
+ if err := reg.Upsert(registry.Route{Host: "ws.example.com", Kind: "instance", Addr: "127.0.0.1:1", Status: "ready", Hibernate: true}); err != nil {
+ t.Fatalf("upsert: %v", err)
+ }
+ e, err := New(reg, hijackCheckingActivator{t: t})
+ if err != nil {
+ t.Fatalf("New: %v", err)
+ }
+ e.EnableCache(NewCache(DefaultCacheMaxEntries, DefaultCacheMaxBodyBytes))
+
+ req := httptest.NewRequest("GET", "http://ws.example.com/socket", nil)
+ req.Header.Set("Connection", "Upgrade")
+ req.Header.Set("Upgrade", "websocket")
+ rec := &hijackableRecorder{ResponseRecorder: httptest.NewRecorder()}
+
+ e.ServeHTTP(rec, req)
+}