diff --git a/sqlx-postgres/src/connection/mod.rs b/sqlx-postgres/src/connection/mod.rs index d594585b6c..10ec336222 100644 --- a/sqlx-postgres/src/connection/mod.rs +++ b/sqlx-postgres/src/connection/mod.rs @@ -27,6 +27,7 @@ pub use self::stream::PgStream; mod describe; mod establish; mod executor; +mod oauth; mod resolve; mod sasl; mod stream; diff --git a/sqlx-postgres/src/connection/oauth.rs b/sqlx-postgres/src/connection/oauth.rs new file mode 100644 index 0000000000..201ff8cc4b --- /dev/null +++ b/sqlx-postgres/src/connection/oauth.rs @@ -0,0 +1,172 @@ +use crate::connection::stream::PgStream; +use crate::error::Error; +use crate::message::{OAuthBearerResponse, SaslInitialResponse, SaslResponse}; +use crate::options::PgOAuthToken; + +/// The only SASL mechanism PostgreSQL's `oauth` HBA method advertises. +/// +/// OAUTHBEARER defines no channel binding, so there is no `-PLUS` variant to negotiate. +pub(crate) const MECHANISM: &str = "OAUTHBEARER"; + +/// RFC 7628 key/value separator (`kvsep`). +const KVSEP: &str = "\x01"; + +/// RFC 5801 gs2-header. OAUTHBEARER has no channel binding and PostgreSQL rejects the `p` +/// specifier outright, so this is always `n` (client does not support channel binding) +/// followed by an empty authzid. +const GS2_HEADER: &str = "n,,"; + +const BEARER_SCHEME: &str = "Bearer "; + +/// Authenticate with a bearer token over SASL `OAUTHBEARER`. +/// +/// This is the "token-first" flow: the token travels in the SASL initial client response, so a +/// successful authentication costs no extra round trip. SQLx never contacts the identity +/// provider, so the discovery flow of RFC 7628 §3.2.2 is not implemented; if the server rejects +/// the token, the exchange is closed out and the server's error is returned. +pub(crate) async fn authenticate(stream: &mut PgStream, token: &PgOAuthToken) -> Result<(), Error> { + let token = token.fetch().await?; + let response = initial_client_response(&token)?; + + stream + .send(SaslInitialResponse { + mechanism: MECHANISM, + response: &response, + }) + .await?; + + match stream.recv_expect::().await? { + // The server validated the token. It sends no mechanism-specific final data, so + // `AuthenticationOk` arrives directly and the exchange is over. + OAuthBearerResponse::Ok => Ok(()), + + OAuthBearerResponse::Failure(document) => { + // RFC 7628 §3.2.3: the only response the server will accept now is a single + // kvsep. Sending it lets the server report the failure as a normal + // `ErrorResponse` instead of leaving the exchange hanging. + stream.send(SaslResponse(KVSEP)).await?; + + // `PgStream::recv` turns `ErrorResponse` into `Err`, which is the expected + // outcome here and carries the server's own diagnostic. + stream.recv().await?; + + // Reached only if the server said something else entirely. + Err(err_protocol!( + "OAUTHBEARER authentication failed; server returned: {}", + String::from_utf8_lossy(&document) + )) + } + } +} + +/// Build the initial client response: `n,,^Aauth=Bearer ^A^A`. +fn initial_client_response(token: &str) -> Result { + validate_token(token)?; + + Ok(format!( + "{GS2_HEADER}{KVSEP}auth={BEARER_SCHEME}{token}{KVSEP}{KVSEP}" + )) +} + +/// Check the token against the `b64token` grammar of RFC 6750 §2.1, which is what the server +/// itself enforces. +/// +/// This is not merely a nicety. The grammar excludes the kvsep byte, whitespace and NUL, so a +/// token that passes cannot forge additional key/value pairs or truncate the message. Checking +/// it here turns a corrupt token into a clear local error instead of a protocol violation from +/// the server. +/// +/// The token value is never included in the error. +fn validate_token(token: &str) -> Result<(), Error> { + // Tokens may end with any number of base64 padding characters. + let unpadded = token.trim_end_matches('='); + + if unpadded.is_empty() { + return Err(Error::Configuration( + "OAuth bearer token is empty".to_string().into(), + )); + } + + if !unpadded.bytes().all(is_b64token_byte) { + return Err(Error::Configuration( + "OAuth bearer token contains characters that are not allowed by the `b64token` \ + grammar of RFC 6750; the token value is omitted from this error" + .to_string() + .into(), + )); + } + + Ok(()) +} + +const fn is_b64token_byte(byte: u8) -> bool { + byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~' | b'+' | b'/') +} + +#[cfg(test)] +mod tests { + use super::initial_client_response; + + #[test] + fn initial_client_response_is_token_first() { + // The token travels in the initial response, so no challenge is needed first. + assert_eq!( + initial_client_response("abc123").unwrap(), + "n,,\x01auth=Bearer abc123\x01\x01" + ); + } + + #[test] + fn initial_client_response_declares_no_channel_binding() { + // PostgreSQL rejects the `p` specifier for OAUTHBEARER outright. + let response = initial_client_response("abc123").unwrap(); + + assert!(response.starts_with("n,,")); + assert!(!response.starts_with('p')); + } + + #[test] + fn padded_token_is_accepted() { + assert_eq!( + initial_client_response("dG9rZW4=").unwrap(), + "n,,\x01auth=Bearer dG9rZW4=\x01\x01" + ); + } + + #[test] + fn b64token_alphabet_is_accepted() { + initial_client_response("aZ09-._~+/").unwrap(); + } + + // The rejection cases go through `initial_client_response` rather than calling the + // check directly, so that dropping the check from the call site fails a test. + + #[test] + fn token_may_not_forge_a_key_value_pair() { + // Without this check the kvsep would let a token append its own kvpairs. + initial_client_response("abc\x01host=evil").unwrap_err(); + } + + #[test] + fn token_may_not_contain_nul_or_whitespace() { + // The server compares the message length against `strlen`, so a NUL is fatal. + initial_client_response("abc\0def").unwrap_err(); + initial_client_response("abc def").unwrap_err(); + initial_client_response("abc\ndef").unwrap_err(); + } + + #[test] + fn empty_token_is_rejected() { + initial_client_response("").unwrap_err(); + initial_client_response("==").unwrap_err(); + } + + #[test] + fn an_error_never_contains_the_token() { + const TOKEN: &str = "sensitive\x01value"; + + let error = initial_client_response(TOKEN).unwrap_err().to_string(); + + assert!(!error.contains("sensitive")); + } +} diff --git a/sqlx-postgres/src/connection/sasl.rs b/sqlx-postgres/src/connection/sasl.rs index 157e1214ab..2b9686c8d5 100644 --- a/sqlx-postgres/src/connection/sasl.rs +++ b/sqlx-postgres/src/connection/sasl.rs @@ -1,3 +1,4 @@ +use crate::connection::oauth; use crate::connection::stream::PgStream; use crate::error::Error; use crate::message::{Authentication, AuthenticationSasl, SaslInitialResponse, SaslResponse}; @@ -24,6 +25,7 @@ pub(crate) async fn authenticate( ) -> Result<(), Error> { let mut has_sasl = false; let mut has_sasl_plus = false; + let mut has_oauth = false; let mut unknown = Vec::new(); for mechanism in data.mechanisms() { @@ -36,12 +38,30 @@ pub(crate) async fn authenticate( has_sasl_plus = true; } + oauth::MECHANISM => { + has_oauth = true; + } + _ => { unknown.push(mechanism.to_owned()); } } } + // PostgreSQL's `oauth` HBA method (added in 18) advertises OAUTHBEARER and nothing else. + if has_oauth { + if let Some(token) = &options.oauth_token { + return oauth::authenticate(stream, token).await; + } + + if !has_sasl && !has_sasl_plus { + return Err(err_protocol!( + "server requested OAUTHBEARER authentication, but no OAuth token is \ + configured; see `PgConnectOptions::oauth_token_provider`" + )); + } + } + if !has_sasl_plus && !has_sasl { return Err(err_protocol!( "unsupported SASL authentication mechanisms: {}", @@ -74,8 +94,8 @@ pub(crate) async fn authenticate( stream .send(SaslInitialResponse { + mechanism: "SCRAM-SHA-256", response: &client_first_message, - plus: false, }) .await?; diff --git a/sqlx-postgres/src/lib.rs b/sqlx-postgres/src/lib.rs index f68c928881..771c6d1b04 100644 --- a/sqlx-postgres/src/lib.rs +++ b/sqlx-postgres/src/lib.rs @@ -56,7 +56,7 @@ pub use database::Postgres; pub use error::{PgDatabaseError, PgErrorPosition}; pub use listener::{PgListener, PgNotification}; pub use message::PgSeverity; -pub use options::{PgConnectOptions, PgSslMode}; +pub use options::{PgConnectOptions, PgOAuthToken, PgSslMode}; pub use query_result::PgQueryResult; pub use row::PgRow; pub use statement::PgStatement; diff --git a/sqlx-postgres/src/message/authentication.rs b/sqlx-postgres/src/message/authentication.rs index 3a3cf7ff6e..1a59d829ff 100644 --- a/sqlx-postgres/src/message/authentication.rs +++ b/sqlx-postgres/src/message/authentication.rs @@ -87,6 +87,40 @@ impl BackendMessage for Authentication { } } +/// The server's reply to an `OAUTHBEARER` initial client response. +/// +/// The OAUTHBEARER exchange does not reuse [`AuthenticationSaslContinue`], whose body is a set +/// of SCRAM attributes: on failure the server sends a JSON error document instead, and it must +/// be read verbatim. +#[derive(Debug)] +pub enum OAuthBearerResponse { + /// The token was accepted. OAUTHBEARER carries no server signature, so the server sends no + /// SASL final message and `AuthenticationOk` arrives directly. + Ok, + + /// The token was rejected. The body is the server's JSON status document, which names the + /// issuer and the required scope. + Failure(Bytes), +} + +impl BackendMessage for OAuthBearerResponse { + const FORMAT: BackendMessageFormat = BackendMessageFormat::Authentication; + + fn decode_body(mut buf: Bytes) -> Result { + Ok(match buf.get_u32() { + 0 => OAuthBearerResponse::Ok, + 11 => OAuthBearerResponse::Failure(buf), + + ty => { + return Err(err_protocol!( + "unexpected authentication message {} during OAUTHBEARER exchange", + ty + )); + } + }) + } +} + /// Body of [Authentication::Md5Password]. #[derive(Debug)] pub struct AuthenticationMd5Password { diff --git a/sqlx-postgres/src/message/mod.rs b/sqlx-postgres/src/message/mod.rs index dedfe7c1bb..ada7608774 100644 --- a/sqlx-postgres/src/message/mod.rs +++ b/sqlx-postgres/src/message/mod.rs @@ -30,7 +30,7 @@ mod startup; mod sync; mod terminate; -pub use authentication::{Authentication, AuthenticationSasl}; +pub use authentication::{Authentication, AuthenticationSasl, OAuthBearerResponse}; pub use backend_key_data::BackendKeyData; pub use bind::Bind; pub use close::Close; diff --git a/sqlx-postgres/src/message/sasl.rs b/sqlx-postgres/src/message/sasl.rs index 5593a9367a..b137566f7b 100644 --- a/sqlx-postgres/src/message/sasl.rs +++ b/sqlx-postgres/src/message/sasl.rs @@ -4,19 +4,9 @@ use sqlx_core::Error; use std::num::Saturating; pub struct SaslInitialResponse<'a> { + /// The name of the SASL mechanism the client selected. + pub mechanism: &'a str, pub response: &'a str, - pub plus: bool, -} - -impl SaslInitialResponse<'_> { - #[inline(always)] - fn selected_mechanism(&self) -> &'static str { - if self.plus { - "SCRAM-SHA-256-PLUS" - } else { - "SCRAM-SHA-256" - } - } } impl FrontendMessage for SaslInitialResponse<'_> { @@ -26,7 +16,7 @@ impl FrontendMessage for SaslInitialResponse<'_> { fn body_size_hint(&self) -> Saturating { let mut size = Saturating(0); - size += self.selected_mechanism().len(); + size += self.mechanism.len(); size += 1; // NUL terminator size += 4; // response_len @@ -37,7 +27,7 @@ impl FrontendMessage for SaslInitialResponse<'_> { fn encode_body(&self, buf: &mut Vec) -> Result<(), Error> { // name of the SASL authentication mechanism that the client selected - buf.put_str_nul(self.selected_mechanism()); + buf.put_str_nul(self.mechanism); let response_len = i32::try_from(self.response.len()).map_err(|_| { err_protocol!( diff --git a/sqlx-postgres/src/options/mod.rs b/sqlx-postgres/src/options/mod.rs index 1432673720..1826d10cea 100644 --- a/sqlx-postgres/src/options/mod.rs +++ b/sqlx-postgres/src/options/mod.rs @@ -1,15 +1,19 @@ use std::borrow::Cow; use std::env::var; use std::fmt::{self, Display, Write}; +use std::future::Future; use std::path::{Path, PathBuf}; use std::sync::{Arc, OnceLock}; +pub use oauth::PgOAuthToken; +use sqlx_core::error::BoxDynError; use sqlx_core::net::tls::TlsConnector; pub use ssl_mode::PgSslMode; use crate::{connection::LogSettings, net::tls::CertificateInput}; mod connect; +mod oauth; mod parse; mod pgpass; mod ssl_mode; @@ -29,6 +33,7 @@ pub struct PgConnectOptions { pub(crate) log_settings: LogSettings, pub(crate) extra_float_digits: Option>, pub(crate) options: Option, + pub(crate) oauth_token: Option, } impl Default for PgConnectOptions { @@ -108,6 +113,7 @@ impl PgConnectOptions { extra_float_digits: Some("2".into()), log_settings: Default::default(), options: var("PGOPTIONS").ok(), + oauth_token: None, } } @@ -124,6 +130,62 @@ impl PgConnectOptions { self } + /// Authenticate with a fixed OAuth 2.0 bearer token, for PostgreSQL's `oauth` + /// authentication method. + /// + /// Prefer [`oauth_token_provider`][Self::oauth_token_provider] where the token can expire: + /// a token set here is used for every connection attempt, including reconnections made by a + /// pool long after the token was minted. + /// + /// The token is used only when the server requests SASL `OAUTHBEARER`; it is never sent to a + /// server asking for a password. + /// + /// # Example + /// + /// ```rust + /// # use sqlx_postgres::PgConnectOptions; + /// # fn f(token: String) -> PgConnectOptions { + /// PgConnectOptions::new().oauth_token(token) + /// # } + /// ``` + pub fn oauth_token(mut self, token: impl Into) -> Self { + self.oauth_token = Some(PgOAuthToken::new(token)); + self + } + + /// Obtain an OAuth 2.0 bearer token once per connection attempt, for PostgreSQL's `oauth` + /// authentication method. + /// + /// SQLx does not talk to an identity provider. `provider` is called immediately before the + /// SASL exchange and should return a currently valid token; how it is acquired, cached and + /// refreshed is the application's choice. + /// + /// The token is used only when the server requests SASL `OAUTHBEARER`; it is never sent to a + /// server asking for a password. + /// + /// Note that there is deliberately no connection-string equivalent of this option: a URL + /// ends up in shell history, in `ps` output and in logs, which is no place for a bearer + /// token. + /// + /// # Example + /// + /// ```rust + /// # use sqlx_postgres::PgConnectOptions; + /// let options = PgConnectOptions::new().oauth_token_provider(|| async { + /// // Called for each connection attempt, so an expired token can be refreshed. + /// let token = std::env::var("EXAMPLE_OAUTH_TOKEN")?; + /// Ok(token) + /// }); + /// ``` + pub fn oauth_token_provider(mut self, provider: F) -> Self + where + F: Fn() -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + self.oauth_token = Some(PgOAuthToken::from_provider(provider)); + self + } + /// Sets the name of the host to connect to. /// /// If a host name begins with a slash, it specifies diff --git a/sqlx-postgres/src/options/oauth.rs b/sqlx-postgres/src/options/oauth.rs new file mode 100644 index 0000000000..65c81cbfcb --- /dev/null +++ b/sqlx-postgres/src/options/oauth.rs @@ -0,0 +1,93 @@ +use std::fmt::{self, Debug, Formatter}; +use std::future::Future; +use std::sync::Arc; + +use futures_core::future::BoxFuture; +use sqlx_core::error::BoxDynError; + +use crate::error::Error; + +type Provider = dyn Fn() -> BoxFuture<'static, Result> + Send + Sync; + +/// A source of OAuth 2.0 bearer tokens for PostgreSQL's `oauth` authentication method. +/// +/// PostgreSQL 18 added the `oauth` HBA method, which authenticates a connection with an +/// OAuth 2.0 bearer token over the SASL `OAUTHBEARER` mechanism (RFC 7628) instead of a +/// password. +/// +/// SQLx does not talk to an identity provider: obtaining a token is the application's job. +/// This type wraps whatever the application already uses to get one. +/// +/// Because tokens expire, the token is requested once per connection attempt rather than +/// stored, so a pool that reconnects hours later presents a fresh token. Use +/// [`PgConnectOptions::oauth_token_provider`] for that. A single token that outlives the +/// connections made with it can be set with [`PgConnectOptions::oauth_token`]. +/// +/// [`PgConnectOptions::oauth_token_provider`]: crate::PgConnectOptions::oauth_token_provider +/// [`PgConnectOptions::oauth_token`]: crate::PgConnectOptions::oauth_token +#[derive(Clone)] +pub struct PgOAuthToken { + provider: Arc, +} + +impl PgOAuthToken { + /// Use a single, fixed token for every connection attempt. + pub fn new(token: impl Into) -> Self { + let token = token.into(); + + Self::from_provider(move || { + let token = token.clone(); + async move { Ok(token) } + }) + } + + /// Call `provider` once per connection attempt to obtain a token. + pub fn from_provider(provider: F) -> Self + where + F: Fn() -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + Self { + provider: Arc::new(move || Box::pin(provider())), + } + } + + pub(crate) async fn fetch(&self) -> Result { + (self.provider)().await.map_err(Error::Configuration) + } +} + +/// Deliberately opaque: `PgConnectOptions` derives `Debug`, and a bearer token is a +/// credential that must not reach a log, a panic message or an error. +impl Debug for PgOAuthToken { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.write_str("PgOAuthToken(..)") + } +} + +#[cfg(test)] +mod tests { + use crate::PgConnectOptions; + + #[test] + fn debug_does_not_leak_the_token() { + // `PgConnectOptions` derives `Debug`, so the token must be opaque at every depth + // rather than merely absent from a hand-written summary. + const TOKEN: &str = "super-secret-bearer-token"; + + let options = PgConnectOptions::new_without_pgpass().oauth_token(TOKEN); + + assert!(!format!("{:?}", options).contains(TOKEN)); + assert!(!format!("{:#?}", options).contains(TOKEN)); + } + + #[test] + fn debug_does_not_leak_a_token_captured_by_a_provider() { + const TOKEN: &str = "super-secret-bearer-token"; + + let options = PgConnectOptions::new_without_pgpass() + .oauth_token_provider(|| async { Ok(TOKEN.to_string()) }); + + assert!(!format!("{:#?}", options).contains(TOKEN)); + } +}