better-auth-api 1.0.0-alpha.5

use super::*;
use josekit::{jwk::Jwk, jws::JwsHeader};
use serde_json::json;
use std::sync::{
    Mutex,
    atomic::{AtomicUsize, Ordering},
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

struct MockServer {
    url: Url,
    body: Arc<Mutex<Value>>,
    requests: Arc<AtomicUsize>,
    task: tokio::task::JoinHandle<()>,
}

impl MockServer {
    async fn start(body: Value) -> Self {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let url = Url::parse(&format!("http://{}/jwks", listener.local_addr().unwrap())).unwrap();
        let body = Arc::new(Mutex::new(body));
        let requests = Arc::new(AtomicUsize::new(0));
        let response_body = body.clone();
        let request_count = requests.clone();
        let task = tokio::spawn(async move {
            loop {
                let (mut stream, _) = listener.accept().await.unwrap();
                let mut bytes = vec![0; 4096];
                let _ = stream.read(&mut bytes).await.unwrap();
                request_count.fetch_add(1, Ordering::Relaxed);
                let body = response_body.lock().unwrap().to_string();
                let response = format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
                    body.len()
                );
                stream.write_all(response.as_bytes()).await.unwrap();
            }
        });
        Self {
            url,
            body,
            requests,
            task,
        }
    }

    fn verifier(&self) -> OidcVerifier {
        let mut verifier = OidcVerifier::new(
            self.url.clone(),
            "https://issuer.example".into(),
            "client".into(),
            None,
        )
        .unwrap();
        verifier.client = Client::builder()
            .no_proxy()
            .redirect(reqwest::redirect::Policy::none())
            .build()
            .unwrap();
        verifier
    }
}

impl Drop for MockServer {
    fn drop(&mut self) {
        self.task.abort();
    }
}

fn rsa_key(kid: &str) -> Jwk {
    let mut key = Jwk::generate_rsa_key(2048).unwrap();
    key.set_key_id(kid);
    key
}

fn public_key(key: &Jwk) -> Value {
    let mut public = key.to_public_key().unwrap();
    if let Some(kid) = key.key_id() {
        public.set_key_id(kid);
    }
    serde_json::to_value(public).unwrap()
}

fn sign(key: &Jwk, claims: Value) -> String {
    let mut header = JwsHeader::new();
    header.set_algorithm("RS256");
    if let Some(kid) = key.key_id() {
        header.set_key_id(kid);
    }
    josekit::jws::serialize_compact(
        claims.to_string().as_bytes(),
        &header,
        &josekit::jws::RS256.signer_from_jwk(key).unwrap(),
    )
    .unwrap()
}

fn claims() -> Value {
    json!({"iss":"https://issuer.example","aud":"client","sub":"subject","nonce":"expected"})
}

fn age_cache(verifier: &OidcVerifier, elapsed: Duration) {
    let mut cache = verifier.cache.write().unwrap();
    let current = cache.as_ref().unwrap();
    *cache = Some(Arc::new(CachedJwks {
        keys: current.keys.clone(),
        fetched_at: Instant::now() - elapsed,
    }));
}

#[tokio::test]
async fn verifies_upstream_jose_signatures_for_all_supported_asymmetric_algorithms() {
    // Generated and verified by the jose version installed with better-auth 1.7.6.
    // Fixtures contain only public keys and tokens for a synthetic issuer.
    #[derive(Deserialize)]
    struct Fixture {
        alg: String,
        key: Value,
        token: String,
    }
    let fixtures: Vec<Fixture> =
        serde_json::from_str(include_str!("oidc_signatures.json")).unwrap();
    for fixture in fixtures {
        let server = MockServer::start(json!({"keys":[fixture.key]})).await;
        let verifier = server.verifier();
        let result = verifier.verify(&fixture.token, Some("expected")).await;
        assert!(result.is_ok(), "{}: {result:?}", fixture.alg);
        let (message, signature) = fixture.token.rsplit_once('.').unwrap();
        let mut bytes = decode_base64(signature).unwrap();
        *bytes.first_mut().unwrap() ^= 1;
        let forged = format!("{message}.{}", URL_SAFE_NO_PAD.encode(bytes));
        assert!(
            verifier.verify(&forged, None).await.is_err(),
            "{} signature tampering",
            fixture.alg
        );
    }
}

