light-openid 2.0.3

use crate::utils::crypt_utils::sha256;
use base64::Engine as _;
use base64::engine::general_purpose::URL_SAFE_NO_PAD as BASE64_URL_SAFE_NO_PAD;
use rand::distr::{Alphanumeric, SampleString};

/// Code challenge, as specified in <https://datatracker.ietf.org/doc/rfc7636/>
///
/// See some implementation help in <https://docs.hidglobal.com/activid-as-v8.5/api/openid/leverage-pkce-auth-code-grant-flow.htm>
#[derive(Debug, Clone, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct CodeChallenge {
    pub code_challenge: String,
    pub code_challenge_method: String,
}

impl CodeChallenge {
    pub fn new_s256() -> (Self, String) {
        let code_verifier = Alphanumeric.sample_string(&mut rand::rng(), 70);
        (
            Self {
                code_challenge: BASE64_URL_SAFE_NO_PAD.encode(sha256(code_verifier.as_bytes())),
                code_challenge_method: "S256".to_string(),
            },
            code_verifier,
        )
    }

    pub fn verify_code(&self, verifier: &str) -> bool {
        // RFC 7636 - Section 4.1 - Client Creates a Code Verifier
        if verifier.len() < 43 || verifier.len() > 128 {
            tracing::error!("Invalid code verifier len ({})!", verifier.len());
            return false;
        }

        match self.code_challenge_method.as_str() {
            "plain" => verifier.eq(&self.code_challenge),
            "S256" => {
                let encoded = BASE64_URL_SAFE_NO_PAD.encode(sha256(verifier.as_bytes()));

                encoded.eq(&self.code_challenge)
            }
            s => {
                tracing::error!("Unknown code challenge method: {s}");
                false
            }
        }
    }
}

#[cfg(test)]
mod test {
    use crate::code_challenge::CodeChallenge;

    #[test]
    fn test_plain() {
        let chal = CodeChallenge {
            code_challenge_method: "plain".to_string(),
            code_challenge: "cho9ohv9fuG9ongei9eegaicoungoh0tuna0pohr8thekaing4".to_string(),
        };

        assert!(chal.verify_code("cho9ohv9fuG9ongei9eegaicoungoh0tuna0pohr8thekaing4"));
        assert!(!chal.verify_code("niz8pha4gu0faibahtahghuofoom1yaagh5ipaiWicei2Roo6u"));
    }

    #[test]
    fn test_s256() {
        let chal = CodeChallenge {
            code_challenge_method: "S256".to_string(),
            code_challenge: "uSOvC48D8TMh6RgW-36XppMlMgys-6KAE_wEIev9W2g".to_string(),
        };

        assert!(chal.verify_code("HIwht3lCHfnsruA+7Sq8NP2mPj5cBZe0Ewf23eK9UQhK4TdCIt3SK7Fr/giCdnfjxYQILOPG2D562emggAa2lA=="));
        assert!(!chal.verify_code("text1"));
    }

    #[test]
    fn test_s256_2() {
        let chal = CodeChallenge {
            code_challenge_method: "S256".to_string(),
            code_challenge: "E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM".to_string(),
        };

        assert!(chal.verify_code("dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"));
        assert!(!chal.verify_code("text1"));
    }
}