Skip to content

Commit d217002

Browse files
benbenbenben
authored
fix: follow forwarded Trino pool result continuations (#1209)
1 parent e3faf93 commit d217002

4 files changed

Lines changed: 174 additions & 2 deletions

File tree

‎controlplane/trino_pool_continuation_test.go‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
//go:build kubernetes
2+
3+
package controlplane
4+
5+
import (
6+
"context"
7+
"encoding/json"
8+
"errors"
9+
"net/http"
10+
"net/http/httptest"
11+
"strings"
12+
"testing"
13+
14+
"github.com/posthog/duckgres/controlplane/provisioner"
15+
)
16+
17+
func TestPoolCandidateConsumesForwardedHTTPSContinuation(t *testing.T) {
18+
for _, explicitPort := range []bool{false, true} {
19+
name := "implicit HTTPS port"
20+
if explicitPort {
21+
name = "explicit HTTPS port"
22+
}
23+
t.Run(name, func(t *testing.T) {
24+
coordinator := newFakeCoordinator(t)
25+
pages := 0
26+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
27+
user, password, ok := r.BasicAuth()
28+
if !ok || user != "observer" || password != "test-password" || r.Header.Get("X-Forwarded-Proto") != "https" || r.Header.Get("X-Forwarded-Port") != "443" {
29+
t.Error("probe lost authentication or forwarded origin headers")
30+
w.WriteHeader(http.StatusUnauthorized)
31+
return
32+
}
33+
next := func(path string) {
34+
host := strings.Split(r.Host, ":")[0]
35+
if explicitPort {
36+
host += ":443"
37+
}
38+
_ = json.NewEncoder(w).Encode(map[string]any{"nextUri": "https://" + host + path})
39+
}
40+
switch r.URL.Path {
41+
case "/v1/statement":
42+
next("/v1/statement/queued/query/token/1")
43+
case "/v1/statement/queued/query/token/1":
44+
pages++
45+
next("/v1/statement/executing/query/token/2")
46+
case "/v1/statement/executing/query/token/2":
47+
pages++
48+
_ = json.NewEncoder(w).Encode(map[string]any{"data": coordinator.nodes})
49+
default:
50+
coordinator.server.Config.Handler.ServeHTTP(w, r)
51+
}
52+
}))
53+
defer server.Close()
54+
validation, err := validateTrinoPoolCandidate(context.Background(), server.Client(), server.URL,
55+
func() (string, string) { return "observer", "test-password" },
56+
trinoPoolObservation{ReadyWorkers: 2, CoordinatorImage: fakeCoordinatorImage},
57+
trinoPoolExpectation{Image: fakeCoordinatorImage, CatalogRevision: 42, InternalHTTP: true,
58+
ProjectionDigest: provisioner.TrinoProjectionDigest(fakePolicyRevision, fakePasswordRevision, fakeGroupRevision)})
59+
if err != nil {
60+
t.Fatalf("candidate could not consume its forwarded result pages: %v", err)
61+
}
62+
if pages != 2 || validation.ReadyWorkers != 2 {
63+
t.Fatalf("pages=%d, workers=%d, want two of each", pages, validation.ReadyWorkers)
64+
}
65+
})
66+
}
67+
}
68+
69+
func TestPoolProbeRejectsUntrustedContinuationOrigins(t *testing.T) {
70+
for _, target := range []string{
71+
"https://foreign.example.test/v1/statement/result",
72+
"https://%s:444/v1/statement/result",
73+
"http://%s:443/v1/statement/result",
74+
"ftp://%s/v1/statement/result",
75+
"//%s/v1/statement/result",
76+
"https://user:secret@%s/v1/statement/result",
77+
"https://%s/v1/statement/result?secret=value",
78+
"https://%s/v1/statement/result?",
79+
"https://%s/v1/statement/result#fragment",
80+
"https://%s/v1/statement-adjacent",
81+
"https://%s/v1/info",
82+
"https://%s/v1/statement",
83+
} {
84+
t.Run(target, func(t *testing.T) {
85+
requests := 0
86+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
87+
requests++
88+
host := strings.Split(r.Host, ":")[0]
89+
_ = json.NewEncoder(w).Encode(map[string]any{"nextUri": strings.ReplaceAll(target, "%s", host)})
90+
}))
91+
defer server.Close()
92+
client := rolloutSQLClient{baseURL: server.URL, client: server.Client(), username: "observer", password: "test-password", internalHTTP: true}
93+
_, err := client.statement(context.Background(), "SELECT 1")
94+
if err == nil || err.Error() != "invalid coordinator response endpoint" || requests != 1 {
95+
t.Fatalf("unsafe continuation accepted: requests=%d, err=%v", requests, err)
96+
}
97+
if strings.Contains(err.Error(), "secret") || strings.Contains(err.Error(), "test-password") || strings.Contains(err.Error(), "example.test") {
98+
t.Fatalf("failure exposed response-controlled or credential text: %v", err)
99+
}
100+
})
101+
}
102+
}
103+
104+
func TestPoolProbeDoesNotFollowHTTPRedirect(t *testing.T) {
105+
foreignRequests := 0
106+
foreign := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { foreignRequests++ }))
107+
defer foreign.Close()
108+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
109+
if r.URL.Path == "/v1/statement" {
110+
_ = json.NewEncoder(w).Encode(map[string]any{"nextUri": "https://" + strings.Split(r.Host, ":")[0] + "/v1/statement/result"})
111+
return
112+
}
113+
http.Redirect(w, r, foreign.URL+"/v1/statement/result", http.StatusTemporaryRedirect)
114+
}))
115+
defer server.Close()
116+
client := rolloutSQLClient{baseURL: server.URL, client: newRolloutHTTPClient(""), username: "observer", password: "test-password", internalHTTP: true}
117+
if _, err := client.statement(context.Background(), "SELECT 1"); err == nil || foreignRequests != 0 {
118+
t.Fatalf("redirect followed: calls=%d, error=%v", foreignRequests, err)
119+
}
120+
}
121+
122+
func TestLegacyProbeDoesNotAcceptForwardedHTTPSPort(t *testing.T) {
123+
requests := 0
124+
server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
125+
requests++
126+
_ = json.NewEncoder(w).Encode(map[string]any{"nextUri": "https://" + strings.Split(r.Host, ":")[0] + "/v1/statement/result"})
127+
}))
128+
defer server.Close()
129+
client := rolloutSQLClient{baseURL: server.URL, client: server.Client(), username: "observer", password: "test-password"}
130+
_, err := client.statement(context.Background(), "SELECT 1")
131+
if err == nil || err.Error() != "invalid coordinator response endpoint" || requests != 1 {
132+
t.Fatalf("legacy probe accepted a different HTTPS port: requests=%d, error=%v", requests, err)
133+
}
134+
}
135+
136+
func TestPoolNodeInventoryFailureKeepsSafeCause(t *testing.T) {
137+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
138+
_ = json.NewEncoder(w).Encode(map[string]any{"nextUri": "https://foreign.example.test/v1/statement/result"})
139+
}))
140+
defer server.Close()
141+
client := rolloutSQLClient{baseURL: server.URL, client: server.Client(), username: "observer", password: "test-password", internalHTTP: true}
142+
_, err := registeredWorkerCount(context.Background(), client, "node-1")
143+
if !errors.Is(err, errTrinoPoolCandidateNotReady) || !strings.Contains(err.Error(), "invalid coordinator response endpoint") {
144+
t.Fatalf("inventory failure lost its safe underlying cause: %v", err)
145+
}
146+
}

