darkwing/server/utils/
jwt_utils.rs

use std::sync::Arc;

use anyhow::Context;
use jsonwebtoken::{decode, DecodingKey, Validation};
use serde::{Deserialize, Serialize};
use tracing::error;

use crate::config::DarkwingConfig;
use crate::server::error::{AppResult, Error};

/// A security service for handling JWT authentication.
pub type DynJwtUtil = Arc<dyn JwtUtil + Send + Sync>;

pub trait JwtUtil {
  fn get_user_id_from_token(&self, token: String) -> AppResult<u64>;
}

/// Our claims struct, it needs to derive `Serialize` and/or `Deserialize`
#[derive(Debug, Serialize, Deserialize)]
struct AccessTokenClaims {
  sub: String,
  exp: f32,
  iat: f32,
  aud: String,
}

pub struct JwtTokenUtil {
  config: Arc<DarkwingConfig>,
}

impl JwtTokenUtil {
  pub fn new(config: Arc<DarkwingConfig>) -> Self {
    Self { config }
  }
}

impl JwtUtil for JwtTokenUtil {
  fn get_user_id_from_token(&self, token: String) -> AppResult<u64> {
    let public_key =
      DecodingKey::from_rsa_pem(self.config.jwt_token_public_key.as_bytes())
        .context(
          "Failed to decode JWT token in JwtTokenUtil::get_user_id_from_token",
        )?;

    let mut validation =
      Validation::new(self.config.get_jwt_token_algorithm()?);
    validation.validate_aud = false;

    let decoded_token =
      decode::<AccessTokenClaims>(token.as_str(), &public_key, &validation)
        .map_err(|err| {
          Error::InternalServerErrorWithContext(err.to_string())
        })?;

    let sub = decoded_token.claims.sub.parse::<u64>().map_err(|err| {
      error!(
        "Failed to parse user ID from JWT token in JwtTokenUtil::get_user_id_from_token: {}",
        err
      );
      Error::InternalServerErrorWithContext(err.to_string())
    })?;

    Ok(sub)
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use base64::{prelude::BASE64_STANDARD, Engine as _};
  use jsonwebtoken::{encode, EncodingKey, Header};
  use std::time::{SystemTime, UNIX_EPOCH};

  fn create_test_config() -> (Arc<DarkwingConfig>, &'static str) {
    // Test RSA keys - these are example keys, never use in production
    let private_key = r#"-----BEGIN PRIVATE KEY-----
MIIEvwIBADANBgkqhkiG9w0BAQEFAASCBKkwggSlAgEAAoIBAQC7VJTUt9Us8cKj
MzEfYyjiWA4R4/M2bS1GB4t7NXp98C3SC6dVMvDuictGeurT8jNbvJZHtCSuYEvu
NMoSfm76oqFvAp8Gy0iz5sxjZmSnXyCdPEovGhLa0VzMaQ8s+CLOyS56YyCFGeJZ
qgtzJ6GR3eqoYSW9b9UMvkBpZODSctWSNGj3P7jRFDO5VoTwCQAWbFnOjDfH5Ulg
p2PKSQnSJP3AJLQNFNe7br1XbrhV//eO+t51mIpGSDCUv3E0DDFcWDTH9cXDTTlR
ZVEiR2BwpZOOkE/Z0/BVnhZYL71oZV34bKfWjQIt6V/isSMahdsAASACp4ZTGtwi
VuNd9tybAgMBAAECggEBAKTmjaS6tkK8BlPXClTQ2vpz/N6uxDeS35mXpqasqskV
laAidgg/sWqpjXDbXr93otIMLlWsM+X0CqMDgSXKejLS2jx4GDjI1ZTXg++0AMJ8
sJ74pWzVDOfmCEQ/7wXs3+cbnXhKriO8Z036q92Qc1+N87SI38nkGa0ABH9CN83H
mQqt4fB7UdHzuIRe/me2PGhIq5ZBzj6h3BpoPGzEP+x3l9YmK8t/1cN0pqI+dQwY
dgfGjackLu/2qH80MCF7IyQaseZUOJyKrCLtSD/Iixv/hzDEUPfOCjFDgTpzf3cw
ta8+oE4wHCo1iI1/4TlPkwmXx4qSXtmw4aQPz7IDQvECgYEA8KNThCO2gsC2I9PQ
DM/8Cw0O983WCDY+oi+7JPiNAJwv5DYBqEZB1QYdj06YD16XlC/HAZMsMku1na2T
N0driwenQQWzoev3g2S7gRDoS/FCJSI3jJ+kjgtaA7Qmzlgk1TxODN+G1H91HW7t
0l7VnL27IWyYo2qRRK3jzxqUiPUCgYEAx0oQs2reBQGMVZnApD1jeq7n4MvNLcPv
t8b/eU9iUv6Y4Mj0Suo/AU8lYZXm8ubbqAlwz2VSVunD2tOplHyMUrtCtObAfVDU
AhCndKaA9gApgfb3xw1IKbuQ1u4IF1FJl3VtumfQn//LiH1B3rXhcdyo3/vIttEk
48RakUKClU8CgYEAzV7W3COOlDDcQd935DdtKBFRAPRPAlspQUnzMi5eSHMD/ISL
DY5IiQHbIH83D4bvXq0X7qQoSBSNP7Dvv3HYuqMhf0DaegrlBuJllFVVq9qPVRnK
xt1Il2HgxOBvbhOT+9in1BzA+YJ99UzC85O0Qz06A+CmtHEy4aZ2kj5hHjECgYEA
mNS4+A8Fkss8Js1RieK2LniBxMgmYml3pfVLKGnzmng7H2+cwPLhPIzIuwytXywh
2bzbsYEfYx3EoEVgMEpPhoarQnYPukrJO4gwE2o5Te6T5mJSZGlQJQj9q4ZB2Dfz
et6INsK0oG8XVGXSpQvQh3RUYekCZQkBBFcpqWpbIEsCgYAnM3DQf3FJoSnXaMhr
VBIovic5l0xFkEHskAjFTevO86Fsz1C2aSeRKSqGFoOQ0tmJzBEs1R6KqnHInicD
TQrKhArgLXX4v3CddjfTRJkFWDbE/CkvKZNOrcf1nhaGCPspRJj2KUkj1Fhl9Cnc
dn/RsYEONbwQSjIfMPkvxF+8HQ==
-----END PRIVATE KEY-----"#;

    let public_key = r#"-----BEGIN PUBLIC KEY-----
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAu1SU1LfVLPHCozMxH2Mo
4lgOEePzNm0tRgeLezV6ffAt0gunVTLw7onLRnrq0/IzW7yWR7QkrmBL7jTKEn5u
+qKhbwKfBstIs+bMY2Zkp18gnTxKLxoS2tFczGkPLPgizskuemMghRniWaoLcyeh
kd3qqGElvW/VDL5AaWTg0nLVkjRo9z+40RQzuVaE8AkAFmxZzow3x+VJYKdjykkJ
0iT9wCS0DRTXu269V264Vf/3jvredZiKRkgwlL9xNAwxXFg0x/XFw005UWVRIkdg
cKWTjpBP2dPwVZ4WWC+9aGVd+Gyn1o0CLelf4rEjGoXbAAEgAqeGUxrcIlbjXfbc
mwIDAQAB
-----END PUBLIC KEY-----"#;

    (
      Arc::new(DarkwingConfig {
        jwt_token_public_key: public_key.to_string(),
        jwt_token_algorithm: "RS256".to_string(),
        ..Default::default() // ... other fields can be filled with dummy values
      }),
      private_key,
    )
  }

  fn create_test_token(user_id: u64, expires_in_secs: i64) -> String {
    let (_config, private_key) = create_test_config();
    let now = SystemTime::now()
      .duration_since(UNIX_EPOCH)
      .expect("Time went backwards")
      .as_secs() as f32;

    let claims = AccessTokenClaims {
      sub: user_id.to_string(),
      exp: now + expires_in_secs as f32,
      iat: now,
      aud: "test".to_string(),
    };

    let encoding_key = EncodingKey::from_rsa_pem(private_key.as_bytes())
      .expect("Failed to create encoding key");

    encode(
      &Header::new(jsonwebtoken::Algorithm::RS256),
      &claims,
      &encoding_key,
    )
    .expect("Failed to create test token")
  }

  #[test]
  fn test_valid_token() -> AppResult<()> {
    let (config, _) = create_test_config();
    let jwt_util = JwtTokenUtil::new(config);
    let token = create_test_token(123, 3600);

    let user_id = jwt_util.get_user_id_from_token(token)?;
    assert_eq!(user_id, 123);
    Ok(())
  }

  #[test]
  fn test_expired_token() {
    let (config, _) = create_test_config();
    let jwt_util = JwtTokenUtil::new(config);
    let token = create_test_token(123, -3600); // expired 1 hour ago

    let result = jwt_util.get_user_id_from_token(token);
    assert!(result.is_err());
    if let Err(Error::InternalServerErrorWithContext(msg)) = result {
      assert!(msg.contains("ExpiredSignature"));
    } else {
      panic!("Expected InternalServerErrorWithContext");
    }
  }

  #[test]
  fn test_invalid_token_format() {
    let (config, _) = create_test_config();
    let jwt_util = JwtTokenUtil::new(config);
    let wrong_token = "invalid.token.format".to_string();
    let wrong_token = BASE64_STANDARD.encode(wrong_token);
    let result = jwt_util.get_user_id_from_token(wrong_token);

    assert!(result.is_err());
    if let Err(Error::InternalServerErrorWithContext(msg)) = result {
      println!("msg: {}", msg);
      assert!(msg.contains("InvalidToken"));
    } else {
      panic!("Expected InternalServerErrorWithContext");
    }
  }

  #[test]
  fn test_malformed_user_id() {
    let (config, private_key) = create_test_config();
    let jwt_util = JwtTokenUtil::new(config.clone());

    // Create a token with invalid user ID format
    let now = SystemTime::now()
      .duration_since(UNIX_EPOCH)
      .expect("Time went backwards")
      .as_secs() as f32;

    let claims = AccessTokenClaims {
      sub: "not_a_number".to_string(),
      exp: now + 3600.0,
      iat: now,
      aud: "test".to_string(),
    };

    let encoding_key = EncodingKey::from_rsa_pem(private_key.as_bytes())
      .expect("Failed to create encoding key");

    let token = encode(
      &Header::new(jsonwebtoken::Algorithm::RS256),
      &claims,
      &encoding_key,
    )
    .expect("Failed to create test token");

    let result = jwt_util.get_user_id_from_token(token);
    assert!(result.is_err());
    if let Err(Error::InternalServerErrorWithContext(msg)) = result {
      assert!(msg.contains("invalid digit"));
    } else {
      panic!("Expected InternalServerErrorWithContext");
    }
  }
}
darkwing/server/utils/jwt_utils.rs

darkwing/server/utils/
jwt_utils.rs
