common_access_token/

cat.rs

use ciborium::value::Value;
use serde_json::json;
use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};

#[cfg(test)]
mod test {
    include!("cat.test.rs");
}

use crate::claims::{ClaimLabel, ClaimValue, Claims};
use crate::cose::{CoseMac, COSE_MAC0_TAG, CWT_TAG};
use crate::error::Error;
use crate::util::{from_base64_url, generate_random_hex, to_base64_no_padding};

/// Validation types for Common Access Tokens.
///
/// Specifies the cryptographic mechanism used to secure the token.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CatValidationTypes {
    /// Message Authentication Code (HMAC)
    Mac,

    /// Digital Signature
    Sign,

    /// No cryptographic protection
    None,
}

impl Default for CatValidationTypes {
    fn default() -> Self {
        Self::Mac
    }
}

/// Configuration options for the CAT validator/generator.
#[derive(Debug, Clone)]
pub struct CatOptions {
    /// Mapping of key IDs to their key material.
    ///
    /// Each key is identified by a string ID and consists of raw bytes.
    /// For HMAC-SHA256, keys should be at least 32 bytes.
    pub keys: HashMap<String, Vec<u8>>,

    /// Whether to expect tokens to have a CWT tag.
    ///
    /// CWT (CBOR Web Token) tags wrap the COSE_Mac0 structure with tag 61.
    /// Set this to true for compatibility with most CAT implementations.
    pub expect_cwt_tag: bool,
}

impl Default for CatOptions {
    fn default() -> Self {
        Self {
            keys: HashMap::new(),
            expect_cwt_tag: true,
        }
    }
}

/// Options for token generation.
#[derive(Debug, Clone)]
pub struct CatGenerateOptions {
    /// Type of cryptographic protection to apply to the token.
    pub token_type: CatValidationTypes,

    /// Algorithm identifier (e.g., "HS256" for HMAC-SHA256).
    pub alg: String,

    /// ID of the key to use from the key store.
    pub kid: String,

    /// Whether to automatically generate a random CWT ID claim.
    ///
    /// When true, a random identifier will be added as the "cti" claim
    /// if one does not already exist in the claims.
    pub generate_cwt_id: bool,
}

impl Default for CatGenerateOptions {
    fn default() -> Self {
        Self {
            token_type: CatValidationTypes::Mac,
            alg: "HS256".to_string(),
            kid: "default".to_string(),
            generate_cwt_id: true,
        }
    }
}

/// Options for token validation.
#[derive(Debug, Clone)]
pub struct CatValidationOptions {
    /// Expected token issuer.
    ///
    /// The token's "iss" claim must match this value.
    pub issuer: String,

    /// List of allowed audiences for the token.
    ///
    /// If provided, the token's "aud" claim must match at least one
    /// of these values. If None, audience validation is skipped.
    pub audience: Option<Vec<String>>,
}

/// Result of token validation.
#[derive(Debug, Clone)]
pub struct CatValidationResult {
    /// The validated token, if successful.
    ///
    /// This will be Some if the token was parsed successfully, even if
    /// validation failed due to expired token, invalid issuer, etc.
    pub cat: Option<CommonAccessToken>,

    /// Error that occurred during validation, if any.
    ///
    /// This will be None if validation succeeded completely.
    pub error: Option<Error>,
}

impl CatValidationResult {
    /// Returns true if the token is valid (no validation errors).
    pub fn is_valid(&self) -> bool {
        self.error.is_none() && self.cat.is_some()
    }

    /// Returns the claims from the token if validation succeeded.
    pub fn claims(&self) -> Option<HashMap<String, serde_json::Value>> {
        self.cat.as_ref().map(|cat| cat.get_claims())
    }
}

/// Common Access Token
#[derive(Debug, Clone)]
pub struct CommonAccessToken {
    /// The claims in the token
    claims: Claims,

    /// The raw token data
    data: Option<Vec<u8>>,

    /// The key ID used to sign/MAC the token
    kid: Option<String>,
}

impl CommonAccessToken {
    /// Create a new CAT with the given claims
    pub fn new(claims_map: HashMap<String, serde_json::Value>) -> Self {
        let claims = Claims::from_map(claims_map);

        CommonAccessToken {
            claims,
            data: None,
            kid: None,
        }
    }

