better-auth-api 1.0.0-alpha.5

//! Persisted JWT signing keys and public JWKS discovery.

use crate::plugins::endpoint_context::EndpointContext;
use better_auth_core::{
    AuthContext, AuthError, AuthInitContext, AuthPlugin, AuthRequest, AuthResponse, AuthResult,
    AuthRoute, AuthSchema, FieldMap, FieldValue, HttpMethod,
};
use chrono::{Duration, Utc};
use josekit::{
    jwk::{self, Jwk},
    jws::{self, JwsSigner},
};
use serde_json::{Map, Value, json};

mod adapter;
mod cache;
mod callbacks;
mod claims;
mod jose;
mod key_generation;
mod keys;
mod native;
mod options;
mod overrides;
mod signing;
mod verification;
pub use callbacks::{JwtAdapterFuture, JwtCallbacks};
pub(crate) use jose::verify_hs256_raw;
pub use native::JwtApi;
pub use options::*;
pub use overrides::{JwtCallOverrides, JwtKeyOptions, JwtTokenOptions};

/// Supported asymmetric signing algorithms.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Deserialize)]
pub enum JwtAlgorithm {
    /// Ed25519, the upstream default.
    #[default]
    #[serde(rename = "EdDSA")]
    EdDsa,
    /// RSA PKCS#1 v1.5 with SHA-256 and a 2048-bit key.
    #[serde(rename = "RS256")]
    Rs256,
    /// ECDSA P-256 with SHA-256.
    #[serde(rename = "ES256")]
    Es256,
    /// ECDSA P-521 with SHA-512.
    #[serde(rename = "ES512")]
    Es512,
    /// RSA PSS with SHA-256.
    #[serde(rename = "PS256")]
    Ps256,
}

impl JwtAlgorithm {
    fn name(self) -> &'static str {
        match self {
            Self::EdDsa => "EdDSA",
            Self::Rs256 => "RS256",
            Self::Es256 => "ES256",
            Self::Es512 => "ES512",
            Self::Ps256 => "PS256",
        }
    }
    fn curve(self) -> Option<&'static str> {
        match self {
            Self::EdDsa => Some("Ed25519"),
            Self::Es256 => Some("P-256"),
            Self::Es512 => Some("P-521"),
            Self::Rs256 | Self::Ps256 => None,
        }
    }
    fn generate(self, modulus_length: u32) -> Result<Jwk, josekit::JoseError> {
        match self {
            Self::EdDsa => Jwk::generate_ed_key(jwk::Ed25519),
            Self::Rs256 | Self::Ps256 => Jwk::generate_rsa_key(modulus_length),
            Self::Es256 => Jwk::generate_ec_key(jwk::P_256),
            Self::Es512 => Jwk::generate_ec_key(jwk::P_521),
        }
    }
}

/// Local signing and discovery options.
#[derive(Clone, better_auth_core::PluginConfig)]
#[plugin(name = "JwtPlugin")]
pub struct JwtPluginConfig {
    /// Use locally managed keys for JWT session cookie caches.
    #[config(default = false)]
    pub session_cookie_cache: bool,
    /// RSA modulus length for the primary algorithm.
    #[config(default = 2048)]
    pub modulus_length: u32,
    /// Additional algorithms that explicit signing requests may provision.
    #[config(default = Vec::new())]
    pub key_pair_configs: Vec<JwtKeyPairConfig>,
    /// Remote discovery URL. Disables local discovery; does not change the verifier's adapter.
    #[config(default = None)]
    pub remote_url: Option<String>,
    /// Application-managed signing callback. Requires `remote_url`.
    #[config(default = None, skip)]
    pub custom_sign: Option<JwtNativeCustomSign>,
    /// Session payload callback, before default claims and subject selection.
    #[config(default = None, skip)]
    pub define_payload: Option<JwtNativeDefinePayload>,
    /// Session subject callback.
    #[config(default = None, skip)]
    pub get_subject: Option<JwtNativeGetSubject>,
    /// Public discovery path.
    #[config(default = "/jwks".to_owned())]
    pub jwks_path: String,
    /// Algorithm used for newly generated keys. Omission uses EdDSA and preserves header fallback during cookie verification.
    #[config(default = None)]
    pub algorithm: Option<JwtAlgorithm>,
    /// Relative token lifetime or absolute expiration timestamp.
    #[config(default = JwtExpiration::After(Duration::minutes(15)), skip)]
    pub expiration_time: JwtExpiration,
    /// Issuer claim; defaults to the configured base URL.
    #[config(default = None)]
    pub issuer: Option<String>,
    /// Audience claim; defaults to the configured base URL.
    #[config(default = None)]
    pub audience: Option<JwtAudience>,
    /// Lifetime of signing keys; no automatic rotation when absent.
    #[config(default = None)]
    pub rotation_interval: Option<Duration>,
    /// Keep expired public keys discoverable for previously issued tokens.
    #[config(default = Duration::days(30))]
    pub grace_period: Duration,
    /// Store plaintext private JWKs instead of secret-encrypted JWKs.
    #[config(default = false)]
    pub disable_private_key_encryption: bool,
    /// Suppress the JWT header on successful session responses.
    #[config(default = false)]
    pub disable_setting_jwt_header: bool,
}