#[tokio::test]
async fn verifies_real_signatures_and_oidc_claim_boundaries() {
    let key = rsa_key("original");
    let server = MockServer::start(json!({"keys":[public_key(&key)]})).await;
    let verifier = server.verifier();
    let token = sign(&key, claims());
    let verified = verifier.verify(&token, Some("expected")).await.unwrap();
    assert_eq!(verified.get("sub").and_then(Value::as_str), Some("subject"));
    assert!(verifier.verify(&token, Some("wrong")).await.is_err());
    assert!(verifier.verify(&token, Some("")).await.is_ok());

    let now = chrono::Utc::now().timestamp();
    for (name, value, accepted) in [
        ("iss", json!("https://attacker.example"), false),
        ("aud", json!("another-client"), false),
        ("aud", json!(["another-client", "client", 42]), true),
        ("exp", json!(now), false),
        ("exp", json!(now as f64 + 60.5), true),
        ("exp", json!("9999999999"), false),
        ("exp", Value::Null, false),
        ("nbf", json!(now + 60), false),
        ("nbf", json!(now - 60), true),
        ("iat", json!(now + 3600), true),
        ("iat", Value::Null, false),
        ("nonce", json!(42), false),
    ] {
        let mut payload = claims();
        let _ = payload.as_object_mut().unwrap().insert(name.into(), value);
        assert_eq!(
            verifier
                .verify(&sign(&key, payload), Some("expected"))
                .await
                .is_ok(),
            accepted,
            "claim {name}"
        );
    }
    let forged = sign(&rsa_key("original"), claims());
    assert!(verifier.verify(&forged, None).await.is_err());
    assert_eq!(server.requests.load(Ordering::Relaxed), 1);
}

struct NativeNonceProfile {
    expected_nonce: Mutex<Option<String>>,
    calls: AtomicUsize,
}

#[async_trait::async_trait]
impl super::super::generic::GenericOAuthUserInfoHandler for NativeNonceProfile {
    async fn get_user_info(
        &self,
        tokens: &super::super::providers::OAuthUserInfoRequest,
    ) -> better_auth_core::AuthResult<Option<Value>> {
        assert!(!tokens.fields().contains_key("expectedIdTokenNonce"));
        assert!(tokens.expected_id_token_nonce.is_undefined());
        let _ = self.calls.fetch_add(1, Ordering::SeqCst);
        Ok(Some(
            json!({"sub":"subject", "email":"subject@example.com"}),
        ))
    }

    async fn get_user_info_with_context(
        &self,
        tokens: &super::super::providers::OAuthUserInfoRequest,
        context: super::super::generic::GenericOAuthProfileContext<'_>,
    ) -> better_auth_core::AuthResult<Option<Value>> {
        assert_eq!(
            context.expected_nonce(),
            self.expected_nonce.lock().unwrap().as_deref()
        );
        assert_eq!(context.verified_claims().unwrap().as_value(), &claims());
        self.get_user_info(tokens).await
    }
}

