use aes_gcm::{Aes256Gcm, KeyInit as _};
use sha2::{Digest as _, Sha256};
use std::path::Path;
use base64::Engine as _;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use knot_types::AccountDid;
use crate::config::ValidatedConfig;
use crate::datastore::{SealedStore, schema};
use super::crypto::{KeyPair, derive_gcm_key, dpop_proof, now_secs, seal_bytes, unseal_bytes};
use super::{GRANULAR_SCOPE, StoredSession, client_key_in, save_session_under, stored_session};
pub(crate) fn seed_test_session(
config: &ValidatedConfig,
owner: &AccountDid,
pds_url: &str,
) -> anyhow::Result<()> {
let store = SealedStore::open(config);
let cipher = derive_gcm_key(config)?;
let key = KeyPair::generate();
let session = StoredSession {
owner_did: owner.clone(),
dpop_scalar: key.to_bytes(),
access_token: "test-access-token".to_string(),
refresh_token: "test-refresh-token".to_string(),
scope: GRANULAR_SCOPE.to_string(),
token_endpoint: format!("{}/oauth/token", pds_url.trim_end_matches('/')),
pds_url: pds_url.to_string(),
expires_at: now_secs() + 3600,
dpop_nonce: None,
};
let record = store.record::<schema::OauthSession>(owner);
save_session_under(&record.guard()?, &cipher, &session)
}
fn payload_of(proof: &str) -> serde_json::Value {
let parts: Vec<&str> = proof.split('.').collect();
let payload = URL_SAFE_NO_PAD.decode(parts[1]).expect("payload base64");
serde_json::from_slice(&payload).expect("payload json")
}
#[test]
fn resource_proof_carries_ath_matching_the_access_token() {
let key = KeyPair::generate();
let token = "knot-static-test-access-token";
let unsigned = payload_of(
&dpop_proof(
&key,
"POST",
"https://pds.test/xrpc/com.atproto.repo.putRecord",
None,
Some(token),
)
.unwrap(),
);
let expected = URL_SAFE_NO_PAD.encode(Sha256::digest(token.as_bytes()));
assert_eq!(unsigned["ath"], serde_json::Value::String(expected));
assert_eq!(unsigned["htm"], "POST");
assert!(unsigned["jti"].as_str().is_some());
let endpoint =
payload_of(&dpop_proof(&key, "POST", "https://pds.test/oauth/par", None, None).unwrap());
assert!(endpoint.get("ath").is_none());
}
#[test]
fn seal_unseal_roundtrip() {
let key = [0x42u8; 32];
let cipher = Aes256Gcm::new_from_slice(&key).unwrap();
let plaintext = b"testing oauth secrets sealing";
let sealed = seal_bytes(&cipher, plaintext).unwrap();
let unsealed = unseal_bytes(&cipher, &sealed).unwrap();
assert_eq!(unsealed, plaintext);
}
fn session(owner: &str, token: &str) -> StoredSession {
StoredSession {
owner_did: AccountDid::new(owner.to_string()).unwrap(),
dpop_scalar: vec![7; 32],
access_token: token.to_string(),
refresh_token: token.to_string(),
scope: GRANULAR_SCOPE.to_string(),
token_endpoint: "https://pds.test/oauth/token".to_string(),
pds_url: "https://pds.test".to_string(),
expires_at: 0,
dpop_nonce: None,
}
}
#[test]
fn concurrent_session_writers_leave_a_session_that_decrypts() {
let dir = tempfile::tempdir().unwrap();
let store = SealedStore::at(dir.path());
let owner = AccountDid::new("did:plc:owner".to_string()).unwrap();
let cipher = Aes256Gcm::new_from_slice(&[0x33u8; 32]).unwrap();
let shared = dir.path().join("sessions/did_plc_owner.tmp");
std::fs::create_dir_all(&shared).unwrap();
let barrier = std::sync::Arc::new(std::sync::Barrier::new(4));
let threads: Vec<_> = (0..4)
.map(|index| {
let barrier = std::sync::Arc::clone(&barrier);
let cipher = cipher.clone();
let store = SealedStore::at(dir.path());
std::thread::spawn(move || {
let session = session("did:plc:owner", &format!("token-{index}").repeat(64));
let record = store.record::<schema::OauthSession>(&session.owner_did);
barrier.wait();
let guard = record.guard().unwrap();
let stored = stored_session(&guard, &cipher).unwrap();
assert!(stored.is_none_or(|stored| stored.scope == session.scope));
save_session_under(&guard, &cipher, &session).unwrap();
})
})
.collect();
for thread in threads {
thread.join().unwrap();
}
let record = store.record::<schema::OauthSession>(&owner);
let stored = stored_session(&record.guard().unwrap(), &cipher)
.expect("the session must decrypt")
.expect("the session must be there");
assert!(stored.access_token.starts_with("token-"));
}
#[test]
fn sealed_records_are_published_readable_by_the_knot_alone() {
use std::os::unix::fs::PermissionsExt as _;
let dir = tempfile::tempdir().unwrap();
let store = SealedStore::at(&dir.path().join("oauth"));
let cipher = Aes256Gcm::new_from_slice(&[0x44u8; 32]).unwrap();
client_key_in(&store, &cipher).unwrap();
let session = session("did:plc:owner", "token");
let record = store.record::<schema::OauthSession>(&session.owner_did);
save_session_under(&record.guard().unwrap(), &cipher, &session).unwrap();
let mode = |path: &Path| std::fs::metadata(path).unwrap().permissions().mode() & 0o777;
let oauth = dir.path().join("oauth");
for (path, expected) in [
(oauth.clone(), 0o700),
(oauth.join("client_key.sealed"), 0o600),
(oauth.join("sessions"), 0o700),
(record.path().to_path_buf(), 0o600),
] {
assert_eq!(mode(&path), expected, "{}", path.display());
}
for entry in std::fs::read_dir(oauth.join("sessions")).unwrap() {
let path = entry.unwrap().path();
assert_eq!(mode(&path), 0o600, "{}", path.display());
}
}
#[test]
fn every_process_that_wants_a_client_key_ends_up_with_the_published_one() {
let dir = tempfile::tempdir().unwrap();
let cipher = Aes256Gcm::new_from_slice(&[0x11u8; 32]).unwrap();
let barrier = std::sync::Arc::new(std::sync::Barrier::new(4));
let threads: Vec<_> = (0..4)
.map(|_| {
let barrier = std::sync::Arc::clone(&barrier);
let cipher = cipher.clone();
let store = SealedStore::at(dir.path());
std::thread::spawn(move || {
barrier.wait();
client_key_in(&store, &cipher).unwrap().to_bytes()
})
})
.collect();
let keys: Vec<_> = threads
.into_iter()
.map(|thread| thread.join().unwrap())
.collect();
let sealed = std::fs::read(dir.path().join("client_key.sealed")).unwrap();
let on_disk = KeyPair::from_bytes(&unseal_bytes(&cipher, &sealed).unwrap())
.unwrap()
.to_bytes();
for key in &keys {
assert_eq!(key, &on_disk, "a key no pds knows about would sign nothing");
}
}
#[test]
fn keypair_bytes_roundtrip() {
let key = KeyPair::generate();
let bytes = key.to_bytes();
let restored = KeyPair::from_bytes(&bytes).unwrap();
assert_eq!(restored.to_bytes(), bytes);
}