From 7af4f1022664362ff5e031154229da45ed40c71d Mon Sep 17 00:00:00 2001 From: Hongming Wang Date: Tue, 14 Apr 2026 18:09:44 -0700 Subject: [PATCH] fix(middleware): tenant guard reads bare UUID from state= (no prefix) Pair to molecule-controlplane PR #8. Fly's proxy returns 502 if the fly-replay state value contains '=', so the control plane now puts the bare UUID in state= (no 'org-id=' prefix). TenantGuard now treats the whole 'state=...' value as the org id. --- platform/internal/middleware/tenant_guard.go | 24 ++++++++----------- .../internal/middleware/tenant_guard_test.go | 23 +++++++++--------- 2 files changed, 21 insertions(+), 26 deletions(-) diff --git a/platform/internal/middleware/tenant_guard.go b/platform/internal/middleware/tenant_guard.go index d59b37af..3bcd010d 100644 --- a/platform/internal/middleware/tenant_guard.go +++ b/platform/internal/middleware/tenant_guard.go @@ -10,12 +10,11 @@ import ( // flyReplaySrcHeader is the header Fly injects on requests it replays via // the `fly-replay: ...;state=...` mechanism. Format is a semicolon- // separated list of k=v pairs, e.g. -// instance=91854...;region=ord;t=1700000000000;state=org-id= -// We care only about the `state=` segment; the control plane encodes -// the org id as `state=org-id=` so we can treat it equivalently -// to the X-Molecule-Org-Id header. +// instance=91854...;region=ord;t=1700000000000;state= +// Control plane puts the bare UUID in state (no prefix) because Fly's +// proxy returns 502 "replay malformed" on any second `=` in the value. +// We read the whole state= segment as the org id. const flyReplaySrcHeader = "Fly-Replay-Src" -const flyReplayStatePrefix = "org-id=" // Tenant-mode guard — public repo's only SaaS hook. // @@ -88,9 +87,10 @@ func TenantGuardWithOrgID(configuredOrgID string) gin.HandlerFunc { } } -// orgIDFromReplaySrc extracts the org id the control plane encoded via -// `state=org-id=` in the fly-replay response header. Returns "" if -// the header is missing, malformed, or the state segment isn't ours. +// orgIDFromReplaySrc extracts the org id the control plane put in the +// fly-replay state= segment. Value is the bare UUID — the control plane +// deliberately doesn't prefix it because Fly 502s on any `=` in the state +// value. Returns "" if the header is missing or has no state segment. // Separated from TenantGuardWithOrgID so tests can round-trip header → // id without spinning a full Gin context. func orgIDFromReplaySrc(header string) string { @@ -100,12 +100,8 @@ func orgIDFromReplaySrc(header string) string { for _, seg := range strings.Split(header, ";") { seg = strings.TrimSpace(seg) const statePrefix = "state=" - if !strings.HasPrefix(seg, statePrefix) { - continue - } - value := seg[len(statePrefix):] - if strings.HasPrefix(value, flyReplayStatePrefix) { - return value[len(flyReplayStatePrefix):] + if strings.HasPrefix(seg, statePrefix) { + return seg[len(statePrefix):] } } return "" diff --git a/platform/internal/middleware/tenant_guard_test.go b/platform/internal/middleware/tenant_guard_test.go index 034e4dda..f82f75ad 100644 --- a/platform/internal/middleware/tenant_guard_test.go +++ b/platform/internal/middleware/tenant_guard_test.go @@ -82,9 +82,9 @@ func TestTenantGuard_AllowlistBypassesCheck(t *testing.T) { } } -// Fly-Replay-Src state path: the production path. Control plane sends the -// org id as `state=org-id=` via fly-replay; Fly injects that into -// the replayed request as a segment of the Fly-Replay-Src header. +// Fly-Replay-Src state path: the production path. Control plane puts the +// bare UUID in state= (no prefix — Fly 502s on `=` in the state value). +// Fly injects the whole Fly-Replay-Src header on the replayed request. func TestTenantGuard_AcceptsFlyReplaySrcState(t *testing.T) { gin.SetMode(gin.TestMode) r := gin.New() @@ -92,7 +92,7 @@ func TestTenantGuard_AcceptsFlyReplaySrcState(t *testing.T) { r.GET("/workspaces", func(c *gin.Context) { c.String(200, "ok") }) req := httptest.NewRequest("GET", "/workspaces", nil) - req.Header.Set("Fly-Replay-Src", "instance=src-123;region=ord;t=1700000000000;state=org-id=org-abc") + req.Header.Set("Fly-Replay-Src", "instance=src-123;region=ord;t=1700000000000;state=org-abc") w := httptest.NewRecorder() r.ServeHTTP(w, req) @@ -108,7 +108,7 @@ func TestTenantGuard_RejectsFlyReplaySrcMismatch(t *testing.T) { r.GET("/workspaces", func(c *gin.Context) { c.String(200, "ok") }) req := httptest.NewRequest("GET", "/workspaces", nil) - req.Header.Set("Fly-Replay-Src", "state=org-id=org-xyz") + req.Header.Set("Fly-Replay-Src", "state=org-xyz") w := httptest.NewRecorder() r.ServeHTTP(w, req) @@ -119,13 +119,12 @@ func TestTenantGuard_RejectsFlyReplaySrcMismatch(t *testing.T) { func TestOrgIDFromReplaySrc(t *testing.T) { cases := map[string]string{ - "instance=x;region=ord;state=org-id=abc-123": "abc-123", - "state=org-id=abc-123;instance=x": "abc-123", - " state=org-id=abc-123 ": "abc-123", - "state=other=foo;instance=x": "", // wrong state key - "instance=x;region=ord": "", // no state - "": "", // empty header - "garbage": "", // unparseable + "instance=x;region=ord;state=abc-123": "abc-123", + "state=abc-123;instance=x": "abc-123", + " state=abc-123 ": "abc-123", + "instance=x;region=ord": "", // no state + "": "", // empty header + "garbage": "", // unparseable } for in, want := range cases { if got := orgIDFromReplaySrc(in); got != want {