junobuild-auth 0.4.1

use crate::openid::jwt::header::decode_jwt_header;
use crate::openid::jwt::types::cert::{JwkParams, JwkType};
use crate::openid::jwt::types::token::JwtClaims;
use crate::openid::jwt::types::{cert::Jwk, errors::JwtVerifyError};
use crate::openid::utils::nonce::build_nonce;
use crate::state::types::state::Salt;
use jsonwebtoken::{decode, Algorithm, DecodingKey, TokenData, Validation};
use serde::de::DeserializeOwned;

fn pick_key<'a>(kid: &str, jwks: &'a [Jwk]) -> Option<&'a Jwk> {
    jwks.iter().find(|j| j.kid.as_deref() == Some(kid))
}

pub fn verify_openid_jwt<Claims, Custom>(
    jwt: &str,
    issuers: &[&str],
    jwks: &[Jwk],
    salt: &Salt,
    assert_custom: Custom,
) -> Result<TokenData<Claims>, JwtVerifyError>
where
    Claims: DeserializeOwned + JwtClaims,
    Custom: FnOnce(&Claims) -> Result<(), JwtVerifyError>,
{
    // 1) Read header to get `kid`
    let header = decode_jwt_header(jwt).map_err(JwtVerifyError::from)?;

    let kid = header.kid.ok_or(JwtVerifyError::MissingKid)?;

    // 2) Find matching RSA key
    let jwk = pick_key(&kid, jwks).ok_or(JwtVerifyError::NoKeyForKid)?;

    // 3) Extract RSA components - We support only Google at the moment,
    // which always uses RSA keys, and we also did so for GitHub.
    let (n, e) = match (&jwk.kty, &jwk.params) {
        (JwkType::Rsa, JwkParams::Rsa(params)) => (&params.n, &params.e),
        _ => return Err(JwtVerifyError::WrongKeyType),
    };

    // 4) Build decoding key
    let key = DecodingKey::from_rsa_components(n, e)
        .map_err(|e| JwtVerifyError::BadSig(e.to_string()))?;

    // 5) Build validation but disable automatic audience validation
    let mut val = Validation::new(Algorithm::RS256);

    // The forked library uses IC time for time-based checks.

    // We treat the ID token as a fresh proof of identity, so we ignore long-lived expiration (`exp`).
    // We enforce our own short window using update timestamp (`iat`) - 10 minutes currently.
    val.validate_exp = false;

    // We use the library to enforce "not before" (`nbf`) if present.
    // Not strictly required (we already check `iat` freshness), but it doesn't hurt.
    // Note: jsonwebtoken skips this validation when `nbf` is absent from the token.
    val.validate_nbf = true;

    // Let the lib check issuer
    val.set_issuer(issuers);

    // Disable aud checks — we’ll handle audience manually
    val.validate_aud = false;

    let token =
        decode::<Claims>(jwt, &key, &val).map_err(|e| JwtVerifyError::BadSig(e.to_string()))?;

    let c = &token.claims;

    // 6) Checks the nonce - i.e. the caller + salt is present in the jwt
    // This ensures the JWT has not been intercepted and submitted with a different identity.
    let nonce = build_nonce(salt);

    if c.nonce() != Some(nonce.as_str()) {
        return Err(JwtVerifyError::BadClaim("nonce".to_string()));
    }

    // 7) Assert custom fields according consumer's flow
    assert_custom(c)?;

    // 8) Assert expiration
    let now_ns = now_ns();
    const MAX_VALIDITY_WINDOW_NS: u64 = 10 * 60 * 1_000_000_000; // 10 min
    const IAT_FUTURE_SKEW_NS: u64 = 2 * 60 * 1_000_000_000; // 2 min

    let iat_s = c.iat().ok_or(JwtVerifyError::BadClaim("iat".to_string()))?;
    let iat_ns = iat_s.saturating_mul(1_000_000_000);

    // Reject if token is from the future
    // Note: in practice we noticed that using exactly iat for this comparison
    // lead to issue. That is why we add a bit of buffer to the comparison.
    if now_ns < iat_ns.saturating_sub(IAT_FUTURE_SKEW_NS) {
        return Err(JwtVerifyError::BadClaim("iat_future".to_string()));
    }

    // Reject if too old
    if now_ns > iat_ns.saturating_add(MAX_VALIDITY_WINDOW_NS) {
        return Err(JwtVerifyError::BadClaim("iat_expired".to_string()));
    }

    Ok(token)
}

