From d770fa656c22819aad9be127a4d983d3c7c8b468 Mon Sep 17 00:00:00 2001 From: Sebastian Bernauer Date: Thu, 30 Jul 2026 17:19:20 +0200 Subject: [PATCH 1/2] fix: Point forwarding headers at the external Trino endpoint Trino builds the absolute URLs it hands out to clients (OAuth 2.0 challenge, infoUri, ackUri) from the forwarding headers. These pointed at trino-lb, so clients got sent to /oauth2/token/{id} on trino-lb, which answers with a 404 and breaks the whole authentication flow. --- CHANGELOG.md | 9 + trino-lb-core/src/config.rs | 13 +- trino-lb/src/cluster_group_manager.rs | 245 ++++++++++++++++++++++- trino-lb/src/http_server/v1/statement.rs | 2 + 4 files changed, 263 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f21f64e..805ea76 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,9 +13,18 @@ All notable changes to this project will be documented in this file. - Handle Redis connection errors (e.g. broken pipe during master failover) gracefully instead of panicking. Previously, `get_queued_query_count` would `.unwrap()` on the Redis result, causing a panic that poisoned the metrics `RwLock`, cascading into further panics and leaving pods unresponsive ([#111]). +- Point the forwarding related headers (`Host`, `X-Forwarded-Host`, `X-Forwarded-Proto` and + `X-Forwarded-Port`) of proxied client requests at the configured `externalEndpoint` of the Trino + cluster. Trino builds the absolute URLs it hands out to clients from these headers, so previously + they pointed at trino-lb. Most notably this broke the OAuth 2.0 flow: Trino put the trino-lb + address into the `x_redirect_server` and `x_token_server` of its `WWW-Authenticate` challenge, so + clients were sent to `https:///oauth2/token/{id}`, which trino-lb answered with a `404`. + The RFC 7239 `Forwarded` header is now removed as well, as it takes precedence over the + `X-Forwarded-*` headers. This is a no-op for Trino clusters without an `externalEndpoint` ([#119]). [#111]: https://github.com/stackabletech/trino-lb/pull/111 [#116]: https://github.com/stackabletech/trino-lb/pull/116 +[#119]: https://github.com/stackabletech/trino-lb/pull/119 ## [0.6.0] - 2026-02-17 diff --git a/trino-lb-core/src/config.rs b/trino-lb-core/src/config.rs index 71745ca..709a7d0 100644 --- a/trino-lb-core/src/config.rs +++ b/trino-lb-core/src/config.rs @@ -192,8 +192,17 @@ pub struct TrinoClusterConfig { pub name: String, pub endpoint: Url, - /// Public endpoint of the Trino cluster. - /// This can e.g. be used to change segment ackUris to. + /// Public endpoint of the Trino cluster, meaning the address clients can reach this Trino + /// cluster at directly (bypassing trino-lb). + /// + /// It is used for two purposes: + /// + /// 1. The forwarding related headers (such as `Host` and `X-Forwarded-Host`) of proxied client + /// requests are pointed at this endpoint. Trino builds the absolute URLs it hands out to + /// clients from these headers, so without this they would point at trino-lb, which e.g. + /// breaks the OAuth 2.0 authentication flow. + /// 2. The segment ackUris are changed to this endpoint, as sometimes Trino gets confused and + /// puts the wrong endpoint (namely the one of trino-lb) in there. pub external_endpoint: Option, pub credentials: TrinoClusterCredentialsConfig, } diff --git a/trino-lb/src/cluster_group_manager.rs b/trino-lb/src/cluster_group_manager.rs index 49f78ea..503ec58 100644 --- a/trino-lb/src/cluster_group_manager.rs +++ b/trino-lb/src/cluster_group_manager.rs @@ -6,7 +6,10 @@ use std::{ use axum::{Json, body::Body, response::IntoResponse}; use futures::future::try_join_all; -use http::{HeaderMap, StatusCode}; +use http::{ + HeaderMap, HeaderName, HeaderValue, StatusCode, + header::{FORWARDED, HOST}, +}; use reqwest::Client; use serde::Serialize; use snafu::{OptionExt, ResultExt, Snafu}; @@ -76,8 +79,31 @@ pub enum Error { source: trino_lb_persistence::Error, cluster_group: String, }, + + #[snafu(display( + "Failed to determine the host of the external Trino endpoint {external_endpoint}" + ))] + ExternalTrinoEndpointWithoutHost { external_endpoint: Url }, + + #[snafu(display( + "Failed to turn {value:?} of the external Trino endpoint {external_endpoint} into a HTTP header value" + ))] + ConvertExternalTrinoEndpointToHeaderValue { + source: http::header::InvalidHeaderValue, + value: String, + external_endpoint: Url, + }, } +/// Not part of [`http::header`], so it needs to be defined here. +const X_FORWARDED_HOST: HeaderName = HeaderName::from_static("x-forwarded-host"); + +/// Not part of [`http::header`], so it needs to be defined here. +const X_FORWARDED_PORT: HeaderName = HeaderName::from_static("x-forwarded-port"); + +/// Not part of [`http::header`], so it needs to be defined here. +const X_FORWARDED_PROTO: HeaderName = HeaderName::from_static("x-forwarded-proto"); + pub struct ClusterGroupManager { groups: HashMap>, persistence: Arc, @@ -174,9 +200,11 @@ impl ClusterGroupManager { pub async fn send_query_to_cluster( &self, query: String, - headers: http::HeaderMap, + mut headers: http::HeaderMap, cluster: &TrinoCluster, ) -> Result { + point_forwarded_headers_to_trino(&mut headers, cluster.external_endpoint.as_ref())?; + // TODO: Enable propagation again. This is disabled, as the POST /v1/statement span runs for the whole // query lifetime and let it look like the initial POST takes multiple minutes. // add_current_context_to_client_request(tracing::Span::current().context(), &mut r_headers); @@ -221,14 +249,16 @@ impl ClusterGroupManager { } #[instrument( - skip(self), + skip(self, external_endpoint), fields(next_uri = %next_uri, headers = ?headers.sanitize()) )] pub async fn ask_for_query_state( &self, next_uri: Url, + external_endpoint: Option<&Url>, mut headers: HeaderMap, ) -> Result<(TrinoQueryApiResponse, HeaderMap), Error> { + point_forwarded_headers_to_trino(&mut headers, external_endpoint)?; add_current_context_to_client_request(tracing::Span::current().context(), &mut headers); let response = self .http_client @@ -255,14 +285,16 @@ impl ClusterGroupManager { /// Sometimes the trino-client HEADs a /executing/xxx endpoint instead of GETing it. /// We need to proxy this as a HEAD request as well. #[instrument( - skip(self), + skip(self, external_endpoint), fields(head_uri = %head_uri, headers = ?headers.sanitize()) )] pub async fn send_head_to_trino( &self, head_uri: Url, + external_endpoint: Option<&Url>, mut headers: HeaderMap, ) -> Result { + point_forwarded_headers_to_trino(&mut headers, external_endpoint)?; add_current_context_to_client_request(tracing::Span::current().context(), &mut headers); let response = self @@ -289,6 +321,10 @@ impl ClusterGroupManager { query: &TrinoQuery, requested_path: &str, ) -> Result<(), Error> { + point_forwarded_headers_to_trino( + &mut request_headers, + query.trino_external_endpoint.as_ref(), + )?; add_current_context_to_client_request( tracing::Span::current().context(), &mut request_headers, @@ -401,6 +437,93 @@ impl ClusterGroupManager { } } +/// Trino builds the absolute URLs it hands out to clients from the host and protocol of the incoming +/// request. Among others this affects +/// +/// 1. the `x_redirect_server` and `x_token_server` of the OAuth 2.0 `WWW-Authenticate` challenge, +/// 2. the `infoUri` and `partialCancelUri` of a query and +/// 3. the `ackUri` of spooled segments. +/// +/// As trino-lb proxies the requests of the clients, all of the relevant headers point at trino-lb, +/// which makes Trino hand out URLs that don't exist on trino-lb. The OAuth 2.0 endpoints are the +/// worst offender here: Clients get sent to `https:///oauth2/token/{id}`, which trino-lb +/// answers with a `404`, breaking the entire authentication flow. +/// +/// To prevent this, the forwarding related headers are overwritten with the `externalEndpoint` of +/// the Trino cluster, so that Trino hands out URLs pointing at itself, which clients can actually +/// reach. The `nextUri` is unaffected by this, as it is always rewritten to point back at trino-lb +/// afterwards (see [`TrinoQueryApiResponse::update_trino_references`]). +/// +/// `X-Forwarded-For` and `X-Real-Ip` are intentionally left alone, as they describe the client +/// itself and not the address the client asked for. +/// +/// This is a no-op in case no `externalEndpoint` is configured for the Trino cluster, as trino-lb +/// has no clue which address clients can reach Trino at in that case. +#[instrument(skip_all)] +fn point_forwarded_headers_to_trino( + headers: &mut HeaderMap, + external_endpoint: Option<&Url>, +) -> Result<(), Error> { + let Some(external_endpoint) = external_endpoint else { + return Ok(()); + }; + + let host = external_endpoint + .host_str() + .context(ExternalTrinoEndpointWithoutHostSnafu { + external_endpoint: external_endpoint.clone(), + })?; + + // The port is only added in case it is not the default one of the scheme, so that Trino hands + // out `https://trino.example.com/...` instead of `https://trino.example.com:443/...`. + let host = match external_endpoint.port() { + Some(port) => format!("{host}:{port}"), + None => host.to_owned(), + }; + let host = to_header_value(&host, external_endpoint)?; + + headers.insert(HOST, host.clone()); + headers.insert(X_FORWARDED_HOST, host); + headers.insert( + X_FORWARDED_PROTO, + to_header_value(external_endpoint.scheme(), external_endpoint)?, + ); + + match external_endpoint.port() { + Some(port) => { + headers.insert( + X_FORWARDED_PORT, + to_header_value(&port.to_string(), external_endpoint)?, + ); + } + // Trino derives the port from the protocol in this case, which is better than leaving the + // port of an upstream proxy (pointing at trino-lb) in place. + None => { + headers.remove(X_FORWARDED_PORT); + } + } + + // The `Forwarded` header (RFC 7239) takes precedence over the `X-Forwarded-*` headers in Jetty, + // which Trino is built upon. It therefore needs to be removed, as it would otherwise override + // everything set above. Every proxy we are aware of sets `X-Forwarded-For` alongside + // `Forwarded`, so the address of the client is not lost by doing so. + headers.remove(FORWARDED); + + debug!( + headers = ?headers.sanitize(), + "Pointed the forwarding related headers at the external Trino endpoint" + ); + + Ok(()) +} + +fn to_header_value(value: &str, external_endpoint: &Url) -> Result { + HeaderValue::from_str(value).context(ConvertExternalTrinoEndpointToHeaderValueSnafu { + value, + external_endpoint: external_endpoint.clone(), + }) +} + fn filter_to_trino_headers(headers: &HeaderMap) -> HeaderMap { let mut trino_headers = HeaderMap::new(); for (name, value) in headers.into_iter() { @@ -422,3 +545,117 @@ fn filter_to_www_authenticate_headers(headers: &HeaderMap) -> HeaderMap { www_headers } + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + /// The headers a Gateway or Ingress in front of trino-lb typically sets, all of them pointing at + /// trino-lb instead of Trino. + fn headers_pointing_at_trino_lb() -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert(HOST, HeaderValue::from_static("trino-lb.example.com")); + headers.insert( + FORWARDED, + HeaderValue::from_static("host=trino-lb.example.com;proto=https"), + ); + headers.insert( + X_FORWARDED_HOST, + HeaderValue::from_static("trino-lb.example.com"), + ); + headers.insert(X_FORWARDED_PORT, HeaderValue::from_static("443")); + headers.insert(X_FORWARDED_PROTO, HeaderValue::from_static("https")); + headers.insert("x-forwarded-for", HeaderValue::from_static("172.30.10.61")); + headers.insert("x-real-ip", HeaderValue::from_static("172.30.10.61")); + headers + } + + #[rstest] + #[case("https://trino.example.com", "trino.example.com", "https", None)] + #[case( + "https://trino.example.com:8443", + "trino.example.com:8443", + "https", + Some("8443") + )] + #[case( + "http://trino.example.com:8080", + "trino.example.com:8080", + "http", + Some("8080") + )] + // Default ports of the scheme are not added, so that Trino hands out the nicer looking + // `https://trino.example.com/...` instead of `https://trino.example.com:443/...`. + #[case("https://trino.example.com:443", "trino.example.com", "https", None)] + #[case("http://trino.example.com:80", "trino.example.com", "http", None)] + #[case( + "https://5.250.182.203:8443", + "5.250.182.203:8443", + "https", + Some("8443") + )] + fn point_forwarded_headers_to_external_trino_endpoint( + #[case] external_endpoint: &str, + #[case] expected_host: &str, + #[case] expected_proto: &str, + #[case] expected_port: Option<&str>, + ) { + let external_endpoint: Url = external_endpoint + .parse() + .expect("test case URL is always valid"); + let mut headers = headers_pointing_at_trino_lb(); + + point_forwarded_headers_to_trino(&mut headers, Some(&external_endpoint)) + .expect("headers built from a valid URL are always valid"); + + assert_eq!( + headers.get(HOST).expect("host is always set"), + expected_host + ); + assert_eq!( + headers + .get(X_FORWARDED_HOST) + .expect("x-forwarded-host is always set"), + expected_host + ); + assert_eq!( + headers + .get(X_FORWARDED_PROTO) + .expect("x-forwarded-proto is always set"), + expected_proto + ); + assert_eq!( + headers + .get(X_FORWARDED_PORT) + .map(|port| port.to_str().expect("the port is always valid ASCII")), + expected_port + ); + + // The `Forwarded` header takes precedence over the `X-Forwarded-*` headers, so it must be gone. + assert_eq!(headers.get(FORWARDED), None); + + // The client address must survive, as Trino uses it for e.g. auditing. + assert_eq!( + headers.get("x-forwarded-for").expect("must be kept"), + "172.30.10.61" + ); + assert_eq!( + headers.get("x-real-ip").expect("must be kept"), + "172.30.10.61" + ); + } + + /// Without an `externalEndpoint` trino-lb has no clue which address clients can reach Trino at, + /// so the headers need to be passed through unchanged. + #[test] + fn keep_forwarded_headers_without_external_trino_endpoint() { + let mut headers = headers_pointing_at_trino_lb(); + + point_forwarded_headers_to_trino(&mut headers, None) + .expect("not touching any headers never fails"); + + assert_eq!(headers, headers_pointing_at_trino_lb()); + } +} diff --git a/trino-lb/src/http_server/v1/statement.rs b/trino-lb/src/http_server/v1/statement.rs index 7ee8810..34da036 100644 --- a/trino-lb/src/http_server/v1/statement.rs +++ b/trino-lb/src/http_server/v1/statement.rs @@ -501,6 +501,7 @@ async fn handle_query_running_on_trino( trino_endpoint: query.trino_endpoint.clone(), }, )?, + query.trino_external_endpoint.as_ref(), headers, ) .await @@ -566,6 +567,7 @@ async fn handle_head_request_to_trino( trino_endpoint: query.trino_endpoint.clone(), }, )?, + query.trino_external_endpoint.as_ref(), headers, ) .await From 23f048c62505a98cfae741982646bedfad3891de Mon Sep 17 00:00:00 2001 From: Sebastian Bernauer Date: Mon, 3 Aug 2026 13:06:54 +0200 Subject: [PATCH 2/2] docs: Mention path is ignored --- trino-lb-core/src/config.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/trino-lb-core/src/config.rs b/trino-lb-core/src/config.rs index 709a7d0..97040b6 100644 --- a/trino-lb-core/src/config.rs +++ b/trino-lb-core/src/config.rs @@ -203,6 +203,8 @@ pub struct TrinoClusterConfig { /// breaks the OAuth 2.0 authentication flow. /// 2. The segment ackUris are changed to this endpoint, as sometimes Trino gets confused and /// puts the wrong endpoint (namely the one of trino-lb) in there. + /// + /// Note that the path of the URL is ignored. pub external_endpoint: Option, pub credentials: TrinoClusterCredentialsConfig, }