    /// Create a MAC for the token
    pub fn mac(
        &mut self,
        key: &[u8],
        kid: &str,
        _alg: &str,
        no_cwt_tag: bool,
    ) -> Result<(), Error> {
        // Serialize the claims to CBOR
        let mut payload_bytes = Vec::new();
        ciborium::ser::into_writer(&self.claims, &mut payload_bytes)
            .map_err(|_| Error::CborEncoding)?;

        // Create the COSE MAC
        let cose_mac_bytes = CoseMac::create_hmac_sha256(kid, &payload_bytes, key)?;

        if !no_cwt_tag {
            // Add the COSE_Mac0 tag
            let cose_mac_tagged = Value::Tag(COSE_MAC0_TAG, Box::new(Value::Bytes(cose_mac_bytes)));

            // Add the CWT tag
            let cwt_tagged = Value::Tag(CWT_TAG, Box::new(cose_mac_tagged));

            // Serialize the tagged value
            let mut cwt_bytes = Vec::new();
            ciborium::ser::into_writer(&cwt_tagged, &mut cwt_bytes)
                .map_err(|_| Error::CborEncoding)?;

            self.data = Some(cwt_bytes);
        } else {
            self.data = Some(cose_mac_bytes);
        }

        self.kid = Some(kid.to_string());
        Ok(())
    }

    /// Parse a token
    pub fn parse(
        &mut self,
        token: &[u8],
        key: &[u8],
        kid: &str,
        expect_cwt_tag: bool,
    ) -> Result<(), Error> {
        // Decode the token
        let value: Value = ciborium::de::from_reader(token).map_err(|_| Error::CborDecoding)?;

        // Check if it has a CWT tag
        if expect_cwt_tag {
            if let Value::Tag(tag, _) = &value {
                if *tag != CWT_TAG {
                    return Err(Error::ExpectedCwtTag);
                }
            } else {
                return Err(Error::ExpectedCwtTag);
            }
        }

        // Extract the COSE_Mac0 structure
        let cose_mac_bytes = if let Value::Tag(tag, content) = value {
            if tag == CWT_TAG {
                // Extract the COSE_Mac0 structure from the CWT
                if let Value::Tag(inner_tag, inner_content) = *content {
                    if inner_tag == COSE_MAC0_TAG {
                        if let Value::Bytes(bytes) = *inner_content {
                            bytes
                        } else {
                            return Err(Error::UnableToParseToken);
                        }
                    } else {
                        return Err(Error::UnableToParseToken);
                    }
                } else {
                    return Err(Error::UnableToParseToken);
                }
            } else {
                // Direct COSE_Mac0 structure
                if let Value::Bytes(bytes) = *content {
                    bytes
                } else {
                    return Err(Error::UnableToParseToken);
                }
            }
        } else {
            // Raw COSE_Mac0 structure
            let mut bytes = Vec::new();
            ciborium::ser::into_writer(&value, &mut bytes).map_err(|_| Error::CborEncoding)?;
            bytes
        };

        // Verify the MAC and get the payload
        let payload = CoseMac::verify_hmac_sha256(&cose_mac_bytes, key)?;

        // Decode the payload
        let claims: Claims =
            ciborium::de::from_reader(&payload[..]).map_err(|_| Error::CborDecoding)?;

        self.claims = claims;
        self.kid = Some(kid.to_string());
        self.data = Some(token.to_vec());

        Ok(())
    }

    /// Check if the token is acceptable according to the validation options
    pub fn is_acceptable(&self, opts: &CatValidationOptions) -> Result<bool, Error> {
        // Check issuer
        if let Some(ClaimValue::String(iss)) = self.claims.get(ClaimLabel::Iss) {
            if iss != &opts.issuer {
                return Err(Error::InvalidIssuer {
                    expected: opts.issuer.clone(),
                    actual: iss.clone(),
                });
            }
        }

        // Check expiration
        if let Some(ClaimValue::Integer(exp)) = self.claims.get(ClaimLabel::Exp) {
            let now = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs() as i64;

            if *exp < now {
                return Err(Error::TokenExpired);
            }
        }

        // Check audience
        if let Some(audiences) = &opts.audience {
            if let Some(claim_value) = self.claims.get(ClaimLabel::Aud) {
                match claim_value {
                    ClaimValue::String(aud) => {
                        if !audiences.contains(aud) {
                            return Err(Error::InvalidAudience);
                        }
                    }
                    ClaimValue::Array(auds) => {
                        let mut found = false;
                        for aud in auds {
                            if let ClaimValue::String(aud_str) = aud {
                                if audiences.contains(aud_str) {
                                    found = true;
                                    break;
                                }
                            }
                        }
                        if !found {
                            return Err(Error::InvalidAudience);
                        }
                    }
                    _ => return Err(Error::InvalidClaimType),
                }
            }
        }

        // Check not before
        if let Some(ClaimValue::Integer(nbf)) = self.claims.get(ClaimLabel::Nbf) {
            let now = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs() as i64;

            if *nbf > now {
                return Err(Error::TokenNotActive);
            }
        }

        Ok(true)
    }

