feat(postgres): support SASL OAUTHBEARER authentication

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;

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::<OAuthBearerResponse>().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 <token>^A^A`.
+fn initial_client_response(token: &str) -> Result<String, Error> {
+    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"));
+    }
+}

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?;
 

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;

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<Self, Error> {
+        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 {

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;

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<usize> {
         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<u8>) -> 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!(