#[cfg(target_arch = "wasm32")]
fn now_ns() -> u64 {
    ic_cdk::api::time()
}

// For test only
#[cfg(not(target_arch = "wasm32"))]
fn now_ns() -> u64 {
    use std::time::{SystemTime, UNIX_EPOCH};
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_nanos() as u64
}

#[cfg(test)]
mod verify_tests {
    use super::verify_openid_jwt;
    use crate::openid::jwt::types::cert::{JwkParams, JwkParamsRsa, JwkType};
    use crate::openid::jwt::types::token::JwtClaims;
    use crate::openid::jwt::types::{cert::Jwk, errors::JwtVerifyError};
    use crate::openid::utils::nonce::build_nonce;
    use crate::state::types::state::Salt;
    use candid::Deserialize;
    use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
    use serde::Serialize;
    use std::time::{SystemTime, UNIX_EPOCH};

    const TEST_RSA_PEM: &str = include_str!("../../../tests/keys/test_rsa.pem");

    const N_B64URL: &str = "qtQHkWpyd489-_bWjRtrvlQX9CwiQreOsi6kNeeySznI8u-8sxyuO3spW1r2pRmu-rc4jnD9vY6eTGZ3WFNIMxe1geXsF_3nQc5fcNJUUZj19BZE4Ud3dCmUQ4ezkslTvBj8RgD-iBJL7BT7YpxpPgvmqQy_9IgYUkDW4I9_e6kME5kVpySvpRznlk73PfAaDkHWmUTN0j2WcxkW09SGJ_f-tStaYXtc4uH5J-PWMRjwsfL66A_sxLxAwUODJ0VUbeDxVFHGJa0L-58_6GYDTqeel1vH4XjezDL8lf53YRyva3aFxGrC_JeLuIUaJOJX1hXWQb2DruB4hVcQX9afrQ";
    const E_B64URL: &str = "AQAB";

    const ISS_GOOGLE: &str = "https://accounts.google.com";
    const AUD_OK: &str = "client-123";
    const KID_OK: &str = "test-kid-1";