    /// Get the raw token data
    pub fn raw(&self) -> Option<&[u8]> {
        self.data.as_deref()
    }

    /// Get the claims as a JSON-like map
    pub fn get_claims(&self) -> HashMap<String, serde_json::Value> {
        self.claims.to_map()
    }
}

/// Common Access Token (CAT) validator and generator
#[derive(Debug, Clone)]
pub struct Cat {
    /// Key ID to key mapping
    keys: HashMap<String, Vec<u8>>,

    /// Whether there should be a CWT tag in the token
    expect_cwt_tag: bool,
}

impl Cat {
    /// Creates a new CAT token validator and generator.
    ///
    /// # Arguments
    ///
    /// * `opts` - Configuration options including keys and token format settings
    ///
    /// # Examples
    ///
    /// ```
    /// use common_access_token::{Cat, CatOptions};
    /// use std::collections::HashMap;
    ///
    /// // Create a key
    /// let key = hex::decode("403697de87af64611c1d32a05dab0fe1fcb715a86ab435f1ec99192d79569388").unwrap();
    ///
    /// // Set up the key store
    /// let mut keys = HashMap::new();
    /// keys.insert("my-key".to_string(), key);
    ///
    /// // Create the CAT handler
    /// let cat = Cat::new(CatOptions {
    ///     keys,
    ///     expect_cwt_tag: true,
    /// });
    /// ```
    pub fn new(opts: CatOptions) -> Self {
        Cat {
            keys: opts.keys,
            expect_cwt_tag: opts.expect_cwt_tag,
        }
    }

    /// Adds a key to the key store.
    ///
    /// # Arguments
    ///
    /// * `kid` - Key ID to associate with the key
    /// * `key` - The key material as raw bytes
    ///
    /// # Examples
    ///
    /// ```
    /// use common_access_token::{Cat, CatOptions};
    /// use std::collections::HashMap;
    ///
    /// let mut cat = Cat::new(CatOptions::default());
    /// let key = hex::decode("403697de87af64611c1d32a05dab0fe1fcb715a86ab435f1ec99192d79569388").unwrap();
    /// cat.add_key("my-key", key);
    /// ```
    pub fn add_key(&mut self, kid: impl Into<String>, key: Vec<u8>) {
        self.keys.insert(kid.into(), key);
    }

