firebase-token 0.5.0

use std::{collections::HashMap, str::FromStr};

use jsonwebtoken::{
    Algorithm, DecodingKey, Validation, decode, decode_header, errors::Error as JwtError,
};
use serde::de::DeserializeOwned;
use thiserror::Error;

use crate::jwk::JwkKey;

const ISSUER_URL: &str = "https://securetoken.google.com/";

#[derive(Error, Debug)]
pub enum VerificationError {
    #[error("InvalidSignature `{0}`")]
    InvalidSignature(JwtError),
    #[error("InvalidToken `{0}`")]
    InvalidToken(JwtError),
    #[error("NoKid")]
    NoKid,
    #[error("InvalidKid")]
    InvalidKid(String),
    #[error("InvalidDecodingKey `{0}`")]
    InvalidDecodingKey(JwtError),
    #[error("UnknownKeyAlgorithm `{0}`")]
    UnknownKeyAlgorithm(String),
}

pub(crate) struct JwkVerifier {
    keys: HashMap<String, JwkKey>,
    config: JwkConfig,
}

impl JwkVerifier {
    pub(crate) fn new(project_id: &str, keys: Vec<JwkKey>) -> JwkVerifier {
        JwkVerifier {
            keys: keys_to_map(keys),
            config: JwkConfig::new(project_id),
        }
    }

    pub(crate) fn set_keys(&mut self, keys: Vec<JwkKey>) {
        self.keys = keys_to_map(keys);
    }

    pub(crate) fn verify<T: DeserializeOwned>(&self, token: &str) -> Result<T, VerificationError> {
        let header = decode_header(token).map_err(VerificationError::InvalidSignature)?;

        let kid = header.kid.ok_or(VerificationError::NoKid)?;

        let public_key = self
            .keys
            .get(&kid)
            .ok_or(VerificationError::InvalidKid(kid))?;

        let decoding_key = DecodingKey::from_rsa_components(&public_key.n, &public_key.e)
            .map_err(VerificationError::InvalidDecodingKey)?;

        let algorithm = Algorithm::from_str(&public_key.alg)
            .map_err(|_| VerificationError::UnknownKeyAlgorithm(public_key.alg.clone()))?;

        let mut validation = Validation::new(algorithm);
        validation.set_audience(&[&self.config.audience]);
        validation.set_issuer(&[&self.config.issuer]);

        decode(token, &decoding_key, &validation)
            .map(|v| v.claims)
            .map_err(VerificationError::InvalidToken)
    }
}

struct JwkConfig {
    audience: String,
    issuer: String,
}

impl JwkConfig {
    fn new(project_id: &str) -> Self {
        let issuer = format!("{ISSUER_URL}{project_id}");

        Self {
            audience: project_id.to_owned(),
            issuer,
        }
    }
}

fn keys_to_map(keys: Vec<JwkKey>) -> HashMap<String, JwkKey> {
    let mut keys_as_map = HashMap::new();

    for key in keys {
        keys_as_map.insert(key.kid.clone(), key);
    }

    keys_as_map
}