gitdataai/libs/transport/token.rs

use base64::Engine;
use hmac::{Hmac, KeyInit, Mac};
use models::UserId;
use session::Session;
use sha2::Sha256;

use crate::AppTransport;
use crate::error::AppTransportError;

type HmacSha256 = Hmac<Sha256>;

const VERSION: u8 = 0;
const TOKEN_TTL_SECS: u64 = 600;
const SESSION_TTL_SECS: u64 = 1800;
const MAX_LIFETIME_SECS: i64 = 3000;

const TOKEN_PREFIX: &str = "token:access:";
const SESSION_PREFIX: &str = "transport:session:";

/// Stateless access token — 10min TTL, session lookup key only.
#[derive(Clone, Debug, serde::Deserialize, serde::Serialize)]
pub struct AppTransportToken {
    pub access_token: String,
}

/// Device/client binding for multi-device login distinction.
#[derive(Clone, Debug, serde::Deserialize, serde::Serialize)]
pub struct AppTransportTokenApply {
    pub client_id: String,
    pub device_id: String,
}

/// Validated token context — WS/HTTP identity, no session internals exposed.
#[derive(Clone, Debug, serde::Deserialize, serde::Serialize)]
pub struct AppTransportTokenContext {
    pub user_id: UserId,
    pub device_id: String,
    pub client_id: String,
}

/// Token binary payload (local decode, no Redis roundtrip):
///   [0]     version (1 byte)
///   [1..17] user_id (16 bytes, UUID big-endian)
///   [17..25] created_at (8 bytes, i64 Unix secs big-endian)
///   [25..57] HMAC-SHA256(bytes[0..25])
struct TokenPayload {
    user_id: UserId,
    created_at: i64,
}

impl TokenPayload {
    const LEN: usize = 57;

    fn encode(&self, signing_key: &[u8]) -> Result<Vec<u8>, AppTransportError> {
        let mut buf = Vec::with_capacity(Self::LEN);
        buf.push(VERSION);
        buf.extend_from_slice(self.user_id.as_bytes());
        buf.extend_from_slice(&self.created_at.to_be_bytes());

        let tag = hmac_sign(signing_key, &buf)?;
        buf.extend_from_slice(&tag);

        Ok(buf)
    }

    fn decode(bytes: &[u8], signing_key: &[u8]) -> Result<Self, AppTransportError> {
        if bytes.len() != Self::LEN || bytes[0] != VERSION {
            return Err(AppTransportError::TokenInvalidOrExpired);
        }

        let expected_tag = hmac_sign(signing_key, &bytes[..25])?;
        if !constant_time_eq(&expected_tag, &bytes[25..]) {
            return Err(AppTransportError::TokenInvalidOrExpired);
        }

        let user_id_bytes: [u8; 16] = bytes[1..17]
            .try_into()
            .map_err(|_: std::array::TryFromSliceError| AppTransportError::TokenInvalidOrExpired)?;
        let user_id = UserId::from_bytes(user_id_bytes);
        let created_at_bytes: [u8; 8] = bytes[17..25]
            .try_into()
            .map_err(|_: std::array::TryFromSliceError| AppTransportError::TokenInvalidOrExpired)?;
        let created_at = i64::from_be_bytes(created_at_bytes);

        Ok(TokenPayload {
            user_id,
            created_at,
        })
    }
}

impl AppTransport {
    fn signing_key(&self) -> Result<[u8; 32], AppTransportError> {
        let secret = self
            .config
            .env
            .get("APP_SESSION_SECRET")
            .map(|s| s.as_str())
            .ok_or(AppTransportError::Internal)?;
        let mut mac = HmacSha256::new_from_slice(secret.as_bytes())
            .map_err(|_| AppTransportError::Internal)?;
        mac.update(b"transport-access-token-signing-key");
        let result = mac.finalize().into_bytes();
        Ok(result.into())
    }

    fn session_hash_key(&self, user_id: &UserId, created_at: i64) -> String {
        format!("{}{}:{}", SESSION_PREFIX, user_id, created_at)
    }

    fn token_redis_key(&self, token_str: &str) -> String {
        format!("{}{}", TOKEN_PREFIX, token_str)
    }

