diff --git a/Cargo.lock b/Cargo.lock index 7a8da222c03..4c5a0077193 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -663,8 +663,8 @@ dependencies = [ [[package]] name = "arrow-pg" -version = "0.15.0" -source = "git+https://github.com/datafusion-contrib/datafusion-postgres.git?rev=ffb14a52b3d7c7489559812e1ec38096f225ce93#ffb14a52b3d7c7489559812e1ec38096f225ce93" +version = "0.16.0" +source = "git+https://github.com/datafusion-contrib/datafusion-postgres.git?rev=3ae0a6cb30b1c5131687f4637247ff4f84ac11e4#3ae0a6cb30b1c5131687f4637247ff4f84ac11e4" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4234,8 +4234,8 @@ dependencies = [ [[package]] name = "datafusion-pg-catalog" -version = "0.18.3" -source = "git+https://github.com/datafusion-contrib/datafusion-postgres.git?rev=ffb14a52b3d7c7489559812e1ec38096f225ce93#ffb14a52b3d7c7489559812e1ec38096f225ce93" +version = "0.19.0" +source = "git+https://github.com/datafusion-contrib/datafusion-postgres.git?rev=3ae0a6cb30b1c5131687f4637247ff4f84ac11e4#3ae0a6cb30b1c5131687f4637247ff4f84ac11e4" dependencies = [ "arrow-pg", "async-trait", @@ -10711,9 +10711,9 @@ dependencies = [ [[package]] name = "pgwire" -version = "0.41.0" +version = "0.42.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c667f8f4046889d52ed11e17aae954b1816d3df48dd23bdebcf199cef332df93" +checksum = "1452e6541bcc64e166ff9d350c25c09334c0c80e8f60fde5e36ff81768eeaa55" dependencies = [ "async-trait", "aws-lc-rs", diff --git a/Cargo.toml b/Cargo.toml index 20582d8eaa4..1674063f0c6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -141,7 +141,7 @@ datafusion-functions-aggregate-common = "=55.1.0" datafusion-functions-window-common = "=55.1.0" datafusion-optimizer = "=55.1.0" datafusion-orc = "0.10.0" -datafusion-pg-catalog = "0.18.3" +datafusion-pg-catalog = "0.19.0" datafusion-physical-expr = "=55.1.0" datafusion-physical-plan = "=55.1.0" datafusion-proto = "=55.1.0" @@ -382,8 +382,8 @@ datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.g datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-pg-catalog = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } -arrow-pg = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } +datafusion-pg-catalog = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "3ae0a6cb30b1c5131687f4637247ff4f84ac11e4" } +arrow-pg = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "3ae0a6cb30b1c5131687f4637247ff4f84ac11e4" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "9e9019bb1c7040ed956f654e39378dd42ab17884" } # pprof 0.15.0 initializes framehop inside SIGPROF; pin the unreleased fix diff --git a/src/auth/src/common.rs b/src/auth/src/common.rs index d25390ed778..4b88cc60fc4 100644 --- a/src/auth/src/common.rs +++ b/src/auth/src/common.rs @@ -361,6 +361,18 @@ impl PgScramSha256Verifier { &self.salt } + /// `SHA256(ClientKey)`, used to verify the client proof in SCRAM. + pub fn stored_key(&self) -> &[u8] { + &self.stored_key + } + + /// `HMAC(SaltedPassword, "Server Key")`, used to compute the SCRAM + /// server signature. This is secret material and must not be exposed to + /// clients. + pub fn server_key(&self) -> &[u8] { + &self.server_key + } + pub fn verify_plain_password(&self, password: &[u8]) -> Result { let salted_password = pg_scram_sha256_salted_password(password, &self.salt, self.iterations)?; diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 0be0e47eab7..802ee46f9d4 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -24,7 +24,7 @@ api.workspace = true arrow.workspace = true arrow-flight.workspace = true arrow-ipc.workspace = true -arrow-pg = "0.15" +arrow-pg = "0.16" arrow-schema.workspace = true async-trait.workspace = true auth.workspace = true @@ -98,7 +98,7 @@ otel-arrow-rust.workspace = true parking_lot.workspace = true partition.workspace = true pg_interval = "0.6" -pgwire = { version = "0.41", default-features = false, features = [ +pgwire = { version = "0.42.1", default-features = false, features = [ "server-api-aws-lc-rs", "pg-ext-types", ] } diff --git a/src/servers/src/postgres.rs b/src/servers/src/postgres.rs index 1cd5907157b..15aae48a494 100644 --- a/src/servers/src/postgres.rs +++ b/src/servers/src/postgres.rs @@ -29,17 +29,18 @@ pub(crate) const METADATA_SCHEMA: &str = "schema"; use std::collections::HashMap; use std::net::SocketAddr; use std::sync::Arc; +use std::sync::atomic::AtomicBool; -use ::auth::UserProviderRef; +use ::auth::{UserInfoRef, UserProviderRef}; use derive_builder::Builder; -use pgwire::api::auth::{ServerParameterProvider, StartupHandler}; +use pgwire::api::auth::{AuthDriver, ServerParameterProvider, StartupHandler}; use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler}; use pgwire::api::{ClientInfo, ErrorHandler, PgWireServerHandlers}; pub use server::PostgresServer; use session::Session; use session::context::Channel; -use self::auth_handler::PgLoginVerifier; +use self::auth_handler::{PgUserProviderSource, ZeroPidSecretKeyGenerator}; use self::handler::DefaultQueryParser; use crate::query_handler::sql::ServerSqlQueryHandlerRef; @@ -76,9 +77,11 @@ impl ServerParameterProvider for GreptimeDBStartupParameters { pub struct PostgresServerHandlerInner { query_handler: ServerSqlQueryHandlerRef, - login_verifier: PgLoginVerifier, + auth_source: Option>, force_tls: bool, - param_provider: Arc, + auth_driver: AuthDriver, UserInfoRef>, + /// SCRAM selected; failures get the uniform metric label. + is_scram_auth: AtomicBool, batching_enabled: bool, session: Arc, @@ -129,9 +132,14 @@ impl MakePostgresServerHandler { )); let handler = PostgresServerHandlerInner { query_handler: self.query_handler.clone(), - login_verifier: PgLoginVerifier::new(self.user_provider.clone()), + auth_source: self + .user_provider + .clone() + .map(|p| Arc::new(PgUserProviderSource::new(p))), force_tls: self.force_tls, - param_provider: self.param_provider.clone(), + auth_driver: AuthDriver::new(self.param_provider.clone()) + .with_pid_secret_key_generator(Arc::new(ZeroPidSecretKeyGenerator)), + is_scram_auth: AtomicBool::new(false), batching_enabled, session: session.clone(), diff --git a/src/servers/src/postgres/auth_handler.rs b/src/servers/src/postgres/auth_handler.rs index 1fa075a8e17..169ce594f5e 100644 --- a/src/servers/src/postgres/auth_handler.rs +++ b/src/servers/src/postgres/auth_handler.rs @@ -13,29 +13,31 @@ // limitations under the License. use std::fmt::Debug; +use std::sync::Arc; +use std::sync::atomic::Ordering; use ::auth::{ BEARER_TOKEN_USER, Identity, Password, PgAuthInfo, PgScramSha256Verifier, UserInfoRef, UserProviderRef, userinfo_by_name, }; use async_trait::async_trait; -use base64::Engine; -use base64::engine::general_purpose::STANDARD as BASE64; use common_catalog::parse_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_time::Timezone; use futures::{Sink, SinkExt}; use pgwire::api::auth::StartupHandler; -use pgwire::api::auth::sasl::SCRAM_SHA_256_METHOD; -use pgwire::api::{ClientInfo, PgWireConnectionState, auth}; +use pgwire::api::auth::cleartext::CleartextPasswordAuth; +use pgwire::api::auth::mechanism::AuthMechanismFactory; +use pgwire::api::auth::sasl::scram::{ScramAuth, ScramSecret}; +use pgwire::api::auth::{AuthOutcome, AuthSource, LoginInfo as WireLoginInfo}; +use pgwire::api::{ClientInfo, PidSecretKeyGenerator, auth}; use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; use pgwire::messages::response::ErrorResponse; -use pgwire::messages::startup::{Authentication, PasswordMessageFamily, SecretKey}; +use pgwire::messages::startup::SecretKey; use pgwire::messages::{PgWireBackendMessage, PgWireFrontendMessage}; use session::Session; use snafu::IntoError; -use tokio::sync::Mutex; use crate::error::{AuthSnafu, Result}; use crate::metrics::METRIC_AUTH_FAILURE; @@ -44,168 +46,183 @@ use crate::postgres::types::PgErrorCode; use crate::postgres::utils::convert_err; use crate::query_handler::sql::ServerSqlQueryHandlerRef; -pub(crate) struct PgLoginVerifier { - user_provider: Option, - state: Mutex, -} +/// Query cancellation is not supported yet, so announce pid/secret key 0 as +/// before the `AuthDriver` adoption. +pub(crate) struct ZeroPidSecretKeyGenerator; -impl PgLoginVerifier { - pub(crate) fn new(user_provider: Option) -> Self { - Self { - user_provider, - state: Mutex::new(PgAuthenticationState::Initial), - } +impl PidSecretKeyGenerator for ZeroPidSecretKeyGenerator { + fn generate(&self, _client: &dyn ClientInfo) -> (i32, SecretKey) { + (0, SecretKey::I32(0)) } } -enum PgAuthenticationState { - Initial, - Cleartext, - SaslInitial { - auth_info: PgAuthInfo, - }, - SaslFinal { - verifier: PgScramSha256Verifier, - user_info: Option, - channel_binding: String, - nonce: String, - client_first_bare: String, - server_first: String, - }, +fn record_auth_failure(status_code: StatusCode) { + METRIC_AUTH_FAILURE + .with_label_values(&[status_code.as_ref()]) + .inc(); } -#[allow(dead_code)] -struct LoginInfo { - user: Option, - catalog: Option, - schema: Option, - host: String, +/// The [`AuthSource`] shared by pgwire's cleartext and SCRAM mechanisms. +/// Stateless by design: it only holds the user provider, everything about +/// the current login arrives with each call, and the verified +/// [`UserInfoRef`] travels back to the startup handler as the source +/// identity, reported by [`AuthOutcome::Authenticated`]. +/// +/// Read-only: it resolves credentials and checks authorization with the +/// provider, but holds no policy โ€” authorization runs in the startup +/// handler after verification, before the startup completes. +pub(crate) struct PgUserProviderSource { + user_provider: UserProviderRef, } -impl LoginInfo { - pub fn from_client_info(client: &C) -> LoginInfo - where - C: ClientInfo, - { - LoginInfo { - user: client.metadata().get(super::METADATA_USER).map(Into::into), - catalog: client - .metadata() - .get(super::METADATA_CATALOG) - .map(Into::into), - schema: client - .metadata() - .get(super::METADATA_SCHEMA) - .map(Into::into), - host: client.socket_addr().ip().to_string(), - } +impl Debug for PgUserProviderSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PgUserProviderSource") + .finish_non_exhaustive() } } -impl PgLoginVerifier { - async fn auth(&self, login: &LoginInfo, password: &str) -> Result> { - let user_provider = match &self.user_provider { - Some(provider) => provider, - None => return Ok(None), - }; - - let user_name = match &login.user { - Some(name) => name, - None => return Ok(None), - }; - let catalog = match &login.catalog { - Some(name) => name, - None => return Ok(None), - }; - let schema = match &login.schema { - Some(name) => name, - None => return Ok(None), - }; - - let result = if user_name == BEARER_TOKEN_USER { - user_provider - .auth_bearer_token(password, catalog, schema) - .await - } else { - user_provider - .auth( - Identity::UserId(user_name, None), - Password::PlainText(password.to_string().into()), - catalog, - schema, - ) - .await - }; - match result { - Err(e) => { - METRIC_AUTH_FAILURE - .with_label_values(&[e.status_code().as_ref()]) - .inc(); - Err(AuthSnafu.into_error(e)) - } - Ok(user_info) => Ok(Some(user_info)), - } +impl PgUserProviderSource { + pub(crate) fn new(user_provider: UserProviderRef) -> Self { + Self { user_provider } } - async fn postgres_auth_info(&self, login: &LoginInfo) -> Result { - let user_provider = match &self.user_provider { - Some(provider) => provider, - None => return Ok(PgAuthInfo::Cleartext), - }; + /// (user, catalog, schema) of this login. Catalog/schema are parsed from + /// the `database` startup parameter, which was validated at connection + /// start. + fn resolve(login: &WireLoginInfo<'_>) -> Option<(String, String, String)> { + let user = login.user()?.to_string(); + let (catalog, schema) = parse_catalog_and_schema_from_db_string(login.database()?); + Some((user, catalog, schema)) + } - let user_name = match &login.user { - Some(name) => name, - None => return Ok(PgAuthInfo::Cleartext), - }; - if user_name == BEARER_TOKEN_USER { + /// Which mechanism to use for this login. Called once at connection + /// start; the SCRAM machinery resolves the verifier itself at + /// client-first time. + async fn postgres_auth_info(&self, login: &WireLoginInfo<'_>) -> Result { + let Some((user, catalog, _)) = Self::resolve(login) else { return Ok(PgAuthInfo::Cleartext); - } - let catalog = match &login.catalog { - Some(name) => name, - None => return Ok(PgAuthInfo::Cleartext), }; - match user_provider - .postgres_auth_info(Identity::UserId(user_name, None), catalog) + match self + .user_provider + .postgres_auth_info(Identity::UserId(&user, None), &catalog) .await { Err(e) => { - METRIC_AUTH_FAILURE - .with_label_values(&[e.status_code().as_ref()]) - .inc(); + record_auth_failure(e.status_code()); Err(AuthSnafu.into_error(e)) } Ok(auth_info) => Ok(auth_info), } } - async fn authorize(&self, login: &LoginInfo, user_info: &UserInfoRef) -> Result<()> { - let user_provider = match &self.user_provider { - Some(provider) => provider, - None => return Ok(()), + /// Checks with the provider that the user may access this + /// catalog/schema. Called by the startup handler after verification, + /// while nothing terminal has been sent yet. + async fn authorize( + &self, + login: &WireLoginInfo<'_>, + user_info: &UserInfoRef, + ) -> PgWireResult<()> { + let Some((_, catalog, schema)) = Self::resolve(login) else { + return Err(PgWireError::UserNameRequired); }; - let catalog = match &login.catalog { - Some(name) => name, - None => return Ok(()), - }; - let schema = match &login.schema { - Some(name) => name, - None => return Ok(()), - }; - - match user_provider.authorize(catalog, schema, user_info).await { + match self + .user_provider + .authorize(&catalog, &schema, user_info) + .await + { Err(e) => { - METRIC_AUTH_FAILURE - .with_label_values(&[e.status_code().as_ref()]) - .inc(); - Err(AuthSnafu.into_error(e)) + record_auth_failure(e.status_code()); + Err(PgWireError::ApiError(Box::new(AuthSnafu.into_error(e)))) } Ok(()) => Ok(()), } } } +#[async_trait] +impl AuthSource for PgUserProviderSource { + async fn get_password( + &self, + _login: &WireLoginInfo, + ) -> PgWireResult { + Err(PgWireError::ApiError( + "no cleartext password available from the user provider".into(), + )) + } + + /// Cleartext and bearer-token verification; the provider verifies the + /// password itself. Authorization is left to the startup handler. + async fn verify_password( + &self, + login: &WireLoginInfo, + presented_password: &[u8], + ) -> PgWireResult> { + let Some((user, catalog, _schema)) = Self::resolve(login) else { + return Err(PgWireError::UserNameRequired); + }; + let password = String::from_utf8_lossy(presented_password).into_owned(); + + let result = if user == BEARER_TOKEN_USER { + self.user_provider + .authenticate_bearer_token(&password, &catalog) + .await + } else { + self.user_provider + .authenticate( + Identity::UserId(&user, None), + Password::PlainText(password.into()), + ) + .await + }; + match result { + Err(e) => { + record_auth_failure(e.status_code()); + Err(PgWireError::ApiError(Box::new(AuthSnafu.into_error(e)))) + } + Ok(user_info) => Ok(Some(user_info)), + } + } + + async fn get_scram_secret( + &self, + login: &WireLoginInfo, + ) -> PgWireResult<(ScramSecret, Option)> { + match self.postgres_auth_info(login).await { + // For unknown users the provider hands out a mock verifier with + // no identity, so their handshake fails at proof verification + // exactly like a wrong password. + Ok(PgAuthInfo::ScramSha256 { + verifier, + user_info, + }) => Ok((to_scram_secret(&verifier), user_info)), + // SCRAM was selected but the provider no longer backs it: fail + // closed with a mock verifier. + _ => { + let user = login.user().unwrap_or_default(); + Ok(( + to_scram_secret(&PgScramSha256Verifier::mock_for_unknown_user( + user.as_bytes(), + )), + None, + )) + } + } + } +} + +fn to_scram_secret(verifier: &PgScramSha256Verifier) -> ScramSecret { + ScramSecret::Verifier { + salt: verifier.salt().to_vec(), + iterations: verifier.iterations() as usize, + stored_key: verifier.stored_key().to_vec(), + server_key: verifier.server_key().to_vec(), + } +} + fn set_client_info(client: &mut C, session: &Session) where C: ClientInfo, @@ -216,12 +233,35 @@ where if let Some(current_schema) = client.metadata().get(super::METADATA_SCHEMA) { session.set_schema(current_schema.clone()); } +} - // pass generated process id and secret key to client, this information will - // be sent to postgres client for query cancellation. - // use all 0 before we actually supported query cancellation - client.set_pid_and_secret_key(0, SecretKey::I32(0)); - // set userinfo outside +enum MechanismSelection { + /// No user provider configured: skip authentication. + NoAuth, + Cleartext(Arc), + Scram(Arc), + /// Resolving the mechanism failed. + Failed, +} + +impl PostgresServerHandlerInner { + async fn select_mechanism(&self, client: &C) -> MechanismSelection + where + C: ClientInfo + Sync, + { + let Some(source) = self.auth_source.clone() else { + return MechanismSelection::NoAuth; + }; + let login = WireLoginInfo::from_client_info(client); + if login.user() == Some(BEARER_TOKEN_USER) { + return MechanismSelection::Cleartext(source); + } + match source.postgres_auth_info(&login).await { + Err(_) => MechanismSelection::Failed, + Ok(PgAuthInfo::Cleartext) => MechanismSelection::Cleartext(source), + Ok(PgAuthInfo::ScramSha256 { .. }) => MechanismSelection::Scram(source), + } + } } #[async_trait] @@ -232,7 +272,7 @@ impl StartupHandler for PostgresServerHandlerInner { message: PgWireFrontendMessage, ) -> PgWireResult<()> where - C: ClientInfo + Sink + Unpin + Send, + C: ClientInfo + Sink + Unpin + Send + Sync, C::Error: Debug, PgWireError: From<>::Error>, { @@ -248,6 +288,7 @@ impl StartupHandler for PostgresServerHandlerInner { return Ok(()); } + auth::protocol_negotiation(client, startup).await?; auth::save_startup_parameters_to_metadata(client, startup); // check if db is valid @@ -280,57 +321,85 @@ impl StartupHandler for PostgresServerHandlerInner { } } - if self.login_verifier.user_provider.is_some() { - let login_info = LoginInfo::from_client_info(client); - let auth_info = match self.login_verifier.postgres_auth_info(&login_info).await - { - Ok(auth_info) => auth_info, - Err(_) => { + match self.select_mechanism(client).await { + MechanismSelection::NoAuth => { + self.session.set_user_info(userinfo_by_name( + client.metadata().get(super::METADATA_USER).cloned(), + )); + set_client_info(client, &self.session); + self.auth_driver + .complete_authentication(client, None) + .await?; + } + MechanismSelection::Cleartext(source) => { + if self + .auth_driver + .begin_authentication( + client, + CleartextPasswordAuth::new(source).create(), + ) + .await + .is_err() + { return send_password_authentication_failed(client).await; } - }; - client.set_state(PgWireConnectionState::AuthenticationInProgress); - match auth_info { - PgAuthInfo::ScramSha256 { .. } => { - *self.login_verifier.state.lock().await = - PgAuthenticationState::SaslInitial { auth_info }; - client - .send(PgWireBackendMessage::Authentication(Authentication::SASL( - vec![SCRAM_SHA_256_METHOD.to_string()], - ))) - .await?; - } - PgAuthInfo::Cleartext => { - *self.login_verifier.state.lock().await = - PgAuthenticationState::Cleartext; - client - .send(PgWireBackendMessage::Authentication( - Authentication::CleartextPassword, - )) - .await?; + } + MechanismSelection::Scram(source) => { + self.is_scram_auth.store(true, Ordering::Relaxed); + if self + .auth_driver + .begin_authentication(client, ScramAuth::new(source).create()) + .await + .is_err() + { + return send_password_authentication_failed(client).await; } } - } else { - self.session.set_user_info(userinfo_by_name( - client.metadata().get(super::METADATA_USER).cloned(), - )); - set_client_info(client, &self.session); - auth::finish_authentication(client, self.param_provider.as_ref()).await?; + MechanismSelection::Failed => { + return send_password_authentication_failed(client).await; + } } } PgWireFrontendMessage::PasswordMessageFamily(pwd) => { - let login_info = LoginInfo::from_client_info(client); - match self - .authenticate_password_message(client, &login_info, pwd) - .await? - { - PgAuthenticationResult::Continue => {} - PgAuthenticationResult::Success(user_info) => { + match self.auth_driver.advance_authentication(client, pwd).await { + Ok(AuthOutcome::InProgress) => {} + Ok(AuthOutcome::Authenticated { + identity, + final_message, + }) => { + self.is_scram_auth.store(false, Ordering::Relaxed); + let Some(user_info) = identity else { + return send_password_authentication_failed(client).await; + }; + + // Nothing terminal has been sent yet: authorize + // before completing, so a rejection reaches the + // client as an authentication failure. Failures are + // recorded by `authorize` itself. + let authorized = { + let login = WireLoginInfo::from_client_info(client); + match &self.auth_source { + Some(source) => source.authorize(&login, &user_info).await.is_ok(), + None => false, + } + }; + if !authorized { + return send_password_authentication_failed(client).await; + } + self.session.set_user_info(user_info); set_client_info(client, &self.session); - auth::finish_authentication(client, self.param_provider.as_ref()).await?; + self.auth_driver + .complete_authentication(client, final_message) + .await?; } - PgAuthenticationResult::Failed => { + Err(_e) => { + // SCRAM failures carry a uniform label so the metric + // cannot distinguish wrong password, unknown and + // unauthorized users. + if self.is_scram_auth.swap(false, Ordering::Relaxed) { + record_auth_failure(StatusCode::UserPasswordMismatch); + } return send_password_authentication_failed(client).await; } } @@ -341,225 +410,6 @@ impl StartupHandler for PostgresServerHandlerInner { } } -enum PgAuthenticationResult { - Continue, - Success(UserInfoRef), - Failed, -} - -impl PostgresServerHandlerInner { - /// Records a rejected SCRAM attempt in [`METRIC_AUTH_FAILURE`]. The label is - /// intentionally uniform (never `UserNotFound`), so the counter cannot be - /// used to distinguish a wrong password from an unknown user. - fn record_scram_failure(result: PgAuthenticationResult) -> PgAuthenticationResult { - if matches!(result, PgAuthenticationResult::Failed) { - METRIC_AUTH_FAILURE - .with_label_values(&[StatusCode::UserPasswordMismatch.as_ref()]) - .inc(); - } - result - } - - async fn authenticate_password_message( - &self, - client: &mut C, - login_info: &LoginInfo, - pwd: PasswordMessageFamily, - ) -> PgWireResult - where - C: ClientInfo + Sink + Unpin + Send, - C::Error: Debug, - PgWireError: From<>::Error>, - { - let mut state = self.login_verifier.state.lock().await; - let current_state = std::mem::replace(&mut *state, PgAuthenticationState::Initial); - - match current_state { - PgAuthenticationState::Cleartext => { - let pwd = pwd.into_password()?; - drop(state); - match self.login_verifier.auth(login_info, &pwd.password).await { - Ok(Some(user_info)) => Ok(PgAuthenticationResult::Success(user_info)), - Ok(None) | Err(_) => Ok(PgAuthenticationResult::Failed), - } - } - PgAuthenticationState::SaslInitial { auth_info } => { - // A labeled block funnels every rejection through a single - // `record_scram_failure`, so failed SCRAM attempts still show up - // in `METRIC_AUTH_FAILURE` without sprinkling the metric call - // across each early return. - let result = 'sasl: { - let sasl_initial = pwd.into_sasl_initial_response()?; - if sasl_initial.auth_method != SCRAM_SHA_256_METHOD { - break 'sasl PgAuthenticationResult::Failed; - } - let Some(data) = sasl_initial.data else { - break 'sasl PgAuthenticationResult::Failed; - }; - let Some(client_first) = ScramClientFirst::parse(&data) else { - break 'sasl PgAuthenticationResult::Failed; - }; - let PgAuthInfo::ScramSha256 { - verifier, - user_info, - } = auth_info - else { - break 'sasl PgAuthenticationResult::Failed; - }; - - let server_nonce = BASE64.encode(rand::random::<[u8; 18]>()); - let nonce = format!("{}{}", client_first.nonce, server_nonce); - let server_first = format!( - "r={},s={},i={}", - nonce, - BASE64.encode(verifier.salt()), - verifier.iterations() - ); - client - .send(PgWireBackendMessage::Authentication( - Authentication::SASLContinue(server_first.clone().into()), - )) - .await?; - *state = PgAuthenticationState::SaslFinal { - verifier, - user_info, - channel_binding: client_first.channel_binding, - nonce, - client_first_bare: client_first.bare, - server_first, - }; - PgAuthenticationResult::Continue - }; - Ok(Self::record_scram_failure(result)) - } - PgAuthenticationState::SaslFinal { - verifier, - user_info, - channel_binding, - nonce, - client_first_bare, - server_first, - } => { - let result = 'sasl: { - let sasl_response = pwd.into_sasl_response()?; - let Some(client_final) = ScramClientFinal::parse(&sasl_response.data) else { - break 'sasl PgAuthenticationResult::Failed; - }; - if client_final.channel_binding != channel_binding { - break 'sasl PgAuthenticationResult::Failed; - } - if client_final.nonce != nonce { - break 'sasl PgAuthenticationResult::Failed; - } - - let auth_message = format!( - "{},{},{}", - client_first_bare, server_first, client_final.without_proof - ); - let Ok(client_proof) = BASE64.decode(client_final.proof.as_bytes()) else { - break 'sasl PgAuthenticationResult::Failed; - }; - let Ok(Some(server_signature)) = - verifier.verify_client_proof(auth_message.as_bytes(), &client_proof) - else { - break 'sasl PgAuthenticationResult::Failed; - }; - let Some(user_info) = user_info else { - break 'sasl PgAuthenticationResult::Failed; - }; - - drop(state); - if self - .login_verifier - .authorize(login_info, &user_info) - .await - .is_err() - { - // `authorize` already recorded this failure; return early - // to bypass `record_scram_failure` and avoid double-counting. - return Ok(PgAuthenticationResult::Failed); - } - - client - .send(PgWireBackendMessage::Authentication( - Authentication::SASLFinal( - format!("v={}", BASE64.encode(server_signature)).into(), - ), - )) - .await?; - PgAuthenticationResult::Success(user_info) - }; - Ok(Self::record_scram_failure(result)) - } - PgAuthenticationState::Initial => Ok(PgAuthenticationResult::Failed), - } - } -} - -struct ScramClientFirst { - channel_binding: String, - bare: String, - nonce: String, -} - -impl ScramClientFirst { - fn parse(data: &[u8]) -> Option { - let message = std::str::from_utf8(data).ok()?; - let mut parts = message.splitn(3, ','); - let cbind = parts.next()?; - if !matches!(cbind, "n" | "y") { - return None; - } - let authzid = parts.next()?; - if !authzid.is_empty() { - return None; - } - let channel_binding = BASE64.encode(format!("{cbind},,")); - let bare = parts.next()?.to_string(); - let nonce = bare - .split(',') - .find_map(|chunk| chunk.strip_prefix("r="))? - .to_string(); - Some(Self { - channel_binding, - bare, - nonce, - }) - } -} - -struct ScramClientFinal { - channel_binding: String, - nonce: String, - without_proof: String, - proof: String, -} - -impl ScramClientFinal { - fn parse(data: &[u8]) -> Option { - let message = std::str::from_utf8(data).ok()?; - let proof_pos = message.rfind(",p=")?; - let without_proof = message[..proof_pos].to_string(); - let proof = message[proof_pos + 3..].to_string(); - let channel_binding = without_proof - .split(',') - .find_map(|chunk| chunk.strip_prefix("c="))?; - let nonce = without_proof - .split(',') - .find_map(|chunk| chunk.strip_prefix("r="))?; - if nonce.is_empty() || proof.is_empty() { - return None; - } - - Some(Self { - channel_binding: channel_binding.to_string(), - nonce: nonce.to_string(), - without_proof, - proof, - }) - } -} - async fn send_error(client: &mut C, err_info: ErrorInfo) -> PgWireResult<()> where C: ClientInfo + Sink + Unpin + Send, @@ -616,43 +466,3 @@ where Ok(DbResolution::NotFound("Database not specified".to_owned())) } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_scram_client_first_parse() { - let message = b"n,,n=greptime,r=clientnonce"; - let client_first = ScramClientFirst::parse(message).unwrap(); - assert_eq!("biws", client_first.channel_binding); - assert_eq!("n=greptime,r=clientnonce", client_first.bare); - assert_eq!("clientnonce", client_first.nonce); - - let message = b"y,,n=greptime,r=clientnonce"; - let client_first = ScramClientFirst::parse(message).unwrap(); - assert_eq!("eSws", client_first.channel_binding); - assert_eq!("n=greptime,r=clientnonce", client_first.bare); - assert_eq!("clientnonce", client_first.nonce); - - assert!(ScramClientFirst::parse(b"p=tls-server-end-point,,n=greptime,r=nonce").is_none()); - assert!(ScramClientFirst::parse(b"n,a=authzid,n=greptime,r=nonce").is_none()); - } - - #[test] - fn test_scram_client_final_parse() { - let message = b"c=biws,r=clientnonceservernonce,p=dGVzdA=="; - let client_final = ScramClientFinal::parse(message).unwrap(); - assert_eq!("biws", client_final.channel_binding); - assert_eq!("clientnonceservernonce", client_final.nonce); - assert_eq!( - "c=biws,r=clientnonceservernonce", - client_final.without_proof - ); - assert_eq!("dGVzdA==", client_final.proof); - - assert!(ScramClientFinal::parse(b"r=nonce,p=dGVzdA==").is_none()); - assert!(ScramClientFinal::parse(b"c=biws,r=nonce").is_none()); - assert!(ScramClientFinal::parse(b"c=biws,r=,p=dGVzdA==").is_none()); - } -} diff --git a/src/servers/tests/postgres/mod.rs b/src/servers/tests/postgres/mod.rs index 6b2995283f8..23971145a86 100644 --- a/src/servers/tests/postgres/mod.rs +++ b/src/servers/tests/postgres/mod.rs @@ -12,18 +12,22 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::io; use std::net::SocketAddr; use std::sync::Arc; +use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use async_trait::async_trait; use auth::tests::{DatabaseAuthInfo, MockUserProvider}; use auth::{ - BEARER_TOKEN_USER, Identity, Password, UserInfoRef, UserProvider, UserProviderRef, - format_pg_scram_sha256_password_verifier, user_provider_from_option, + BEARER_TOKEN_USER, Identity, Password, PgAuthInfo, PgScramSha256Verifier, UserInfoRef, + UserProvider, UserProviderRef, format_pg_scram_sha256_password_verifier, + user_provider_from_option, }; +use common_base::secrets::ExposeSecret; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_runtime::Builder as RuntimeBuilder; use common_runtime::runtime::BuilderBuild; @@ -40,6 +44,9 @@ use servers::server::Server; use servers::tls::{ReloadableTlsServerConfig, TlsOption}; use table::TableRef; use table::test_util::MemTable; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpStream; +use tokio_postgres::error::SqlState; use tokio_postgres::{Client, Error as PgError, NoTls, SimpleQueryMessage}; use crate::create_testing_instance; @@ -89,6 +96,116 @@ impl UserProvider for BearerProvider { } } +const MIXED_AUTH_SCRAM_USER: &str = "scram_user"; +const MIXED_AUTH_SCRAM_PASSWORD: &[u8] = b"scram-password"; +const MIXED_AUTH_CLEARTEXT_USER: &str = "clear_user"; +const MIXED_AUTH_CLEARTEXT_PASSWORD: &str = "clear-password"; +const MIXED_AUTH_DENIED_USER: &str = "denied_user"; + +/// Selects SCRAM-SHA-256 for [`MIXED_AUTH_SCRAM_USER`] and +/// [`MIXED_AUTH_DENIED_USER`] (whose authorization is denied) and cleartext +/// for everyone else, so one endpoint serves both mechanisms. +struct MixedAuthProvider { + scram_verifier: PgScramSha256Verifier, + selections: Mutex>, + authorize_calls: Mutex>, +} + +#[async_trait] +impl UserProvider for MixedAuthProvider { + fn name(&self) -> &str { + "mixed-auth-test" + } + + async fn authenticate( + &self, + id: Identity<'_>, + password: Password<'_>, + ) -> auth::error::Result { + let Identity::UserId(username, _) = id; + let Password::PlainText(password) = password else { + return auth::error::UnsupportedPasswordTypeSnafu { + password_type: "plain_text", + } + .fail(); + }; + + let verified = if username == MIXED_AUTH_SCRAM_USER || username == MIXED_AUTH_DENIED_USER { + self.scram_verifier + .verify_plain_password(password.expose_secret().as_bytes())? + } else if username == MIXED_AUTH_CLEARTEXT_USER { + password.expose_secret() == MIXED_AUTH_CLEARTEXT_PASSWORD + } else { + return auth::error::UserNotFoundSnafu { + username: username.to_string(), + } + .fail(); + }; + + if verified { + Ok(auth::userinfo_by_name(Some(username.to_string()))) + } else { + auth::error::UserPasswordMismatchSnafu { + username: username.to_string(), + } + .fail() + } + } + + async fn authorize( + &self, + catalog: &str, + schema: &str, + user_info: &UserInfoRef, + ) -> auth::error::Result<()> { + *self + .authorize_calls + .lock() + .unwrap() + .entry(user_info.username().to_string()) + .or_insert(0) += 1; + + if user_info.username() == MIXED_AUTH_DENIED_USER { + return auth::error::AccessDeniedSnafu { + catalog: catalog.to_string(), + schema: schema.to_string(), + username: user_info.username().to_string(), + } + .fail(); + } + Ok(()) + } + + async fn postgres_auth_info( + &self, + id: Identity<'_>, + _catalog: &str, + ) -> auth::error::Result { + let Identity::UserId(username, _) = id; + let auth_info = if username == MIXED_AUTH_SCRAM_USER || username == MIXED_AUTH_DENIED_USER { + PgAuthInfo::ScramSha256 { + verifier: self.scram_verifier.clone(), + user_info: Some(auth::userinfo_by_name(Some(username.to_string()))), + } + } else { + PgAuthInfo::Cleartext + }; + + // Record the mechanism from the returned variant, so the map can + // never claim a mechanism the provider did not hand out. + let method = match &auth_info { + PgAuthInfo::ScramSha256 { .. } => "scram", + PgAuthInfo::Cleartext => "cleartext", + }; + self.selections + .lock() + .unwrap() + .insert(username.to_string(), method); + + Ok(auth_info) + } +} + fn create_postgres_server( table: TableRef, check_pwd: bool, @@ -244,6 +361,92 @@ async fn test_pg_scram_sha256_auth() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_pg_scram_rejects_channel_binding_request() -> Result<()> { + // The server supports no channel binding (no certificate configured), + // so RFC 5802 ยง6 requires rejecting a client-first with a "p" flag. + let verifier = + format_pg_scram_sha256_password_verifier(b"greptime", b"pg-scram-salt", 4096).unwrap(); + let user_provider = + user_provider_from_option(&format!("static_user_provider:cmd:greptime={verifier}")) + .unwrap(); + let (postgres_server, server_port) = + start_test_server_with_user_provider(user_provider, Default::default()).await?; + + // "p" flag: rejected with an error + let response = raw_scram_client_first( + server_port, + b"p=tls-server-end-point,,n=greptime,r=clientnonce", + ) + .await?; + assert_eq!( + b'E', response, + "expected an ErrorResponse for a p-flag client-first" + ); + + // "n" flag (tokio_postgres sends this): handshake continues normally + let response = raw_scram_client_first(server_port, b"n,,n=greptime,r=clientnonce").await?; + assert_eq!( + b'R', response, + "expected a SASLContinue for an n-flag client-first" + ); + + postgres_server.shutdown().await.unwrap(); + Ok(()) +} + +/// A minimal raw SCRAM client: sends the startup message and a +/// SASLInitialResponse carrying `client_first`, returns the message type +/// byte of the server's answer ('R' for Authentication::SASL*, 'E' for +/// ErrorResponse). +async fn raw_scram_client_first(port: u16, client_first: &[u8]) -> Result { + let mut stream = TcpStream::connect(("127.0.0.1", port)).await.unwrap(); + + fn put_i32(buf: &mut Vec, v: i32) { + buf.extend_from_slice(&v.to_be_bytes()); + } + + async fn read_message(stream: &mut TcpStream) -> (u8, Vec) { + let mut header = [0u8; 5]; + stream.read_exact(&mut header).await.unwrap(); + let len = i32::from_be_bytes(header[1..].try_into().unwrap()) as usize - 4; + let mut body = vec![0u8; len]; + stream.read_exact(&mut body).await.unwrap(); + (header[0], body) + } + + // startup message: i32 len | i32 196608 | user/database params + let mut body = Vec::new(); + put_i32(&mut body, 196608); + body.extend_from_slice(b"user\0greptime\0database\0public\0\0"); + let mut message = Vec::new(); + put_i32(&mut message, (body.len() + 4) as i32); + message.extend_from_slice(&body); + stream.write_all(&message).await.unwrap(); + + // Authentication::SASL advertisement: 'R' | len | i32 10 | mechanism cstrings + let (typ, body) = read_message(&mut stream).await; + assert_eq!(b'R', typ, "expected the SASL mechanism advertisement"); + assert_eq!(10, i32::from_be_bytes(body[..4].try_into().unwrap())); + assert!( + body.windows(b"SCRAM-SHA-256".len()) + .any(|w| w == b"SCRAM-SHA-256") + ); + + // SASLInitialResponse: 'p' | len | "SCRAM-SHA-256"\0 | i32 len | client-first + let mut body = Vec::new(); + body.extend_from_slice(b"SCRAM-SHA-256\0"); + put_i32(&mut body, client_first.len() as i32); + body.extend_from_slice(client_first); + let mut message = Vec::new(); + message.push(b'p'); + put_i32(&mut message, (body.len() + 4) as i32); + message.extend_from_slice(&body); + stream.write_all(&message).await.unwrap(); + + Ok(read_message(&mut stream).await.0) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_pg_cleartext_auth_fallback() -> Result<()> { common_telemetry::init_default_ut_logging(); @@ -269,6 +472,103 @@ async fn test_pg_cleartext_auth_fallback() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_pg_mixed_auth_scram_and_cleartext() -> Result<()> { + let provider = Arc::new(MixedAuthProvider { + scram_verifier: PgScramSha256Verifier::from_password( + MIXED_AUTH_SCRAM_PASSWORD, + b"mixed-auth-salt", + 4096, + ) + .unwrap(), + selections: Mutex::new(HashMap::new()), + authorize_calls: Mutex::new(HashMap::new()), + }); + let (postgres_server, server_port) = + start_test_server_with_user_provider(provider.clone(), Default::default()).await?; + + // SCRAM-selected user + let client = create_plain_connection_with_credentials( + server_port, + MIXED_AUTH_SCRAM_USER, + "scram-password", + ) + .await + .unwrap(); + let rows = client + .simple_query("SELECT uint32s FROM numbers LIMIT 1") + .await + .unwrap(); + assert_eq!("0", unwrap_results(&rows)[0]); + + assert!( + create_plain_connection_with_credentials(server_port, MIXED_AUTH_SCRAM_USER, "wrong") + .await + .is_err() + ); + + // cleartext-selected user + let client = create_plain_connection_with_credentials( + server_port, + MIXED_AUTH_CLEARTEXT_USER, + MIXED_AUTH_CLEARTEXT_PASSWORD, + ) + .await + .unwrap(); + let rows = client + .simple_query("SELECT uint32s FROM numbers LIMIT 1") + .await + .unwrap(); + assert_eq!("0", unwrap_results(&rows)[0]); + + assert!( + create_plain_connection_with_credentials(server_port, MIXED_AUTH_CLEARTEXT_USER, "wrong",) + .await + .is_err() + ); + + // unknown user falls back to cleartext and fails authentication + assert!( + create_plain_connection_with_credentials(server_port, "nobody", "whatever") + .await + .is_err() + ); + + // SCRAM-selected user denied by authorization: the client proof + // verifies, authorize rejects, and the failure reaches the client as + // 28P01 before SASLFinal/AuthenticationOk/ReadyForQuery were sent โ€” + // had the completion fired first, this connection would succeed. + let denied = create_plain_connection_with_credentials( + server_port, + MIXED_AUTH_DENIED_USER, + "scram-password", + ) + .await + .unwrap_err(); + assert_eq!(Some(&SqlState::INVALID_PASSWORD), denied.code()); + + // both mechanisms were served by the same endpoint + let selections = provider.selections.lock().unwrap().clone(); + assert_eq!(Some(&"scram"), selections.get(MIXED_AUTH_SCRAM_USER)); + assert_eq!(Some(&"scram"), selections.get(MIXED_AUTH_DENIED_USER)); + assert_eq!( + Some(&"cleartext"), + selections.get(MIXED_AUTH_CLEARTEXT_USER) + ); + + // authorize ran exactly once for every password-verified user โ€” + // including the denied one, proving it was rejected by authorization + // after successful SCRAM verification, not by authentication + let authorize_calls = provider.authorize_calls.lock().unwrap().clone(); + assert_eq!(Some(&1), authorize_calls.get(MIXED_AUTH_SCRAM_USER)); + assert_eq!(Some(&1), authorize_calls.get(MIXED_AUTH_DENIED_USER)); + assert_eq!(Some(&1), authorize_calls.get(MIXED_AUTH_CLEARTEXT_USER)); + assert!(!authorize_calls.contains_key("nobody")); + + postgres_server.shutdown().await.unwrap(); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_bearer_token_auth_over_cleartext() -> Result<()> { let provider = Arc::new(BearerProvider::default());