    /// Validates a CAT token.
    ///
    /// This function attempts to parse and validate the token using the provided options.
    /// It will try each available key in the key store until one succeeds or all fail.
    ///
    /// # Arguments
    ///
    /// * `token` - The base64-encoded token to validate
    /// * `validation_type` - The type of validation to perform (Mac, Sign, None)
    /// * `opts` - Validation options including issuer and audience requirements
    ///
    /// # Returns
    ///
    /// A `Result` containing a `CatValidationResult` with the parsed token and any validation
    /// error that occurred, or an `Error` if token parsing completely failed.
    ///
    /// # Examples
    ///
    /// ```
    /// use common_access_token::{Cat, CatOptions, CatValidationOptions, CatValidationTypes};
    /// use std::collections::HashMap;
    ///
    /// // Create a CAT validator with keys
    /// let mut keys = HashMap::new();
    /// keys.insert("Symmetric256".to_string(),
    ///     hex::decode("403697de87af64611c1d32a05dab0fe1fcb715a86ab435f1ec99192d79569388").unwrap());
    ///
    /// let cat = Cat::new(CatOptions {
    ///     keys,
    ///     expect_cwt_tag: true,
    /// });
    ///
    /// // Validate a token
    /// // In a real scenario, you would use an actual token, not a placeholder
    /// # // We skip the actual validation in the doctest to avoid errors
    /// # let token_example = || {
    /// let token = "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJ1c2VyLTEyMyJ9.SIGNATURE";
    /// let result = cat.validate(
    ///     token,
    ///     CatValidationTypes::Mac,
    ///     CatValidationOptions {
    ///         issuer: "example-issuer".to_string(),
    ///         audience: Some(vec!["api".to_string()]),
    ///     },
    /// );
    /// #     // Example of how to check the result (not actually run in doctest)
    /// #     match result {
    /// #         Ok(validation_result) => {
    /// #             if validation_result.is_valid() {
    /// #                 println!("Token is valid!");
    /// #                 if let Some(claims) = validation_result.claims() {
    /// #                     println!("Subject: {}", claims.get("sub").unwrap_or(&serde_json::json!("none")));
    /// #                 }
    /// #             } else {
    /// #                 println!("Validation failed: {:?}", validation_result.error);
    /// #             }
    /// #         },
    /// #         Err(err) => println!("Error parsing token: {:?}", err),
    /// #     }
    /// # };
    /// ```
    pub fn validate(
        &self,
        token: &str,
        validation_type: CatValidationTypes,
        opts: CatValidationOptions,
    ) -> Result<CatValidationResult, Error> {
        let token_bytes = from_base64_url(token)?;

        match validation_type {
            CatValidationTypes::Mac => {
                let mut cat = None;
                let mut error = None;

                // Try each key
                for (kid, key) in &self.keys {
                    let mut token = CommonAccessToken::new(HashMap::new());
                    match token.parse(&token_bytes, key, kid, self.expect_cwt_tag) {
                        Ok(_) => {
                            cat = Some(token);
                            break;
                        }
                        Err(err) => {
                            error = Some(err);
                        }
                    }
                }

                // If we found a valid token, check if it's acceptable
                if let Some(ref token) = cat {
                    match token.is_acceptable(&opts) {
                        Ok(_) => {
                            return Ok(CatValidationResult { cat, error: None });
                        }
                        Err(err) => {
                            return Ok(CatValidationResult {
                                cat,
                                error: Some(err),
                            });
                        }
                    }
                }

                // If we didn't find a valid token, return the last error
                Ok(CatValidationResult { cat: None, error })
            }
            _ => Err(Error::UnsupportedValidationType),
        }
    }

    /// Generates a CAT token with the provided claims.
    ///
    /// # Arguments
    ///
    /// * `claims` - A map of claim names to values
    /// * `opts` - Options for token generation
    ///
    /// # Returns
    ///
    /// A base64-encoded token string, or an error if generation fails.
    ///
    /// # Examples
    ///
    /// ```
    /// use common_access_token::{Cat, CatOptions, CatGenerateOptions, CatValidationTypes};
    /// use std::collections::HashMap;
    /// use std::time::{SystemTime, UNIX_EPOCH};
    ///
    /// // Create a CAT generator with keys
    /// let mut keys = HashMap::new();
    /// keys.insert("Symmetric256".to_string(),
    ///     hex::decode("403697de87af64611c1d32a05dab0fe1fcb715a86ab435f1ec99192d79569388").unwrap());
    ///
    /// let cat = Cat::new(CatOptions {
    ///     keys,
    ///     expect_cwt_tag: true,
    /// });
    ///
    /// // Create claims
    /// let now = SystemTime::now()
    ///     .duration_since(UNIX_EPOCH)
    ///     .unwrap()
    ///     .as_secs() as i64;
    ///
    /// let mut claims = HashMap::new();
    /// claims.insert("iss".to_string(), serde_json::json!("example-issuer"));
    /// claims.insert("sub".to_string(), serde_json::json!("user-123"));
    /// claims.insert("exp".to_string(), serde_json::json!(now + 3600));
    ///
    /// // Generate the token
    /// let token = cat.generate(
    ///     claims,
    ///     CatGenerateOptions {
    ///         token_type: CatValidationTypes::Mac,
    ///         alg: "HS256".to_string(),
    ///         kid: "Symmetric256".to_string(),
    ///         generate_cwt_id: true,
    ///     },
    /// ).unwrap();
    ///
    /// println!("Generated token: {}", token);
    /// ```
    pub fn generate(
        &self,
        claims: HashMap<String, serde_json::Value>,
        opts: CatGenerateOptions,
    ) -> Result<String, Error> {
        let mut claims = claims.clone();

        // Generate a random CWT ID if requested
        if opts.generate_cwt_id && !claims.contains_key("cti") {
            claims.insert("cti".to_string(), json!(generate_random_hex(16)));
        }

        // Create the token
        let mut token = CommonAccessToken::new(claims);

        match opts.token_type {
            CatValidationTypes::Mac => {
                let key = self
                    .keys
                    .get(&opts.kid)
                    .ok_or_else(|| Error::KeyNotFound(opts.kid.clone()))?;

                token.mac(key, &opts.kid, &opts.alg, !self.expect_cwt_tag)?;

                if let Some(raw) = token.raw() {
                    Ok(to_base64_no_padding(raw))
                } else {
                    Err(Error::MacFailed)
                }
            }
            _ => Err(Error::UnsupportedValidationType),
        }
    }