#[tokio::test]
async fn generic_profile_uses_request_nonce_for_callback_proxy_and_native_extensions() {
    use super::super::{
        generic::GenericOAuthConfig,
        handlers::callback_user_request,
        providers::{OAuthProvider, OAuthTokenSet},
        resolved::{ResolvedGenericOAuth, ResolvedProvider},
        social_profile::fetch_user_info_for_code,
    };
    use better_auth_core::{FieldDate, FieldMap, Utf16String};

    let fixtures: Vec<Value> = serde_json::from_str(include_str!("oidc_signatures.json")).unwrap();
    let fixture = &fixtures[0];
    let token = fixture["token"].as_str().unwrap();
    let server = MockServer::start(json!({"keys":[fixture["key"]]})).await;
    let verifier = Arc::new(server.verifier());
    let handler = Arc::new(NativeNonceProfile {
        expected_nonce: Mutex::new(None),
        calls: AtomicUsize::new(0),
    });
    let provider = ResolvedProvider {
        config: OAuthProvider::custom("client", "", "", ""),
        generic: Some(ResolvedGenericOAuth {
            config: GenericOAuthConfig {
                get_user_info: Some(handler.clone()),
                ..Default::default()
            },
            issuer: Some("https://issuer.example".into()),
            is_oidc: true,
            verifier: Some(verifier.clone()),
        }),
    };
    let mut accepted_count = 0;
    for (index, (label, extension, state_nonce, accepted)) in [
        (
            "callback override",
            Some("wrong".into()),
            Some("expected"),
            true,
        ),
        (
            "callback mismatch",
            Some("expected".into()),
            Some("wrong"),
            false,
        ),
        ("callback empty", Some("expected".into()), Some(""), true),
        ("proxy extension", Some("expected".into()), None, true),
        ("proxy absent", None, None, true),
        ("own undefined", Some(FieldValue::Undefined), None, true),
        ("null", Some(FieldValue::Null), None, true),
        ("false", Some(false.into()), None, true),
        ("zero", Some(0.0.into()), None, true),
        ("negative zero", Some((-0.0).into()), None, true),
        ("NaN", Some(f64::NAN.into()), None, true),
        ("empty string", Some("".into()), None, true),
        (
            "empty UTF-16",
            Some(FieldValue::Utf16String(Utf16String::from(""))),
            None,
            true,
        ),
        (
            "UTF-16 match",
            Some(FieldValue::Utf16String(Utf16String::from("expected"))),
            None,
            true,
        ),
        ("wrong string", Some("wrong".into()), None, false),
        ("true", Some(true.into()), None, false),
        ("number", Some(17.0.into()), None, false),
        (
            "array",
            Some(vec![FieldValue::from("expected")].into()),
            None,
            false,
        ),
        ("object", Some(FieldMap::new().into()), None, false),
        (
            "date",
            Some(FieldDate::from_milliseconds(0.0).into()),
            None,
            false,
        ),
        (
            "lone surrogate",
            Some(FieldValue::Utf16String(Utf16String::from_units(vec![
                0xd800,
            ]))),
            None,
            false,
        ),
    ]
    .into_iter()
    .enumerate()
    {
        let mut fields = FieldMap::from([("idToken".into(), token.into())]);
        if let Some(extension) = extension {
            let _ = fields.insert("expectedIdTokenNonce".into(), extension);
        }
        let subject_tokens = OAuthTokenSet::from_fields(fields);
        let request = callback_user_request(
            &subject_tokens.fields().into(),
            &better_auth_core::FieldValue::Undefined,
            state_nonce,
        )
        .unwrap();
        *handler.expected_nonce.lock().unwrap() = request
            .expected_id_token_nonce
            .field_value()
            .as_str()
            .map(str::to_owned);
        let response = fetch_user_info_for_code(&provider, request, Some("unused-flow-nonce"))
            .await
            .unwrap();
        assert_eq!(!response.value.is_null(), accepted, "{label}");
        accepted_count += usize::from(accepted);
        assert_eq!(
            handler.calls.load(Ordering::SeqCst),
            accepted_count,
            "{label}"
        );
        assert_eq!(
            verifier.verification_calls.load(Ordering::SeqCst),
            index + 1,
            "{label}"
        );
    }
    assert_eq!(server.requests.load(Ordering::Relaxed), 1);
}

#[tokio::test]
async fn caches_keys_and_reloads_only_after_rotation_cooldown_or_expiry() {
    let original = rsa_key("original");
    let rotated = rsa_key("rotated");
    let server = MockServer::start(json!({"keys":[public_key(&original)]})).await;
    let verifier = server.verifier();
    let original_token = sign(&original, claims());
    let rotated_token = sign(&rotated, claims());
    let (first, second) = tokio::join!(
        verifier.verify(&original_token, None),
        verifier.verify(&original_token, None)
    );
    assert!(first.is_ok() && second.is_ok());
    assert_eq!(server.requests.load(Ordering::Relaxed), 1);

    *server.body.lock().unwrap() = json!({"keys":[public_key(&rotated)]});
    assert!(verifier.verify(&rotated_token, None).await.is_err());
    assert_eq!(server.requests.load(Ordering::Relaxed), 1);
    age_cache(&verifier, RELOAD_COOLDOWN + Duration::from_secs(1));
    assert!(verifier.verify(&rotated_token, None).await.is_ok());
    assert_eq!(server.requests.load(Ordering::Relaxed), 2);

    *server.body.lock().unwrap() = json!({"keys":[]});
    assert!(verifier.verify(&rotated_token, None).await.is_ok());
    age_cache(&verifier, CACHE_MAX_AGE + Duration::from_secs(1));
    assert!(verifier.verify(&rotated_token, None).await.is_err());
    assert_eq!(server.requests.load(Ordering::Relaxed), 3);
}

