jwtsec 0.3.0

/*!
JWT Manager with Access and Refresh support, IP and User-Agent verification.

This module allows:
- Generate Access and Refresh tokens with different TTLs and Cryptographic algorithms
- Check token validity, including IP and user-agent
- Automatically check the outflow of token
- Get a new Access Token by Valid Refresh-Token

# Example of use

```rust
use std::time::Duration;
use jwtsec::jwt::{JwtManager, TokenType, CryptAlgo};

let manager = JwtManager::new("my_secret", Duration::from_secs(60), Duration::from_secs(3600), CryptAlgo::Sha256);
let ip = "127.0.0.1";
let ua = "Mozilla/5.0";
let accept = "text/html";
let accept_encoding = "gzip";
let accept_language = "en-US";
let cache_control = "no-cache";

let access_token = manager.create_token("user1", ip, ua, TokenType::Access, accept, accept_encoding, accept_language, cache_control).unwrap();
let refresh_token = manager.create_token("user1", ip, ua, TokenType::Refresh, accept, accept_encoding, accept_language, cache_control).unwrap();

let claims = manager.verify_token(&access_token, ip, ua, accept, accept_encoding, accept_language, cache_control).unwrap();
let new_access = manager.refresh_access_token(&refresh_token, ip, ua, accept, accept_encoding, accept_language, cache_control).unwrap();
```
*/

use jsonwebtoken::{
    DecodingKey, EncodingKey, Header, TokenData, Validation, decode, encode, errors::ErrorKind,
    errors::Result,
};
use serde::{Deserialize, Serialize};
use sha2::{
    Sha224, Sha256, Sha384, Sha512,
    digest::{Digest, DynDigest},
};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

/// Token type: access or refresh
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
pub enum TokenType {
    Access,
    Refresh,
}

/// Cryptographic algorithm: Sha224, Sha256, Sha384, Sha512
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
pub enum CryptAlgo {
    Sha224,
    Sha256,
    Sha384,
    Sha512,
}

/// JWT manager tokens
#[derive(Debug)]
pub struct JwtManager {
    /// Base secret in bytes
    base_secret: Vec<u8>,
    /// TTL for Access token
    access_ttl: Duration,
    /// TTL for Refresh token
    refresh_ttl: Duration,
    /// Cryptographic algorithm
    crypt_algo: CryptAlgo,
}

/// Claims structure in token
#[derive(Debug, Serialize, Deserialize)]
pub struct Claims {
    /// User ID
    pub sub: String,
    /// Token expiration time (Unix timestamp)
    pub exp: usize,
    /// User IP address
    pub ip: String,
    /// User-Agent of the client
    pub user_agent: String,
    /// Type of token (access or refresh)
    pub token_type: TokenType,
}

impl JwtManager {
    /// Create new JWT manager
    pub fn new(
        base_secret: &str,
        access_ttl: Duration,
        refresh_ttl: Duration,
        crypt_algo: CryptAlgo,
    ) -> Self {
        JwtManager {
            base_secret: base_secret.as_bytes().to_vec(),
            access_ttl,
            refresh_ttl,
            crypt_algo,
        }
    }

    /// Get Hasher instance
    fn get_hasher(&self) -> Box<dyn DynDigest> {
        match self.crypt_algo {
            CryptAlgo::Sha224 => Box::new(Sha224::new()),
            CryptAlgo::Sha256 => Box::new(Sha256::new()),
            CryptAlgo::Sha384 => Box::new(Sha384::new()),
            CryptAlgo::Sha512 => Box::new(Sha512::new()),
        }
    }

    /// Encryption key generation on the fly
    fn generate_key(
        &self,
        ip: &[u8],
        user_agent: &[u8],
        accept: &[u8],
        accept_encoding: &[u8],
        accept_language: &[u8],
        cache_control: &[u8],
    ) -> [u8; 32] {
        let mut hasher = self.get_hasher();
        hasher.update(ip);
        hasher.update(user_agent);
        hasher.update(accept);
        hasher.update(accept_encoding);
        hasher.update(accept_language);
        hasher.update(cache_control);
        hasher.update(&self.base_secret);

        let result = hasher.finalize();

        let mut key = [0u8; 32];
        key.copy_from_slice(&result[..32.min(result.len())]);
        key
    }

    /// Returns the current time in Unix seconds
    fn now_ts() -> usize {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_else(|_| Duration::from_secs(0))
            .as_secs() as usize
    }

