use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use tempfile::tempdir;
use crate::{
api_keys::{ApiKeyGenerator, ApiKeyModel},
models::ServiceModel,
proxy::AuthenticatedProxy,
test_client::TestClient,
types::{
ChallengeRequest, ChallengeResponse, KeyType, ServiceId, VerifyChallengeResponse, headers,
},
};
#[tokio::test]
async fn test_api_key_rotation() {
let mut rng = blueprint_std::BlueprintRng::new();
let tmp = tempdir().unwrap();
let proxy = AuthenticatedProxy::new(tmp.path()).unwrap();
let db = proxy.db();
let service_id = ServiceId::new(1);
let mut service = ServiceModel {
api_key_prefix: "rot_".to_string(),
owners: Vec::new(),
upstream_url: "http://localhost:8080".to_string(),
tls_profile: None,
};
let signing_key = k256::ecdsa::SigningKey::random(&mut rng);
let public_key = signing_key.verifying_key().to_sec1_bytes();
service.add_owner(KeyType::Ecdsa, public_key.into());
service.save(service_id, &db).unwrap();
let key_gen = ApiKeyGenerator::with_prefix("rot_");
let expires_at = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs()
+ 3600;
let old_key = key_gen.generate_key(service_id, expires_at, BTreeMap::new(), &mut rng);
let mut old_model = ApiKeyModel::from(&old_key);
old_model.save(&db).unwrap();
let found = ApiKeyModel::find_by_key_id(&old_model.key_id, &db).unwrap();
assert!(found.is_some());
assert!(found.unwrap().validates_key(old_key.full_key()));
old_model.is_enabled = false;
old_model.save(&db).unwrap();
let new_expires_at = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs()
+ (90 * 24 * 3600);
let new_key = key_gen.generate_key(service_id, new_expires_at, BTreeMap::new(), &mut rng);
let mut new_model = ApiKeyModel::from(&new_key);
new_model.save(&db).unwrap();
let old_check = ApiKeyModel::find_by_key_id(&old_model.key_id, &db)
.unwrap()
.unwrap();
assert!(!old_check.is_enabled, "Old key should be disabled");
let new_check = ApiKeyModel::find_by_key_id(&new_model.key_id, &db)
.unwrap()
.unwrap();
assert!(new_check.is_enabled, "New key should be enabled");
assert!(new_check.validates_key(new_key.full_key()));
}
#[tokio::test]
async fn test_api_key_renewal() {
let mut rng = blueprint_std::BlueprintRng::new();
let tmp = tempdir().unwrap();
let proxy = AuthenticatedProxy::new(tmp.path()).unwrap();
let db = proxy.db();
let service_id = ServiceId::new(1);
let service = ServiceModel {
api_key_prefix: "ren_".to_string(),
owners: Vec::new(),
upstream_url: "http://localhost:8080".to_string(),
tls_profile: None,
};
service.save(service_id, &db).unwrap();
let key_gen = ApiKeyGenerator::with_prefix("ren_");
let short_expires = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs()
+ 300;
let api_key = key_gen.generate_key(service_id, short_expires, BTreeMap::new(), &mut rng);
let mut model = ApiKeyModel::from(&api_key);
model.save(&db).unwrap();
assert_eq!(model.expires_at, short_expires);
let new_expires = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs()
+ (30 * 24 * 3600);
model.expires_at = new_expires;
model.save(&db).unwrap();
let renewed = ApiKeyModel::find_by_key_id(&model.key_id, &db)
.unwrap()
.unwrap();
assert_eq!(renewed.expires_at, new_expires);
assert!(!renewed.is_expired());
}
#[tokio::test]
async fn test_api_key_revocation() {
let mut rng = blueprint_std::BlueprintRng::new();
let tmp = tempdir().unwrap();
let proxy = AuthenticatedProxy::new(tmp.path()).unwrap();
let db = proxy.db();
let service_id = ServiceId::new(1);
let mut service = ServiceModel {
api_key_prefix: "rev_".to_string(),
owners: Vec::new(),
upstream_url: "http://localhost:8080".to_string(),
tls_profile: None,
};
let signing_key = k256::ecdsa::SigningKey::random(&mut rng);
let public_key = signing_key.verifying_key().to_sec1_bytes();
service.add_owner(KeyType::Ecdsa, public_key.clone().into());
service.save(service_id, &db).unwrap();
let router = proxy.router();
let client = TestClient::new(router);
let challenge_req = ChallengeRequest {
pub_key: public_key.into(),
key_type: KeyType::Ecdsa,
};
let res = client
.post("/v1/auth/challenge")
.header(headers::X_SERVICE_ID, service_id.to_string())
.json(&challenge_req)
.await;
let challenge_res: ChallengeResponse = res.json().await;
let (signature, _) = signing_key
.sign_prehash_recoverable(&challenge_res.challenge)
.unwrap();
let verify_req = crate::types::VerifyChallengeRequest {
challenge: challenge_res.challenge,
signature: signature.to_bytes().into(),
challenge_request: challenge_req,
expires_at: 0,
additional_headers: BTreeMap::new(),
};
let res = client
.post("/v1/auth/verify")
.header(headers::X_SERVICE_ID, service_id.to_string())
.json(&verify_req)
.await;
let verify_res: VerifyChallengeResponse = res.json().await;
let api_key = match verify_res {
VerifyChallengeResponse::Verified { api_key, .. } => api_key,
_ => panic!("Expected verified response"),
};
let res = client
.get("/test")
.header(headers::AUTHORIZATION, format!("Bearer {api_key}"))
.await;
assert_ne!(res.status(), 401, "Key should be valid initially");
let key_id = api_key.split('.').next().unwrap();
let mut model = ApiKeyModel::find_by_key_id(key_id, &db).unwrap().unwrap();
model.is_enabled = false;
model.save(&db).unwrap();
let res = client
.get("/test")
.header(headers::AUTHORIZATION, format!("Bearer {api_key}"))
.await;
assert_eq!(res.status(), 401, "Revoked key should be rejected");
}
#[tokio::test]
async fn test_expired_api_key_rejection() {
let mut rng = blueprint_std::BlueprintRng::new();
let tmp = tempdir().unwrap();
let proxy = AuthenticatedProxy::new(tmp.path()).unwrap();
let db = proxy.db();
let service_id = ServiceId::new(1);
let service = ServiceModel {
api_key_prefix: "exp_".to_string(),
owners: Vec::new(),
upstream_url: "http://localhost:8080".to_string(),
tls_profile: None,
};
service.save(service_id, &db).unwrap();
let key_gen = ApiKeyGenerator::with_prefix("exp_");
let past_time = 1;
let api_key = key_gen.generate_key(service_id, past_time, BTreeMap::new(), &mut rng);
let mut model = ApiKeyModel::from(&api_key);
model.save(&db).unwrap();
assert!(model.is_expired(), "Key should be expired");
let router = proxy.router();
let client = TestClient::new(router);
let res = client
.get("/test")
.header(
headers::AUTHORIZATION,
format!("Bearer {}", api_key.full_key()),
)
.await;
assert_eq!(res.status(), 401, "Expired key should be rejected");
}
#[tokio::test]
async fn test_concurrent_api_key_operations() {
let tmp = tempdir().unwrap();
let proxy = AuthenticatedProxy::new(tmp.path()).unwrap();
let db = proxy.db();
let service_id = ServiceId::new(1);
let service = ServiceModel {
api_key_prefix: "con_".to_string(),
owners: Vec::new(),
upstream_url: "http://localhost:8080".to_string(),
tls_profile: None,
};
service.save(service_id, &db).unwrap();
let mut handles = Vec::new();
for i in 0..10 {
let db = db.clone();
let handle = tokio::spawn(async move {
let mut rng = blueprint_std::BlueprintRng::new();
let key_gen = ApiKeyGenerator::with_prefix(&format!("con{i}_"));
let expires = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs()
+ 3600;
let key = key_gen.generate_key(service_id, expires, BTreeMap::new(), &mut rng);
let mut model = ApiKeyModel::from(&key);
model.save(&db).unwrap();
(model.key_id.clone(), key.full_key().to_string())
});
handles.push(handle);
}
let mut results = Vec::new();
for handle in handles {
results.push(handle.await.unwrap());
}
assert_eq!(results.len(), 10);
let key_ids: std::collections::HashSet<_> = results.iter().map(|(id, _)| id).collect();
assert_eq!(key_ids.len(), 10, "All key IDs should be unique");
for (key_id, full_key) in results {
let model = ApiKeyModel::find_by_key_id(&key_id, &db).unwrap().unwrap();
assert!(model.validates_key(&full_key));
}
}