/// Issue session JWTs with keys persisted through `JwksStore`.
pub struct JwtPlugin {
    config: JwtPluginConfig,
}

impl JwtPluginConfig {
    fn primary_algorithm(&self) -> JwtAlgorithm {
        self.algorithm.unwrap_or_default()
    }
}

impl JwtPlugin {
    /// Set a relative duration, an absolute date, or a NumericDate in seconds.
    pub fn expiration_time(mut self, expiration: impl Into<JwtExpiration>) -> Self {
        self.config.expiration_time = expiration.into();
        self
    }
}

#[async_trait::async_trait]
impl<S: AuthSchema> AuthPlugin<S> for JwtPlugin {
    fn name(&self) -> &'static str {
        "jwt"
    }
    fn routes(&self) -> Vec<AuthRoute> {
        vec![
            AuthRoute::get(&self.config.jwks_path, "getJSONWebKeySet"),
            AuthRoute::get("/token", "getJSONWebToken").require_headers(true),
        ]
    }
    fn schema(&self, fields: &mut better_auth_core::plugin_runtime::ModelFields) -> AuthResult<()> {
        let role = better_auth_core::store::schema::EntityRole::Jwk;
        fields.register_model_fields(
            role,
            better_auth_core::plugin_runtime::ModelFields::plugin_native_fields(role),
        )?;
        Ok(())
    }

    async fn on_init(&self, ctx: &mut AuthInitContext<S>) -> AuthResult<()> {
        ctx.extensions.insert(self.config.clone());
        if self.config.custom_sign.is_some() && self.config.remote_url.is_none() {
            return Err(AuthError::better_auth(
                "options.jwks.remoteUrl must be set when using options.jwt.sign",
            ));
        }
        if self.config.remote_url.is_some() && self.config.algorithm.is_none() {
            return Err(AuthError::better_auth(
                "options.jwks.keyPairConfig.alg must be specified when options.jwks.remoteUrl is used for OpenID metadata",
            ));
        }
        if !self.config.jwks_path.starts_with('/') || self.config.jwks_path.contains("..") {
            return Err(AuthError::better_auth(
                "options.jwks.jwksPath must be a non-empty string starting with '/' and not contain '..'",
            ));
        }
        if self.config.session_cookie_cache {
            if !ctx
                .config
                .session
                .cookie_cache
                .as_ref()
                .is_some_and(|cache| {
                    matches!(
                        cache.strategy(),
                        better_auth_core::config::CookieCacheStrategy::Jwt
                    )
                })
            {
                return Err(AuthError::better_auth(
                    "`jwt({ sessionCookieCache: true })` requires `session.cookieCache.strategy = \"jwt\"`.",
                ));
            }
            if self.config.custom_sign.is_some() {
                return Err(AuthError::better_auth(
                    "`jwt({ sessionCookieCache: true })` requires locally managed JWT plugin keys and does not support `jwt.sign`.",
                ));
            }
            let signer: std::sync::Arc<dyn better_auth_core::session::SessionCookieSigner<S>> =
                std::sync::Arc::new(cache::CookieSigner {
                    plugin: Self::with_config(self.config.clone()),
                    runtime: ctx.runtime(),
                });
            ctx.extensions.insert(signer);
        }
        Ok(())
    }
    async fn on_request(
        &self,
        req: &AuthRequest,
        ctx: &AuthContext<S>,
    ) -> AuthResult<Option<AuthResponse>> {
        if req.method() != &HttpMethod::Get
            || (req.path() != self.config.jwks_path && req.path() != "/token")
        {
            return Ok(None);
        }
        let mut endpoint = EndpointContext::new(Some(req), request_body(req)?, ctx);
        if req.path() == self.config.jwks_path {
            return self.jwks(&endpoint).await.map(Some);
        }
        if req.path() == "/token" {
            let session = ctx
                .require_native_session(req)
                .await
                .map_err(|error| match error {
                    AuthError::Unauthenticated => {
                        AuthError::from_code("UNAUTHORIZED", "UNAUTHORIZED", "Unauthorized")
                    }
                    error => error,
                })?;
            endpoint.session = Some(session.clone());
            let token = self.sign_session(Some(session), &endpoint).await?;
            return Ok(Some(AuthResponse::native(
                None,
                FieldMap::from([("token".into(), token)]).into(),
            )));
        }
        Ok(None)
    }
    async fn after_request(
        &self,
        req: &AuthRequest,
        response: &mut AuthResponse,
        ctx: &AuthContext<S>,
    ) -> AuthResult<()> {
        if !(req.path() == "/get-session") {
            return Ok(());
        }
        better_auth_core::observability::instrumentation::with_endpoint_hook(
            &ctx.config,
            req,
            "after",
            "plugin:jwt",
            async {
                if self.config.disable_setting_jwt_header || req.path() != "/get-session" {
                    return Ok(());
                }
                let snapshot = req.native_session_snapshot()?;
                let candidate = match snapshot
                    .as_ref()
                    .filter(|data| data.field_value().is_truthy())
                {
                    Some(data) => Some(data.clone()),
                    None => req.new_session()?,
                };
                let Some(candidate) = candidate.filter(|data| data.field_value().is_truthy())
                else {
                    return Ok(());
                };
                if !candidate.session_value()?.is_truthy() {
                    return Ok(());
                }
                let mut endpoint = EndpointContext::new(Some(req), request_body(req)?, ctx);
                endpoint.session = snapshot.clone();
                endpoint.response = Some(response);
                let token = self.sign_session(snapshot, &endpoint).await?;
                let mut exposed: Vec<_> = response
                    .headers
                    .get("access-control-expose-headers")
                    .map(|value| {
                        value
                            .split(',')
                            .map(str::trim)
                            .filter(|value| !value.is_empty())
                            .map(str::to_owned)
                            .collect()
                    })
                    .unwrap_or_default();
                if !exposed.iter().any(|value| value == "set-auth-jwt") {
                    exposed.push("set-auth-jwt".into());
                }
                response
                    .headers
                    .merge(better_auth_core::Headers::from_field_value(
                        &FieldMap::from([("set-auth-jwt".into(), token)]).into(),
                    )?);
                let _ = response
                    .headers
                    .insert("Access-Control-Expose-Headers", exposed.join(", "));
                Ok(())
            },
        )
        .await
    }
}