    /// Creates a builder for constructing claims conveniently
    ///
    /// # Returns
    ///
    /// A new `ClaimsBuilder` instance that can be used to build token claims
    ///
    /// # Examples
    ///
    /// ```
    /// use common_access_token::{Cat, CatOptions, CatGenerateOptions, CatValidationTypes};
    /// use std::collections::HashMap;
    /// use std::time::{SystemTime, UNIX_EPOCH};
    ///
    /// let cat = Cat::new(CatOptions::default());
    ///
    /// let now = SystemTime::now()
    ///     .duration_since(UNIX_EPOCH)
    ///     .unwrap()
    ///     .as_secs() as i64;
    ///
    /// // Build claims using the builder pattern
    /// let claims = cat.claims_builder()
    ///     .issuer("example-issuer")
    ///     .subject("user-123")
    ///     .audience("api-service")
    ///     .expiration(now + 3600)
    ///     .issued_at(now)
    ///     .build();
    /// ```
    pub fn claims_builder(&self) -> ClaimsBuilder {
        ClaimsBuilder::new()
    }
}

/// Builder for constructing Common Access Token claims
pub struct ClaimsBuilder {
    claims: HashMap<String, serde_json::Value>,
}

impl ClaimsBuilder {
    /// Creates a new ClaimsBuilder
    pub fn new() -> Self {
        Self {
            claims: HashMap::new(),
        }
    }

    /// Sets the token issuer (iss claim)
    pub fn issuer(mut self, issuer: impl Into<String>) -> Self {
        self.claims.insert("iss".to_string(), json!(issuer.into()));
        self
    }

    /// Sets the token subject (sub claim)
    pub fn subject(mut self, subject: impl Into<String>) -> Self {
        self.claims.insert("sub".to_string(), json!(subject.into()));
        self
    }

    /// Sets the token audience (aud claim)
    pub fn audience(mut self, audience: impl Into<String>) -> Self {
        self.claims
            .insert("aud".to_string(), json!(audience.into()));
        self
    }

    /// Sets multiple audiences for the token (aud claim as array)
    pub fn audiences(mut self, audiences: impl IntoIterator<Item = impl Into<String>>) -> Self {
        let audiences: Vec<String> = audiences.into_iter().map(Into::into).collect();
        self.claims.insert("aud".to_string(), json!(audiences));
        self
    }

    /// Sets the token expiration time (exp claim)
    pub fn expiration(mut self, exp: i64) -> Self {
        self.claims.insert("exp".to_string(), json!(exp));
        self
    }

    /// Sets the token not-before time (nbf claim)
    pub fn not_before(mut self, nbf: i64) -> Self {
        self.claims.insert("nbf".to_string(), json!(nbf));
        self
    }

    /// Sets the token issued-at time (iat claim)
    pub fn issued_at(mut self, iat: i64) -> Self {
        self.claims.insert("iat".to_string(), json!(iat));
        self
    }

    /// Sets the token ID (cti claim)
    pub fn token_id(mut self, cti: impl Into<String>) -> Self {
        self.claims.insert("cti".to_string(), json!(cti.into()));
        self
    }

    /// Sets the CAT version (catv claim)
    pub fn cat_version(mut self, version: u8) -> Self {
        self.claims.insert("catv".to_string(), json!(version));
        self
    }

    /// Sets a custom claim with any JSON-serializable value
    pub fn claim(mut self, name: impl Into<String>, value: impl Into<serde_json::Value>) -> Self {
        self.claims.insert(name.into(), value.into());
        self
    }

    /// Builds the final claims map
    pub fn build(self) -> HashMap<String, serde_json::Value> {
        self.claims
    }
}

impl Default for ClaimsBuilder {
    fn default() -> Self {
        Self::new()
    }
}