‎controlplane/trino_pool_validate.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -427,7 +427,7 @@ func fetchCatalogSync(ctx context.Context, client *http.Client, coordinatorURL,
427427
func registeredWorkerCount(ctx context.Context, sql rolloutSQLClient, coordinatorNodeID string) (int, error) {
428428
rows, err := sql.statement(ctx, "SELECT node_id, coordinator, state FROM system.runtime.nodes")
429429
if err != nil {
430-
return 0, fmt.Errorf("%w: node inventory unavailable", errTrinoPoolCandidateNotReady)
430+
return 0, fmt.Errorf("%w: node inventory unavailable: %v", errTrinoPoolCandidateNotReady, err)
431431
}
432432
workers, coordinators := 0, 0
433433
for _, row := range rows {

‎controlplane/trino_rollout_probe.go‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,18 @@ func (c rolloutSQLClient) scheme() string {
142142
func (c rolloutSQLClient) read(ctx context.Context, method, endpoint, sql string) ([]byte, error) {
143143
base, baseErr := url.Parse(c.baseURL)
144144
parsed, err := url.Parse(endpoint)
145-
if baseErr != nil || err != nil || parsed.Scheme != c.scheme() || parsed.Scheme != base.Scheme || !strings.EqualFold(parsed.Hostname(), base.Hostname()) || rolloutHTTPSPort(parsed) != rolloutHTTPSPort(base) || parsed.User != nil || parsed.RawQuery != "" || parsed.ForceQuery || parsed.Fragment != "" || (parsed.Path != "/v1/statement" && !strings.HasPrefix(parsed.Path, "/v1/statement/") && parsed.Path != "/v1/info") {
145+
if baseErr != nil || err != nil || base.Scheme != c.scheme() || !strings.EqualFold(parsed.Hostname(), base.Hostname()) || parsed.User != nil || parsed.RawQuery != "" || parsed.ForceQuery || parsed.Fragment != "" || (parsed.Path != "/v1/statement" && !strings.HasPrefix(parsed.Path, "/v1/statement/") && parsed.Path != "/v1/info") {
146+
return nil, errors.New("invalid coordinator response endpoint")
147+
}
148+
internalContinuation := c.internalHTTP && method == http.MethodGet &&
149+
strings.HasPrefix(parsed.Path, "/v1/statement/") &&
150+
parsed.Scheme == forwardedScheme && rolloutHTTPSPort(parsed) == "443"
151+
if internalContinuation {
152+
// Trino returns the HTTPS origin declared by our forwarded headers.
153+
// Keep validated result paths on the configured internal coordinator transport.
154+
parsed.Scheme, parsed.Host = base.Scheme, base.Host
155+
endpoint = parsed.String()
156+
} else if parsed.Scheme != base.Scheme || rolloutHTTPSPort(parsed) != rolloutHTTPSPort(base) {
146157
return nil, errors.New("invalid coordinator response endpoint")
147158
}
148159
if c.username == "" || c.password == "" {

‎tests/mw-dev/README.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -899,6 +899,21 @@ username through `DUCKGRES_TRINO_MANAGED_GATEWAY_USERNAME`. The same active
899899
acceptance path requires successful API and capability authentication before
900900
the operator can create and admit compute instances.
901901

902+
Internal HTTP coordinator probes declare forwarded HTTPS on port 443. Trino
903+
uses that advertised origin in statement continuation URLs. Duckgres validates
904+
the same coordinator hostname, HTTPS port 443, and statement-result path before
905+
mapping those continuations back to the configured internal HTTP endpoint.
906+
Other origins, user information, query strings, fragments, and HTTP redirects
907+
remain rejected. Fixed HTTPS coordinator probes retain their exact-origin rule.
908+
No new configuration or credential is needed.
909+
910+
`TestPoolCandidateConsumesForwardedHTTPSContinuation` exercises full candidate
911+
validation with queued and executing result pages and both HTTPS port spellings.
912+
The fixture follows Trino's request-derived URI construction, including forwarded
913+
headers; returning all node rows from the initial POST would miss this failure.
914+
Run the active-pool acceptance stage after deployment to verify real coordinator
915+
admission and querying. This local regression does not replace that cluster check.
916+
902917
The default in-Job fixture does not configure a shared pool, and its harness runs
903918
only after the control plane starts. It cannot reproduce this startup failure by
904919
changing its own environment: the control plane reads these settings at process

0 commit comments

Comments
 (0)