impl JwtPlugin {
    async fn jwks<S: AuthSchema>(
        &self,
        endpoint: &EndpointContext<'_, S>,
    ) -> AuthResult<AuthResponse> {
        if self.config.remote_url.is_some() {
            let mut response = AuthResponse::new(404);
            let _ = response.headers.insert("content-type", "application/json");
            return Ok(response);
        }
        let mut keys = self.read_keys(endpoint).await?;
        if !keys.is_truthy() || keys.model_property("length")?.strict_equals(&0.into()) {
            let _ = self.create_key(endpoint).await?;
            keys = self.read_keys(endpoint).await?;
        }
        if keys.is_null() || keys.is_undefined() || !keys.model_property("length")?.is_truthy() {
            return Err(AuthError::better_auth(
                "No key sets found. Make sure you have a key in your database.",
            ));
        }
        let now = Utc::now().timestamp_millis() as f64;
        let grace = self.config.grace_period.num_milliseconds() as f64;
        let live_keys = keys.filter_method(
            &better_auth_core::FieldFunction::contextual(move |_, arguments| {
                let key = arguments.first().unwrap_or(&FieldValue::Undefined);
                if key.model_property("expiresAt")?.is_truthy() {
                    let expiry = better_auth_core::query::field_add(
                        &key.model_property("expiresAt")?
                            .get_time_method("key.expiresAt")?,
                        &grace.into(),
                    )?;
                    return Ok(
                        (better_auth_core::query::field_compare(&expiry, &now.into())?
                            == Some(std::cmp::Ordering::Greater))
                        .into(),
                    );
                }
                Ok(true.into())
            }),
            "keySets",
        )?;
        let default_algorithm: FieldValue = self.config.primary_algorithm().name().into();
        let public = live_keys.map_method(
            &better_auth_core::FieldFunction::contextual(move |_, arguments| {
                let key = arguments.first().unwrap_or(&FieldValue::Undefined);
                let algorithm = keys::algorithm(key, default_algorithm.clone())?;
                let curve = key.model_property("crv")?;
                let mut value = FieldMap::from([
                    ("alg".into(), algorithm),
                    (
                        "crv".into(),
                        if curve.is_null() {
                            FieldValue::Undefined
                        } else {
                            curve
                        },
                    ),
                ]);
                keys::spread(
                    &mut value,
                    keys::parse_json(&key.model_property("publicKey")?)?,
                )?;
                let _ = value.insert("kid".into(), key.model_property("id")?);
                Ok(FieldValue::from(value))
            }),
            "keys",
            None,
        )?;
        Ok(AuthResponse::native(
            None,
            FieldMap::from([("keys".into(), public)]).into(),
        ))
    }
}

fn request_body(request: &AuthRequest) -> AuthResult<better_auth_core::FieldValue> {
    if let Some(context) = better_auth_core::hooks::current_request_hook_context() {
        return Ok(context.body);
    }
    request.input_field_value()
}

fn jose_error(error: josekit::JoseError) -> AuthError {
    AuthError::internal(format!("JWT signing key: {error}"))
}

#[cfg(test)]
mod native_sort_tests;
#[cfg(test)]
mod native_test_support;
#[cfg(test)]
mod native_tests;
#[cfg(test)]
mod session_contract_tests;
#[cfg(test)]
mod tests;

#[cfg(test)]
mod callback_contract_tests;

#[cfg(test)]
mod native_options_tests;