quicproquo/crates/quicprochat-server/src/domain/keys.rs

//! Key management domain logic — KeyPackage and hybrid key operations.

use std::sync::Arc;

use sha2::{Digest, Sha256};

use crate::storage::Store;

use super::types::*;

const MAX_KEYPACKAGE_BYTES: usize = 1024 * 1024; // 1 MB

/// Domain service for MLS KeyPackage and hybrid (PQ) key management.
pub struct KeyService {
    pub store: Arc<dyn Store>,
}

impl KeyService {
    /// Upload an MLS KeyPackage for the given identity. Returns a SHA-256 fingerprint.
    pub fn upload_key_package(
        &self,
        req: UploadKeyPackageReq,
        _auth: &CallerAuth,
    ) -> Result<UploadKeyPackageResp, DomainError> {
        if req.identity_key.len() != 32 {
            return Err(DomainError::InvalidIdentityKey(req.identity_key.len()));
        }
        if req.package.is_empty() {
            return Err(DomainError::EmptyPackage);
        }
        if req.package.len() > MAX_KEYPACKAGE_BYTES {
            return Err(DomainError::PackageTooLarge(req.package.len()));
        }

        let fingerprint: Vec<u8> = Sha256::digest(&req.package).to_vec();
        let start = std::time::Instant::now();
        self.store
            .upload_key_package(&req.identity_key, req.package)?;
        crate::metrics::record_storage_latency("key_package_upload", start.elapsed());

        Ok(UploadKeyPackageResp { fingerprint })
    }

    /// Fetch the stored MLS KeyPackage for an identity. Returns empty if none exists.
    pub fn fetch_key_package(
        &self,
        req: FetchKeyPackageReq,
        _auth: &CallerAuth,
    ) -> Result<FetchKeyPackageResp, DomainError> {
        let start = std::time::Instant::now();
        let package = self.store.fetch_key_package(&req.identity_key)?;
        crate::metrics::record_storage_latency("key_package_fetch", start.elapsed());
        Ok(FetchKeyPackageResp {
            package: package.unwrap_or_default(),
        })
    }

    /// Upload a hybrid (ML-KEM-768) public key for post-quantum key exchange.
    pub fn upload_hybrid_key(
        &self,
        req: UploadHybridKeyReq,
        _auth: &CallerAuth,
    ) -> Result<(), DomainError> {
        if req.identity_key.len() != 32 {
            return Err(DomainError::InvalidIdentityKey(req.identity_key.len()));
        }
        if req.hybrid_public_key.is_empty() {
            return Err(DomainError::EmptyHybridKey);
        }

        self.store
            .upload_hybrid_key(&req.identity_key, req.hybrid_public_key)?;
        Ok(())
    }

    /// Fetch the hybrid public key for a single identity. Returns empty if none exists.
    pub fn fetch_hybrid_key(
        &self,
        req: FetchHybridKeyReq,
        _auth: &CallerAuth,
    ) -> Result<FetchHybridKeyResp, DomainError> {
        let hybrid_public_key = self
            .store
            .fetch_hybrid_key(&req.identity_key)?
            .unwrap_or_default();
        Ok(FetchHybridKeyResp { hybrid_public_key })
    }