    /// Creating a specified type token
    pub fn create_token(
        &self,
        subject: &str,
        ip: &str,
        user_agent: &str,
        token_type: TokenType,
        accept: &str,
        accept_encoding: &str,
        accept_language: &str,
        cache_control: &str,
    ) -> Result<String> {
        let ttl = match token_type {
            TokenType::Access => self.access_ttl,
            TokenType::Refresh => self.refresh_ttl,
        };
        let exp = Self::now_ts() + ttl.as_secs() as usize;

        let claims = Claims {
            sub: subject.to_string(),
            exp,
            ip: ip.to_string(),
            user_agent: user_agent.to_string(),
            token_type,
        };

        let key = self.generate_key(
            ip.as_bytes(),
            user_agent.as_bytes(),
            accept.as_bytes(),
            accept_encoding.as_bytes(),
            accept_language.as_bytes(),
            cache_control.as_bytes(),
        );

        encode(&Header::default(), &claims, &EncodingKey::from_secret(&key))
    }

    /// Checking token for validity and expiration
    pub fn verify_token(
        &self,
        token: &str,
        ip: &str,
        user_agent: &str,
        accept: &str,
        accept_encoding: &str,
        accept_language: &str,
        cache_control: &str,
    ) -> Result<TokenData<Claims>> {
        let key = self.generate_key(
            ip.as_bytes(),
            user_agent.as_bytes(),
            accept.as_bytes(),
            accept_encoding.as_bytes(),
            accept_language.as_bytes(),
            cache_control.as_bytes(),
        );

        let data = decode::<Claims>(
            token,
            &DecodingKey::from_secret(&key),
            &Validation::default(),
        )?;

        if data.claims.exp < Self::now_ts() {
            return Err(jsonwebtoken::errors::Error::from(
                ErrorKind::ExpiredSignature,
            ));
        }

        Ok(data)
    }

    /// Creation of a new Access Token by Refresh token
    pub fn refresh_access_token(
        &self,
        refresh_token: &str,
        ip: &str,
        user_agent: &str,
        accept: &str,
        accept_encoding: &str,
        accept_language: &str,
        cache_control: &str,
    ) -> Result<String> {
        let data = self.verify_token(
            refresh_token,
            ip,
            user_agent,
            accept,
            accept_encoding,
            accept_language,
            cache_control,
        )?;

        if data.claims.token_type != TokenType::Refresh {
            return Err(jsonwebtoken::errors::Error::from(ErrorKind::InvalidToken));
        }

        self.create_token(
            &data.claims.sub,
            ip,
            user_agent,
            TokenType::Access,
            accept,
            accept_encoding,
            accept_language,
            cache_control,
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    #[test]
    fn test_access_and_refresh() {
        let manager = JwtManager::new(
            "secret",
            Duration::from_secs(2),
            Duration::from_secs(10),
            CryptAlgo::Sha256,
        );
        let ip = "127.0.0.1";
        let ua = "Mozilla/5.0";
        let accept = "text/html";
        let accept_encoding = "gzip";
        let accept_language = "en-US";
        let cache_control = "no-cache";

        let access_token = manager
            .create_token(
                "user1",
                ip,
                ua,
                TokenType::Access,
                accept,
                accept_encoding,
                accept_language,
                cache_control,
            )
            .unwrap();

        let refresh_token = manager
            .create_token(
                "user1",
                ip,
                ua,
                TokenType::Refresh,
                accept,
                accept_encoding,
                accept_language,
                cache_control,
            )
            .unwrap();

        let claims = manager
            .verify_token(
                &access_token,
                ip,
                ua,
                accept,
                accept_encoding,
                accept_language,
                cache_control,
            )
            .unwrap();
        assert_eq!(claims.claims.token_type, TokenType::Access);

        let claims = manager
            .verify_token(
                &refresh_token,
                ip,
                ua,
                accept,
                accept_encoding,
                accept_language,
                cache_control,
            )
            .unwrap();
        assert_eq!(claims.claims.token_type, TokenType::Refresh);

        let new_access = manager
            .refresh_access_token(
                &refresh_token,
                ip,
                ua,
                accept,
                accept_encoding,
                accept_language,
                cache_control,
            )
            .unwrap();

        let new_claims = manager
            .verify_token(
                &new_access,
                ip,
                ua,
                accept,
                accept_encoding,
                accept_language,
                cache_control,
            )
            .unwrap();

        assert_eq!(new_claims.claims.token_type, TokenType::Access);
        assert_eq!(new_claims.claims.sub, "user1");
    }
}