server-mark2-dev/src/websocket/session.rs

use std::sync::Arc;

use aes_gcm::{Aes256Gcm, aead::AeadMut};
use log::{error, info, warn};
use p256::elliptic_curve::sec1::{FromEncodedPoint, ToEncodedPoint};
use rand::Rng;
use tokio::time::Instant;

use crate::websocket::core::GOOGLE_PUBLIC_ENDPOINT;

pub(crate) struct SecureSession {
    pub uid: String,
    pub cipher: Aes256Gcm,
    pub key_established_at: Instant,
}

#[derive(serde::Deserialize)]
pub(crate) struct HandshakePayload {
    pub token: String,
    pub client_public_key: String, // BASE 64
}

#[derive(serde::Serialize)]
pub(crate) struct HandshakeAck {
    pub status: String,
    pub server_public_key: String,
}

#[derive(serde::Deserialize, serde::Serialize)]
pub(crate) struct EncryptedFrame {
    pub iv: String,         // Initialized vector per message
    pub ciphertext: String, // Encrypted application message
}

#[derive(serde::Deserialize)]
pub(crate) struct FirebaseJwtClaims {
    aud: String, // Audience (Expect Firebase project id),
    sub: String, // Subject (Firebase user uid),
    exp: u64,    // Expiration timestamp
}

pub(crate) async fn refresh_jwk_cache(
    state: Arc<crate::app::AppState>,
) -> Result<(), Box<dyn std::error::Error>> {
    let response: serde_json::Value = reqwest::get(GOOGLE_PUBLIC_ENDPOINT).await?.json().await?;
    let mut new_keys = Vec::new();

    if let Some(obj) = response.as_object() {
        for (_, cert_pem) in obj {
            if let Some(pem_str) = cert_pem.as_str() {
                if let Ok(key) = jsonwebtoken::DecodingKey::from_rsa_pem(pem_str.as_bytes()) {
                    new_keys.push(key);
                }
            }
        }
    }

    {
        let mut cache = state.jwk_encoding_keys.write().unwrap();
        *cache = new_keys;
        info!("Google Jwk Identity cache updated!");
    }

    Ok(())
}

pub(crate) async fn verify_token(
    token: &str,
    state: Arc<crate::app::AppState>,
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
    let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::RS256);
    validation.set_audience(&[&state.firebase_project_id]);
    validation.validate_exp = true;

    let keys = { state.jwk_encoding_keys.read().unwrap() };

    for key in keys.iter() {
        match jsonwebtoken::decode::<FirebaseJwtClaims>(token, key, &validation) {
            Ok(token_data) => {
                return Ok(token_data.claims.sub);
            }
            Err(e) => {
                error!("failed to decode jwt: {e:?}");
            }
        }
    }

    if let Ok(td) = jsonwebtoken::dangerous::insecure_decode::<serde_json::Value>(token) {
        let kid = td.header.kid;
        error!(
            "Invalid Firebase Token signature or metadata mismatch, kid: {:?}, from {:?} ({:?})",
            kid.clone(),
            td.claims.get("name"),
            td.claims.get("email")
        );
        warn!("current token: {token}");
    }

    Err(Box::from(
        "Invalid Firebase Token signature or metadata mismatch",
    ))
}

pub(crate) fn execute_dh_handshake(
    client_pub_b64: &str,
) -> Result<(String, aes_gcm::Aes256Gcm), Box<dyn std::error::Error + Send + Sync>> {
    use aes_gcm::KeyInit;
    use base64::{Engine, engine::general_purpose::STANDARD as BASE64};
    use p256::{EncodedPoint, PublicKey, ecdh::EphemeralSecret};
    use rand_core::OsRng;

    // Step: decode client public key
    // info!("client_pub_b64: {client_pub_b64}");
    let client_bytes = BASE64.decode(client_pub_b64)?;
    let encoded_point = EncodedPoint::from_bytes(&client_bytes)?;
    let client_public = PublicKey::from_encoded_point(&encoded_point).unwrap();

    // Generate server ephemeral keypair

    let server_secret = EphemeralSecret::random(&mut OsRng);
    let server_public = PublicKey::from(&server_secret);

    // Compute symmetric shared secret
    let shared_secret = server_secret.diffie_hellman(&client_public);

    let secret_bytes = shared_secret.raw_secret_bytes();
    // Instantiate AES-256 GCM Core Cipher block natively using derived 32-byte hash block
    let cipher = aes_gcm::Aes256Gcm::new_from_slice(&secret_bytes)
        .map_err(|_| "failed allocating cipher payload context init")?;

    let server_pub_bytes = server_public.to_encoded_point(false);
    let server_public_b64 = BASE64.encode(server_pub_bytes.as_bytes());

    Ok((server_public_b64, cipher))
}

pub(crate) fn decrypt_message(
    cipher: &aes_gcm::Aes256Gcm,
    frame: &EncryptedFrame,
) -> Result<Vec<u8>, Box<dyn std::error::Error + Send + Sync>> {
    use aes_gcm::aead::Aead;
    use base64::{Engine, engine::general_purpose::STANDARD as BASE64};

    let iv_bytes = BASE64.decode(&frame.iv)?;
    let ciphertext_bytes = BASE64.decode(&frame.ciphertext)?;

    let nonce = aes_gcm::Nonce::from_slice(&iv_bytes);
    let decrypted = cipher
        .decrypt(nonce, ciphertext_bytes.as_slice())
        .map_err(|_| "Decryption routine validation assertion failed")?;

    Ok(decrypted)
}

pub(crate) fn encrypt_server_message(
    mut cipher: aes_gcm::Aes256Gcm,
    plain_text: &str,
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
    use base64::{Engine, engine::general_purpose::STANDARD as BASE64};

    let mut iv_bytes = [0u8; 12];
    rand::rng().fill_bytes(&mut iv_bytes);
    let nonce = aes_gcm::Nonce::from_slice(&iv_bytes);

    let ciphertext_bytes = cipher
        .encrypt(nonce, plain_text.as_bytes())
        .map_err(|_| "Encryption execution routine failed")?;

    let frame = EncryptedFrame {
        iv: BASE64.encode(iv_bytes),
        ciphertext: BASE64.encode(ciphertext_bytes),
    };

    let json_output = serde_json::to_string(&frame)?;

    Ok(json_output)
}