#[tokio::test]
async fn rejects_ambiguous_keys_and_non_verification_key_metadata() {
    let key = rsa_key("original");
    let public = public_key(&key);
    let server = MockServer::start(json!({"keys":[public.clone(), public.clone()]})).await;
    assert!(
        server
            .verifier()
            .verify(&sign(&key, claims()), None)
            .await
            .is_err()
    );
    for (field, value) in [
        ("use", json!("enc")),
        ("alg", json!("RS512")),
        ("key_ops", json!(["sign"])),
        ("key_ops", json!(["sign", "verify"])),
        ("key_ops", json!(["verify", "verify"])),
        ("ext", json!("true")),
        ("d", json!("private")),
    ] {
        let mut invalid = public.clone();
        let _ = invalid.as_object_mut().unwrap().insert(field.into(), value);
        *server.body.lock().unwrap() = json!({"keys":[invalid]});
        assert!(
            server
                .verifier()
                .verify(&sign(&key, claims()), None)
                .await
                .is_err(),
            "key field {field}"
        );
    }
    *server.body.lock().unwrap() = json!({"keys":[public]});
    let mut verifier = server.verifier();
    verifier.algorithms = Some(vec!["ES256".into()]);
    assert!(verifier.verify(&sign(&key, claims()), None).await.is_err());
}

#[tokio::test]
async fn rejects_jwks_redirects_without_contacting_the_redirect_target() {
    let target = MockServer::start(json!({"keys":[]})).await;
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let redirect = format!("http://{}/jwks", listener.local_addr().unwrap());
    let location = target.url.to_string();
    let task = tokio::spawn(async move {
        let (mut stream, _) = listener.accept().await.unwrap();
        let mut bytes = vec![0; 4096];
        let _ = stream.read(&mut bytes).await.unwrap();
        stream.write_all(format!("HTTP/1.1 302 Found\r\nLocation: {location}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n").as_bytes()).await.unwrap();
    });
    let mut verifier = target.verifier();
    verifier.jwks_url = Url::parse(&redirect).unwrap();
    assert!(matches!(
        verifier
            .verify(&sign(&rsa_key("original"), claims()), None)
            .await,
        Err(OidcError::Status(StatusCode::FOUND))
    ));
    task.await.unwrap();
    assert_eq!(target.requests.load(Ordering::Relaxed), 0);
}

#[tokio::test]
async fn discovery_rejects_invalid_issuer_and_preserves_endpoint_overrides_for_the_caller() {
    let server = MockServer::start(json!({"issuer":"not an absolute URL"})).await;
    assert!(
        fetch_discovery(server.url.as_str(), &HeaderMap::new())
            .await
            .is_err()
    );
    *server.body.lock().unwrap() = json!({"issuer":"https://issuer.example", "authorization_endpoint":"https://issuer.example/authorize", "jwks_uri":"./keys", "id_token_signing_alg_values_supported":["RS256"]});
    let document = fetch_discovery(server.url.as_str(), &HeaderMap::new())
        .await
        .unwrap();
    assert_eq!(document.issuer.as_deref(), Some("https://issuer.example"));
    assert_eq!(document.jwks_uri.as_deref(), Some("./keys"));
}

#[tokio::test]
async fn discovery_signing_metadata_does_not_discard_usable_endpoints() {
    let key = rsa_key("original");
    let server = MockServer::start(Value::Null).await;
    let token = sign(&key, claims());
    for (algorithms, accepted) in [
        (json!("RS256"), true),
        (json!({"alg":"RS256"}), true),
        (json!([]), true),
        (json!(["RS256", 42]), false),
    ] {
        *server.body.lock().unwrap() = json!({
            "issuer":"https://issuer.example",
            "authorization_endpoint":"https://issuer.example/authorize",
            "jwks_uri":server.url,
            "id_token_signing_alg_values_supported":algorithms
        });
        let document = fetch_discovery(server.url.as_str(), &HeaderMap::new())
            .await
            .unwrap();
        assert_eq!(
            document.authorization_endpoint.as_deref(),
            Some("https://issuer.example/authorize")
        );
        let mut verifier = server.verifier();
        verifier.algorithms = document
            .id_token_signing_alg_values_supported
            .filter(|values| !values.is_empty());
        *server.body.lock().unwrap() = json!({"keys":[public_key(&key)]});
        assert_eq!(
            verifier.verify(&token, Some("expected")).await.is_ok(),
            accepted
        );
    }
}