    /// Batch-fetch hybrid public keys for multiple identities. Missing keys return empty.
    pub fn fetch_hybrid_keys(
        &self,
        req: FetchHybridKeysReq,
        _auth: &CallerAuth,
    ) -> Result<FetchHybridKeysResp, DomainError> {
        let mut keys = Vec::with_capacity(req.identity_keys.len());
        for ik in &req.identity_keys {
            let pk = self.store.fetch_hybrid_key(ik)?.unwrap_or_default();
            keys.push(pk);
        }
        Ok(FetchHybridKeysResp { keys })
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use crate::storage::FileBackedStore;

    fn test_service() -> (tempfile::TempDir, KeyService) {
        let dir = tempfile::tempdir().unwrap();
        let store = Arc::new(FileBackedStore::open(dir.path()).unwrap());
        let svc = KeyService { store };
        (dir, svc)
    }

    fn test_auth() -> CallerAuth {
        CallerAuth {
            identity_key: vec![1u8; 32],
            token: vec![0u8; 16],
            device_id: None,
        }
    }

    #[test]
    fn upload_and_fetch_key_package() {
        let (_dir, svc) = test_service();
        let auth = test_auth();
        let ik = vec![1u8; 32];
        let package = vec![42u8; 128];

        let resp = svc
            .upload_key_package(
                UploadKeyPackageReq {
                    identity_key: ik.clone(),
                    package: package.clone(),
                },
                &auth,
            )
            .unwrap();
        // Fingerprint is SHA-256 of the package
        assert_eq!(resp.fingerprint.len(), 32);

        let fetched = svc
            .fetch_key_package(FetchKeyPackageReq { identity_key: ik }, &auth)
            .unwrap();
        assert_eq!(fetched.package, package);
    }

    #[test]
    fn fetch_key_package_missing() {
        let (_dir, svc) = test_service();
        let auth = test_auth();

        let resp = svc
            .fetch_key_package(
                FetchKeyPackageReq { identity_key: vec![99u8; 32] },
                &auth,
            )
            .unwrap();
        assert!(resp.package.is_empty());
    }

    #[test]
    fn upload_key_package_rejects_invalid_identity_key() {
        let (_dir, svc) = test_service();
        let auth = test_auth();

        let err = svc
            .upload_key_package(
                UploadKeyPackageReq {
                    identity_key: vec![1u8; 31],
                    package: vec![1u8; 10],
                },
                &auth,
            )
            .unwrap_err();
        assert!(matches!(err, DomainError::InvalidIdentityKey(31)));
    }

    #[test]
    fn upload_key_package_rejects_empty_package() {
        let (_dir, svc) = test_service();
        let auth = test_auth();

        let err = svc
            .upload_key_package(
                UploadKeyPackageReq {
                    identity_key: vec![1u8; 32],
                    package: vec![],
                },
                &auth,
            )
            .unwrap_err();
        assert!(matches!(err, DomainError::EmptyPackage));
    }

    #[test]
    fn upload_key_package_rejects_oversized() {
        let (_dir, svc) = test_service();
        let auth = test_auth();

        let err = svc
            .upload_key_package(
                UploadKeyPackageReq {
                    identity_key: vec![1u8; 32],
                    package: vec![0u8; MAX_KEYPACKAGE_BYTES + 1],
                },
                &auth,
            )
            .unwrap_err();
        assert!(matches!(err, DomainError::PackageTooLarge(_)));
    }

    #[test]
    fn upload_and_fetch_hybrid_key() {
        let (_dir, svc) = test_service();
        let auth = test_auth();
        let ik = vec![2u8; 32];
        let hk = vec![0xABu8; 1184]; // ML-KEM-768 public key size

        svc.upload_hybrid_key(
            UploadHybridKeyReq {
                identity_key: ik.clone(),
                hybrid_public_key: hk.clone(),
            },
            &auth,
        )
        .unwrap();

        let resp = svc
            .fetch_hybrid_key(FetchHybridKeyReq { identity_key: ik }, &auth)
            .unwrap();
        assert_eq!(resp.hybrid_public_key, hk);
    }

    #[test]
    fn fetch_hybrid_key_missing() {
        let (_dir, svc) = test_service();
        let auth = test_auth();

        let resp = svc
            .fetch_hybrid_key(
                FetchHybridKeyReq { identity_key: vec![99u8; 32] },
                &auth,
            )
            .unwrap();
        assert!(resp.hybrid_public_key.is_empty());
    }

    #[test]
    fn upload_hybrid_key_rejects_invalid_identity() {
        let (_dir, svc) = test_service();
        let auth = test_auth();

        let err = svc
            .upload_hybrid_key(
                UploadHybridKeyReq {
                    identity_key: vec![1u8; 10],
                    hybrid_public_key: vec![1u8; 100],
                },
                &auth,
            )
            .unwrap_err();
        assert!(matches!(err, DomainError::InvalidIdentityKey(10)));
    }

    #[test]
    fn upload_hybrid_key_rejects_empty() {
        let (_dir, svc) = test_service();
        let auth = test_auth();

        let err = svc
            .upload_hybrid_key(
                UploadHybridKeyReq {
                    identity_key: vec![1u8; 32],
                    hybrid_public_key: vec![],
                },
                &auth,
            )
            .unwrap_err();
        assert!(matches!(err, DomainError::EmptyHybridKey));
    }

    #[test]
    fn fetch_hybrid_keys_batch() {
        let (_dir, svc) = test_service();
        let auth = test_auth();
        let ik1 = vec![1u8; 32];
        let ik2 = vec![2u8; 32];
        let ik3 = vec![3u8; 32]; // no hybrid key uploaded

        svc.upload_hybrid_key(
            UploadHybridKeyReq {
                identity_key: ik1.clone(),
                hybrid_public_key: vec![0xAAu8; 64],
            },
            &auth,
        )
        .unwrap();
        svc.upload_hybrid_key(
            UploadHybridKeyReq {
                identity_key: ik2.clone(),
                hybrid_public_key: vec![0xBBu8; 64],
            },
            &auth,
        )
        .unwrap();

        let resp = svc
            .fetch_hybrid_keys(
                FetchHybridKeysReq {
                    identity_keys: vec![ik1, ik2, ik3],
                },
                &auth,
            )
            .unwrap();
        assert_eq!(resp.keys.len(), 3);
        assert_eq!(resp.keys[0], vec![0xAAu8; 64]);
        assert_eq!(resp.keys[1], vec![0xBBu8; 64]);
        assert!(resp.keys[2].is_empty()); // missing key returns empty
    }

    #[test]
    fn upload_key_package_at_max_size() {
        let (_dir, svc) = test_service();
        let auth = test_auth();

        // Exactly at max should succeed
        let resp = svc
            .upload_key_package(
                UploadKeyPackageReq {
                    identity_key: vec![1u8; 32],
                    package: vec![0u8; MAX_KEYPACKAGE_BYTES],
                },
                &auth,
            );
        assert!(resp.is_ok());
    }
}