    fn now_secs() -> u64 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs()
    }

    fn test_salt() -> Salt {
        [42u8; 32]
    }

    fn header(typ: Option<&str>, kid: Option<&str>) -> Header {
        let mut h = Header::new(Algorithm::RS256);
        h.typ = typ.map(|t| t.to_string());
        h.kid = kid.map(|k| k.to_string());
        h
    }

    #[derive(Debug, Clone, Deserialize, Serialize)]
    pub struct GoogleClaims {
        pub iss: String,
        pub sub: String,
        pub aud: String,
        pub exp: Option<u64>,
        pub nbf: Option<u64>,
        pub iat: Option<u64>,

        pub nonce: Option<String>,

        pub email: Option<String>,
        pub name: Option<String>,
        pub given_name: Option<String>,
        pub family_name: Option<String>,
        pub preferred_username: Option<String>,
        pub picture: Option<String>,
        pub locale: Option<String>,
    }

    impl JwtClaims for GoogleClaims {
        fn iat(&self) -> Option<u64> {
            self.iat
        }
        fn nonce(&self) -> Option<&str> {
            self.nonce.as_deref()
        }
    }

    fn claims(
        iss: &str,
        aud: &str,
        iat: Option<u64>,
        nbf: Option<u64>,
        nonce: Option<&str>,
        exp: Option<u64>,
    ) -> GoogleClaims {
        GoogleClaims {
            iss: iss.into(),
            sub: "sub".into(),
            aud: aud.into(),
            exp,
            nbf,
            iat,
            email: None,
            name: None,
            given_name: None,
            family_name: None,
            preferred_username: None,
            picture: None,
            nonce: nonce.map(|s| s.into()),
            locale: None,
        }
    }

    fn sign_token(h: &Header, c: &GoogleClaims) -> String {
        let enc = EncodingKey::from_rsa_pem(TEST_RSA_PEM.as_bytes()).expect("valid pem");
        encode(h, c, &enc).expect("jwt encode")
    }

    fn jwk_with_kid(kid: &str) -> Jwk {
        Jwk {
            kty: JwkType::Rsa,
            alg: Some("RS256".into()),
            kid: Some(kid.into()),
            params: JwkParams::Rsa(JwkParamsRsa {
                n: N_B64URL.into(),
                e: E_B64URL.into(),
            }),
        }
    }

    fn assert_audience(claims: &GoogleClaims) -> Result<(), JwtVerifyError> {
        if claims.aud != AUD_OK {
            return Err(JwtVerifyError::BadClaim("aud".to_string()));
        }
        Ok(())
    }

    #[test]
    fn verifies_ok() {
        let now = now_secs();
        let salt = test_salt();

        let nonce = build_nonce(&salt);

        let token = sign_token(
            &header(Some("JWT"), Some(KID_OK)),
            &claims(
                ISS_GOOGLE,
                AUD_OK,
                Some(now),
                Some(now - 5),
                Some(&nonce),
                Some(now + 600),
            ),
        );

        let out = verify_openid_jwt(
            &token,
            &[ISS_GOOGLE],
            &[jwk_with_kid(KID_OK)],
            &salt,
            assert_audience,
        )
        .expect("should verify");

        assert_eq!(out.claims.iss, ISS_GOOGLE);
        assert_eq!(out.claims.aud, AUD_OK);
        assert_eq!(out.claims.nonce.as_deref(), Some(nonce.as_str()));
    }

    #[test]
    fn missing_kid() {
        let now = now_secs();
        let salt = test_salt();

        let nonce = build_nonce(&salt);

        let token = sign_token(
            &header(Some("JWT"), None),
            &claims(
                ISS_GOOGLE,
                AUD_OK,
                Some(now),
                Some(now - 5),
                Some(&nonce),
                Some(now + 600),
            ),
        );

        let err = verify_openid_jwt(
            &token,
            &[ISS_GOOGLE],
            &[jwk_with_kid(KID_OK)],
            &salt,
            assert_audience,
        )
        .unwrap_err();
        assert!(matches!(err, JwtVerifyError::MissingKid));
    }

    #[test]
    fn no_key_for_kid() {
        let now = now_secs();
        let salt = test_salt();

        let nonce = build_nonce(&salt);

        let token = sign_token(
            &header(Some("JWT"), Some("kid-unknown")),
            &claims(
                ISS_GOOGLE,
                AUD_OK,
                Some(now),
                Some(now - 5),
                Some(&nonce),
                Some(now + 600),
            ),
        );

        let err = verify_openid_jwt(
            &token,
            &[ISS_GOOGLE],
            &[jwk_with_kid(KID_OK)],
            &salt,
            assert_audience,
        )
        .unwrap_err();
        assert!(matches!(err, JwtVerifyError::NoKeyForKid));
    }

    #[test]
    fn wrong_issuer_is_badsig_from_lib() {
        let now = now_secs();
        let salt = test_salt();

        let nonce = build_nonce(&salt);

        let token = sign_token(
            &header(Some("JWT"), Some(KID_OK)),
            &claims(
                "https://not.google.example",
                AUD_OK,
                Some(now),
                Some(now - 5),
                Some(&nonce),
                Some(now + 600),
            ),
        );

        // Validation.set_issuer rejects this → mapped to BadSig
        let err = verify_openid_jwt(
            &token,
            &[ISS_GOOGLE],
            &[jwk_with_kid(KID_OK)],
            &salt,
            assert_audience,
        )
        .unwrap_err();
        assert!(matches!(err, JwtVerifyError::BadSig(_)));
    }

    #[test]
    fn wrong_typ_is_badclaim_typ() {
        let now = now_secs();
        let salt = test_salt();

        let nonce = build_nonce(&salt);

        let token = sign_token(
            &header(Some("JOT"), Some(KID_OK)),
            &claims(
                ISS_GOOGLE,
                AUD_OK,
                Some(now),
                Some(now - 5),
                Some(&nonce),
                Some(now + 600),
            ),
        );

        let err = verify_openid_jwt(
            &token,
            &[ISS_GOOGLE],
            &[jwk_with_kid(KID_OK)],
            &salt,
            assert_audience,
        )
        .unwrap_err();
        assert!(matches!(err, JwtVerifyError::BadClaim(ref f) if f == "typ"));
    }

    #[test]
    fn bad_audience() {
        let now = now_secs();
        let salt = test_salt();

        let nonce = build_nonce(&salt);

        let token = sign_token(
            &header(Some("JWT"), Some(KID_OK)),
            &claims(
                ISS_GOOGLE,
                "wrong-aud",
                Some(now),
                Some(now - 5),
                Some(&nonce),
                Some(now + 600),
            ),
        );

        let err = verify_openid_jwt(
            &token,
            &[ISS_GOOGLE],
            &[jwk_with_kid(KID_OK)],
            &salt,
            assert_audience,
        )
        .unwrap_err();
        assert!(matches!(err, JwtVerifyError::BadClaim(ref f) if f == "aud"));
    }

    #[test]
    fn bad_nonce() {
        let now = now_secs();
        let salt = test_salt();

        let token = sign_token(
            &header(Some("JWT"), Some(KID_OK)),
            &claims(
                ISS_GOOGLE,
                AUD_OK,
                Some(now),
                Some(now - 5),
                Some("wrong-nonce"),
                Some(now + 600),
            ),
        );

        let err = verify_openid_jwt(
            &token,
            &[ISS_GOOGLE],
            &[jwk_with_kid(KID_OK)],
            &salt,
            assert_audience,
        )
        .unwrap_err();
        assert!(matches!(err, JwtVerifyError::BadClaim(ref f) if f == "nonce"));
    }

    #[test]
    fn iat_too_far_in_future() {
        let now = now_secs();
        let salt = test_salt();

        let nonce = build_nonce(&salt);

        let future = now + 4 * 60; // +4min, threshold is 2min skew
        let token = sign_token(
            &header(Some("JWT"), Some(KID_OK)),
            &claims(
                ISS_GOOGLE,
                AUD_OK,
                Some(future),
                Some(now - 5),
                Some(&nonce),
                Some(now + 600),
            ),
        );

        let err = verify_openid_jwt(
            &token,
            &[ISS_GOOGLE],
            &[jwk_with_kid(KID_OK)],
            &salt,
            assert_audience,
        )
        .unwrap_err();
        assert!(matches!(err, JwtVerifyError::BadClaim(ref f) if f == "iat_future"));
    }

    #[test]
    fn iat_too_old() {
        let now = now_secs();
        let salt = test_salt();

        let nonce = build_nonce(&salt);

        let old = now.saturating_sub(11 * 60); // >10 min
        let token = sign_token(
            &header(Some("JWT"), Some(KID_OK)),
            &claims(
                ISS_GOOGLE,
                AUD_OK,
                Some(old),
                Some(now - 5),
                Some(&nonce),
                Some(now + 600),
            ),
        );

        let err = verify_openid_jwt(
            &token,
            &[ISS_GOOGLE],
            &[jwk_with_kid(KID_OK)],
            &salt,
            assert_audience,
        )
        .unwrap_err();
        assert!(matches!(err, JwtVerifyError::BadClaim(ref f) if f == "iat_expired"));
    }

    #[test]
    fn nbf_in_future_is_rejected_by_lib() {
        let now = now_secs();
        let salt = test_salt();

        let nonce = build_nonce(&salt);

        let nbf_future = now + 300; // +5 min
        let token = sign_token(
            &header(Some("JWT"), Some(KID_OK)),
            &claims(
                ISS_GOOGLE,
                AUD_OK,
                Some(now),
                Some(nbf_future),
                Some(&nonce),
                Some(now + 600),
            ),
        );

        // validate_nbf = true -> jsonwebtoken rejects -> BadSig
        let err = verify_openid_jwt(
            &token,
            &[ISS_GOOGLE],
            &[jwk_with_kid(KID_OK)],
            &salt,
            assert_audience,
        )
        .unwrap_err();
        assert!(matches!(err, JwtVerifyError::BadSig(_)));
    }

    #[test]
    fn bad_signature_with_wrong_key_material() {
        let now = now_secs();
        let salt = test_salt();

        let nonce = build_nonce(&salt);

        let token = sign_token(
            &header(Some("JWT"), Some(KID_OK)),
            &claims(
                ISS_GOOGLE,
                AUD_OK,
                Some(now),
                Some(now - 5),
                Some(&nonce),
                Some(now + 600),
            ),
        );

        // Break the modulus slightly so the provided JWK doesn't match the signature
        // (change the last char; still valid base64url but wrong key)
        let mut bad_n = N_B64URL.to_string();
        let last = bad_n.pop().unwrap();
        bad_n.push(if last == 'A' { 'B' } else { 'A' });

        let bad_jwk = Jwk {
            kty: JwkType::Rsa,
            alg: Some("RS256".into()),
            kid: Some(KID_OK.into()),
            params: JwkParams::Rsa(JwkParamsRsa {
                n: bad_n,
                e: E_B64URL.into(),
            }),
        };

        let err = verify_openid_jwt(&token, &[ISS_GOOGLE], &[bad_jwk], &salt, assert_audience)
            .unwrap_err();
        assert!(matches!(err, JwtVerifyError::BadSig(_)));
    }

    #[test]
    fn decodes_optional_profile_claims() {
        let now = now_secs();
        let salt = test_salt();

        let nonce = build_nonce(&salt);

        let c = GoogleClaims {
            iss: ISS_GOOGLE.into(),
            sub: "sub-123".into(),
            aud: AUD_OK.into(),
            exp: Some(now + 600),
            nbf: Some(now - 5),
            iat: Some(now),
            email: Some("hello@example.com".into()),
            name: Some("World".into()),
            given_name: Some("Hello".into()),
            family_name: Some("World".into()),
            preferred_username: Some("hello_world".into()),
            picture: Some("https://example.com/world.png".into()),
            nonce: Some(nonce.clone()),
            locale: Some("fr-CH".into()),
        };

        let token = sign_token(&header(Some("JWT"), Some(KID_OK)), &c);

        let out = verify_openid_jwt(
            &token,
            &[ISS_GOOGLE],
            &[jwk_with_kid(KID_OK)],
            &salt,
            assert_audience,
        )
        .expect("should verify");

        let claims = out.claims;
        assert_eq!(claims.email.as_deref(), Some("hello@example.com"));
        assert_eq!(claims.name.as_deref(), Some("World"));
        assert_eq!(claims.given_name.as_deref(), Some("Hello"));
        assert_eq!(claims.family_name.as_deref(), Some("World"));
        assert_eq!(claims.preferred_username.as_deref(), Some("hello_world"));
        assert_eq!(
            claims.picture.as_deref(),
            Some("https://example.com/world.png")
        );
        assert_eq!(claims.locale.as_deref(), Some("fr-CH"));
    }
}