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