fix(pg): classify authentication failures safely

Signed-off-by: DefinitelyNotJosh1 <krasnogo27@up.edu>

Author: jokrasno

crates/client-api/src/auth.rs

@@ -8,9 +8,8 @@ use http::{request, HeaderValue, StatusCode};
 use serde::{Deserialize, Serialize};
 use spacetimedb::auth::identity::{ConnectionAuthCtx, SpacetimeIdentityClaims};
 use spacetimedb::auth::identity::{JwtError, JwtErrorKind};
-use spacetimedb::auth::token_validation::{
-    new_validator, DefaultValidator, TokenSigner, TokenValidationError, TokenValidator,
-};
+use spacetimedb::auth::token_validation::{new_validator, DefaultValidator, TokenSigner, TokenValidator};
+pub use spacetimedb::auth::token_validation::{TokenValidationError, TokenValidationErrorCategory};
 use spacetimedb::auth::JwtKeys;
 use spacetimedb::energy::FunctionBudget;
 use spacetimedb::identity::Identity;

crates/core/src/auth/token_validation.rs

@@ -20,7 +20,7 @@ use super::JwtKeys;
 
 #[derive(thiserror::Error, Debug)]
 pub enum TokenValidationError {
-    // TODO: Add real error types.
+    // TODO: Replace the remaining dependency-specific and catch-all variants with domain errors.
 
     // TODO: If we had our own errors defined we wouldn't be locked into this lib.
     #[error("Invalid token: {0}")]
@@ -29,15 +29,50 @@ pub enum TokenValidationError {
     #[error("Specified key ID not found in JWKs")]
     KeyIDNotFound,
 
+    /// The token was decoded, but its claims are not acceptable.
+    #[error(transparent)]
+    InvalidClaims(anyhow::Error),
+
     #[error(transparent)]
     JwkError(#[from] jwks::JwkError),
     #[error(transparent)]
     JwksError(#[from] jwks::JwksError),
+
+    /// The identity provider's validation material could not be obtained.
+    #[error(transparent)]
+    IdentityProviderUnavailable(anyhow::Error),
+
     // The other case is a catch-all for unexpected errors.
     #[error(transparent)]
     Other(#[from] anyhow::Error),
 }
 
+/// The operational category of a token validation failure.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum TokenValidationErrorCategory {
+    /// The supplied token is invalid and the client may retry with different credentials.
+    InvalidCredentials,
+    /// Validation failed because identity-provider data could not be obtained or used.
+    IdentityProvider,
+    /// Validation failed because of an unexpected local condition.
+    Internal,
+}
+
+impl TokenValidationError {
+    /// Classifies this failure for operational handling such as log severity.
+    pub fn category(&self) -> TokenValidationErrorCategory {
+        match self {
+            Self::TokenError(_) | Self::KeyIDNotFound | Self::InvalidClaims(_) => {
+                TokenValidationErrorCategory::InvalidCredentials
+            }
+            Self::JwkError(_) | Self::JwksError(_) | Self::IdentityProviderUnavailable(_) => {
+                TokenValidationErrorCategory::IdentityProvider
+            }
+            Self::Other(_) => TokenValidationErrorCategory::Internal,
+        }
+    }
+}
+
 // A token signer is responsible for signing tokens without doing any validation.
 pub trait TokenSigner: Sync + Send {
     // Serialize the given claims and sign a JWT token with them as the payload.
@@ -158,7 +193,7 @@ impl TokenValidator for DecodingKey {
 
         let data = decode::<IncomingClaims>(token, self, &validation)?;
         let claims = data.claims;
-        claims.try_into().map_err(TokenValidationError::Other)
+        claims.try_into().map_err(TokenValidationError::InvalidClaims)
     }
 }
 
@@ -170,7 +205,7 @@ impl TokenValidator for BasicTokenValidator {
         if let Some(expected_issuer) = &self.issuer
             && *claims.issuer != **expected_issuer
         {
-            return Err(TokenValidationError::Other(anyhow::anyhow!(
+            return Err(TokenValidationError::InvalidClaims(anyhow::anyhow!(
                 "Issuer mismatch: got {:?}, expected {:?}",
                 claims.issuer,
                 expected_issuer
@@ -233,7 +268,11 @@ impl TokenValidator for CachingOidcTokenValidator {
             .cache
             .get(String::from(raw_issuer.clone()).into())
             .await
-            .ok_or_else(|| anyhow::anyhow!("Error fetching public key for issuer {raw_issuer}"))?;
+            .ok_or_else(|| {
+                TokenValidationError::IdentityProviderUnavailable(anyhow::anyhow!(
+                    "Error fetching public key for issuer {raw_issuer}"
+                ))
+            })?;
         validator.validate_token(token).await
     }
 }
@@ -300,7 +339,7 @@ impl TokenValidator for JwksValidator {
         log::debug!("No key id in header. Trying all keys.");
         // TODO: Consider returning an error if no kid is given?
         // For now, lets just try all the keys.
-        let mut last_error = TokenValidationError::Other(anyhow::anyhow!("No kid found"));
+        let mut last_error = TokenValidationError::InvalidClaims(anyhow::anyhow!("No kid found"));
         for (kid, key) in &self.keyset.keys {
             log::debug!("Trying key {kid}");
             let validator = BasicTokenValidator {
@@ -328,14 +367,31 @@ mod tests {
     use crate::auth::identity::{IncomingClaims, SpacetimeIdentityClaims};
     use crate::auth::token_validation::{
         BasicTokenValidator, CachingOidcTokenValidator, FullTokenValidator, OidcTokenValidator, TokenSigner,
-        TokenValidator,
+        TokenValidationError, TokenValidationErrorCategory, TokenValidator,
     };
     use crate::auth::JwtKeys;
     use base64::Engine;
     use openssl::ec::{EcGroup, EcKey};
     use serde_json;
     use spacetimedb_lib::Identity;
 
+    #[test]
+    fn token_validation_error_categories_distinguish_authentication_failures() {
+        assert_eq!(
+            TokenValidationError::KeyIDNotFound.category(),
+            TokenValidationErrorCategory::InvalidCredentials
+        );
+        assert_eq!(
+            TokenValidationError::IdentityProviderUnavailable(anyhow::anyhow!("controlled provider failure"))
+                .category(),
+            TokenValidationErrorCategory::IdentityProvider
+        );
+        assert_eq!(
+            TokenValidationError::Other(anyhow::anyhow!("controlled internal failure")).category(),
+            TokenValidationErrorCategory::Internal
+        );
+    }
+
     #[tokio::test]
     async fn test_local_validator_checks_issuer() -> anyhow::Result<()> {
         // Test that the issuer must match the expected issuer for LocalTokenValidator.

crates/pg/src/authentication_logging.rs

@@ -0,0 +1,129 @@
+use std::fmt::Display;
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub(crate) enum AuthenticationFailureKind {
+    InvalidCredentials,
+    IdentityProvider,
+    Internal,
+}
+
+pub(crate) fn log_authentication_failure(database: &str, kind: AuthenticationFailureKind, err: &impl Display) {
+    match kind {
+        AuthenticationFailureKind::InvalidCredentials => {
+            log::warn!("PG: Authentication failed on database {database}: {err}");
+        }
+        AuthenticationFailureKind::IdentityProvider => {
+            log::error!("PG: Identity provider failed while authenticating to database {database}: {err}");
+        }
+        AuthenticationFailureKind::Internal => {
+            log::error!("PG: Internal authentication failure on database {database}: {err}");
+        }
+    }
+}
+
+#[cfg(test)]
+fn log_authentication_failure_with_supplied_token(
+    database: &str,
+    _supplied_token: &str,
+    kind: AuthenticationFailureKind,
+    err: &impl Display,
+) {
+    log_authentication_failure(database, kind, err);
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use log::{Level, LevelFilter, Log, Metadata, Record};
+    use std::sync::{Mutex, Once};
+
+    const SYNTHETIC_TOKEN: &str = "SYNTHETIC_PG_AUTH_TOKEN_DO_NOT_LOG_5696";
+
+    #[derive(Clone, Debug)]
+    struct CapturedRecord {
+        level: Level,
+        target: String,
+        message: String,
+    }
+
+    struct CapturingLogger {
+        records: Mutex<Vec<CapturedRecord>>,
+    }
+
+    impl Log for CapturingLogger {
+        fn enabled(&self, _metadata: &Metadata<'_>) -> bool {
+            true
+        }
+
+        fn log(&self, record: &Record<'_>) {
+            self.records
+                .lock()
+                .unwrap_or_else(|poisoned| poisoned.into_inner())
+                .push(CapturedRecord {
+                    level: record.level(),
+                    target: record.target().to_owned(),
+                    message: record.args().to_string(),
+                });
+        }
+
+        fn flush(&self) {}
+    }
+
+    static LOGGER: CapturingLogger = CapturingLogger {
+        records: Mutex::new(Vec::new()),
+    };
+    static INSTALL_LOGGER: Once = Once::new();
+    static CAPTURE_LOCK: Mutex<()> = Mutex::new(());
+
+    fn capture_authentication_failure(kind: AuthenticationFailureKind) -> CapturedRecord {
+        let _capture_guard = CAPTURE_LOCK.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
+        INSTALL_LOGGER.call_once(|| {
+            log::set_logger(&LOGGER).expect("test logger should only be installed once");
+            log::set_max_level(LevelFilter::Trace);
+        });
+
+        let mut records = LOGGER.records.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
+        records.clear();
+        drop(records);
+
+        log_authentication_failure_with_supplied_token(
+            "synthetic-database",
+            SYNTHETIC_TOKEN,
+            kind,
+            &"controlled authentication failure",
+        );
+
+        let records = LOGGER.records.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
+        assert_eq!(records.len(), 1, "expected exactly one authentication log record");
+        records[0].clone()
+    }
+
+    fn assert_token_absent(record: &CapturedRecord) {
+        assert!(!record.message.contains(SYNTHETIC_TOKEN), "token leaked in log message");
+        assert!(!record.target.contains(SYNTHETIC_TOKEN), "token leaked in log target");
+    }
+
+    #[test]
+    fn invalid_credentials_are_warned_without_logging_the_token() {
+        let record = capture_authentication_failure(AuthenticationFailureKind::InvalidCredentials);
+
+        assert_eq!(record.level, Level::Warn);
+        assert_token_absent(&record);
+    }
+
+    #[test]
+    fn identity_provider_failures_are_errors_without_logging_the_token() {
+        let record = capture_authentication_failure(AuthenticationFailureKind::IdentityProvider);
+
+        assert_eq!(record.level, Level::Error);
+        assert_token_absent(&record);
+    }
+
+    #[test]
+    fn internal_authentication_failures_are_errors_without_logging_the_token() {
+        let record = capture_authentication_failure(AuthenticationFailureKind::Internal);
+
+        assert_eq!(record.level, Level::Error);
+        assert_token_absent(&record);
+    }
+}

crates/pg/src/lib.rs

@@ -1,2 +1,3 @@
+mod authentication_logging;
 mod encoder;
 pub mod pg_server;

crates/pg/src/pg_server.rs

@@ -23,7 +23,7 @@ use pgwire::messages::startup::Authentication;
 use pgwire::messages::{PgWireBackendMessage, PgWireFrontendMessage};
 use pgwire::tokio::process_socket;
 use spacetimedb_auth::identity::ConnectionAuthCtx;
-use spacetimedb_client_api::auth::validate_token;
+use spacetimedb_client_api::auth::{validate_token, TokenValidationError, TokenValidationErrorCategory};
 use spacetimedb_client_api::routes::database;
 use spacetimedb_client_api::routes::database::{SqlParams, SqlQueryParams};
 use spacetimedb_client_api::{Authorization, ControlStateReadAccess, ControlStateWriteAccess, NodeDelegate};
@@ -37,6 +37,8 @@ use thiserror::Error;
 use tokio::net::TcpListener;
 use tokio::sync::{Mutex, Notify};
 
+use crate::authentication_logging::{log_authentication_failure, AuthenticationFailureKind};
+
 #[derive(Error, Debug)]
 pub(crate) enum PgError {
     #[error("(metadata) {0}")]
@@ -51,6 +53,14 @@ pub(crate) enum PgError {
     Other(#[from] anyhow::Error),
 }
 
+fn authentication_failure_kind(err: &TokenValidationError) -> AuthenticationFailureKind {
+    match err.category() {
+        TokenValidationErrorCategory::InvalidCredentials => AuthenticationFailureKind::InvalidCredentials,
+        TokenValidationErrorCategory::IdentityProvider => AuthenticationFailureKind::IdentityProvider,
+        TokenValidationErrorCategory::Internal => AuthenticationFailureKind::Internal,
+    }
+}
+
 impl From<PgError> for PgWireError {
     fn from(err: PgError) -> Self {
         if let PgError::Pg(err) = err {
@@ -271,10 +281,8 @@ impl<T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDel
                 let claims = match validate_token(&self.ctx, &pwd.password).await {
                     Ok(claims) => claims,
                     Err(err) => {
-                        log::error!(
-                            "PG: Authentication failed for identity `{}` on database {database}: {err}",
-                            pwd.password
-                        );
+                        let kind = authentication_failure_kind(&err);
+                        log_authentication_failure(&database, kind, &err);
                         let err = ErrorInfo::new("FATAL".to_owned(), "28P01".to_owned(), err.to_string());
                         return close_client(client, err).await;
                     }