    /// Apply access token: session→user_id, generate HMAC token, bind in Redis.
    /// token:access:{token} → session_hash_key (TTL=10min)
    /// transport:session:{user_id}:{created_at} → Hash (TTL=30min)
    pub async fn apply_transport_access_token(
        &self,
        session: Session,
        apply: AppTransportTokenApply,
    ) -> Result<AppTransportToken, AppTransportError> {
        let user_id = session.user().ok_or(AppTransportError::Unauthorized)?;

        let created_at = chrono::Utc::now().timestamp();
        let signing_key = self.signing_key()?;

        let payload = TokenPayload {
            user_id,
            created_at,
        };
        let token_bytes = payload.encode(&signing_key)?;
        let access_token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(&token_bytes);

        let session_key = self.session_hash_key(&user_id, created_at);
        let token_key = self.token_redis_key(&access_token);

        let mut pipe = redis::Pipeline::new();
        pipe.hset(&session_key, "device_id", &apply.device_id)
            .hset(&session_key, "client_id", &apply.client_id)
            .expire(&session_key, SESSION_TTL_SECS as i64)
            .set_ex(&token_key, &session_key, TOKEN_TTL_SECS);

        let mut conn = self
            .cache
            .conn()
            .await
            .map_err(|_| AppTransportError::Internal)?;
        pipe.query_async::<()>(&mut conn)
            .await
            .map_err(|_| AppTransportError::Internal)?;

        Ok(AppTransportToken { access_token })
    }

    /// Validate access_token: local decode (HMAC) → Redis HGETALL session → user context.
    /// No online status check, no TTL inference.
    pub async fn check_access_token(
        &self,
        access_token: String,
    ) -> Result<AppTransportTokenContext, AppTransportError> {
        let token_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .decode(&access_token)
            .map_err(|_| AppTransportError::TokenInvalidOrExpired)?;

        let signing_key = self.signing_key()?;
        let payload = TokenPayload::decode(&token_bytes, &signing_key)?;

        let elapsed = chrono::Utc::now().timestamp() - payload.created_at;
        if elapsed > MAX_LIFETIME_SECS {
            return Err(AppTransportError::TokenInvalidOrExpired);
        }

        let session_key = self.session_hash_key(&payload.user_id, payload.created_at);

        let mut conn = self
            .cache
            .conn()
            .await
            .map_err(|_| AppTransportError::Internal)?;
        let hash_data: std::collections::HashMap<String, String> = redis::Cmd::new()
            .arg("HGETALL")
            .arg(&session_key)
            .query_async(&mut conn)
            .await
            .map_err(|_| AppTransportError::Internal)?;

        let device_id = hash_data
            .get("device_id")
            .cloned()
            .ok_or(AppTransportError::TokenInvalidOrExpired)?;
        let client_id = hash_data
            .get("client_id")
            .cloned()
            .ok_or(AppTransportError::TokenInvalidOrExpired)?;

        Ok(AppTransportTokenContext {
            user_id: payload.user_id,
            device_id,
            client_id,
        })
    }

    /// Renew access_token TTL — no token value change, max 10min, renew cap 50min.
    pub async fn renew_access_token(
        &self,
        access_token: String,
    ) -> Result<AppTransportToken, AppTransportError> {
        let token_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .decode(&access_token)
            .map_err(|_| AppTransportError::TokenInvalidOrExpired)?;

        let signing_key = self.signing_key()?;
        let payload = TokenPayload::decode(&token_bytes, &signing_key)?;

        let elapsed = chrono::Utc::now().timestamp() - payload.created_at;
        if elapsed > MAX_LIFETIME_SECS {
            return Err(AppTransportError::RenewalLimitExceeded);
        }

        let session_key = self.session_hash_key(&payload.user_id, payload.created_at);
        let token_key = self.token_redis_key(&access_token);

        let mut pipe = redis::Pipeline::new();
        pipe.expire(&session_key, SESSION_TTL_SECS as i64)
            .expire(&token_key, TOKEN_TTL_SECS as i64);

        let mut conn = self
            .cache
            .conn()
            .await
            .map_err(|_| AppTransportError::Internal)?;
        pipe.query_async::<()>(&mut conn)
            .await
            .map_err(|_| AppTransportError::Internal)?;

        Ok(AppTransportToken { access_token })
    }
}

fn hmac_sign(key: &[u8], payload: &[u8]) -> Result<[u8; 32], AppTransportError> {
    let mut mac = HmacSha256::new_from_slice(key).map_err(|_| AppTransportError::Internal)?;
    mac.update(payload);
    Ok(mac.finalize().into_bytes().into())
}

fn constant_time_eq(expected: &[u8; 32], actual: &[u8]) -> bool {
    if actual.len() != 32 {
        return false;
    }
    let mut diff = 0u8;
    for i in 0..32 {
        diff |= expected[i] ^ actual[i];
    }
    diff == 0
}