zoe_wire_protocol/

challenge.rs

use crate::{Signature, VerifyingKey};
use forward_compatible_enum::ForwardCompatibleEnum;
use serde::{Deserialize, Serialize};

#[cfg(feature = "server")]
pub mod server;

#[cfg(feature = "client")]
pub mod client;

/// Default challenge timeout in seconds
#[cfg(feature = "server")]
const DEFAULT_CHALLENGE_TIMEOUT_SECS: u64 = 30;

/// Maximum size for challenge messages (to prevent DoS)
#[cfg(any(feature = "client", feature = "server"))]
const MAX_PACKAGE_SIZE: usize = 1024 * 1024; // Should be enough for challenge data

/// Forward-compatible challenge system for connection-level authentication
///
/// The Zoe protocol uses a challenge-response handshake immediately after QUIC connection
/// establishment to verify possession of cryptographic private keys. This happens before
/// any service streams are created, ensuring all connections have verified credentials.
///
/// ## Protocol Flow
///
/// 1. **QUIC Connection**: Client connects using ML-DSA-44 mutual TLS
/// 2. **Challenge Phase**: Server sends `ZoeChallenge` on first bi-directional stream
/// 3. **Response Phase**: Client responds with `ZoeChallengeResponse`
/// 4. **Verification**: Server verifies proofs and sends `ZoeChallengeResult`
/// 5. **Service Phase**: Normal service streams can now be established
///
/// ## Wire Format
///
/// All challenge messages are serialized using postcard format for compact binary encoding.
///
/// ### Challenge Message (Server → Client)
/// ```text
/// | Field                | Type              | Description                    |
/// |---------------------|-------------------|--------------------------------|
/// | challenge_type      | u8                | Forward-compatible enum tag    |
/// | challenge_data      | Vec<u8>           | Serialized challenge content   |
/// ```
///
/// ### Response Message (Client → Server)  
/// ```text
/// | Field                | Type              | Description                    |
/// |---------------------|-------------------|--------------------------------|
/// | response_type       | u8                | Forward-compatible enum tag    |
/// | response_data       | Vec<u8>           | Serialized response content    |
/// ```
///
/// ### Result Message (Server → Client)
/// ```text
/// | Field                | Type              | Description                    |
/// |---------------------|-------------------|--------------------------------|
/// | result_type         | u8                | Forward-compatible enum tag    |
/// | result_data         | Vec<u8>           | Serialized result content      |
/// ```
///
/// ## Security Properties
///
/// - **Replay Protection**: Each challenge includes a unique nonce
/// - **Server Binding**: Signatures include server's public key
/// - **Time Bounds**: Challenges have expiration timestamps
/// - **Connection Scoped**: Verified keys are tied to specific QUIC connections
/// - **Forward Secrecy**: New challenges generated for each connection
///
/// ## Example Usage
///
/// ```rust
/// use zoe_wire_protocol::{ZoeChallenge, ZoeChallengeResponse, KeyChallenge};
///
/// // Server sends challenge
/// let challenge = ZoeChallenge::Key(KeyChallenge {
///     nonce: generate_nonce(),
///     server_public_key: server_key.to_bytes().to_vec(),
///     expires_at: current_time() + 30,
/// });
///
/// // Client creates multiple key proofs
/// let response = ZoeChallengeResponse::Key(KeyResponse {
///     key_proofs: vec![
///         KeyProof { public_key: key1_bytes, signature: sig1_bytes },
///         KeyProof { public_key: key2_bytes, signature: sig2_bytes },
///     ],
/// });
/// ```
#[derive(Debug, Clone, ForwardCompatibleEnum)]
pub enum ZoeChallenge {
    /// Multi-key ML-DSA challenge allowing clients to prove multiple private keys
    ///
    /// This challenge type allows clients to prove possession of multiple ML-DSA
    /// private keys in a single handshake round-trip. This is useful for:
    ///
    /// - **Role-based authentication**: Different keys for personal, work, admin roles
    /// - **Key rotation**: Proving both old and new keys during transition periods  
    /// - **Delegation**: Proving keys for multiple identities or organizations
    /// - **Batch verification**: Efficient verification of multiple keys at once
    ///
    /// The client must sign `(nonce || server_public_key)` with each private key
    /// they wish to prove possession of.
    #[discriminant(1)]
    Key(Box<KeyChallenge>),

    /// Unknown challenge type for forward compatibility
    Unknown { discriminant: u32, data: Vec<u8> },
}

#[derive(Debug, Clone, ForwardCompatibleEnum)]
pub enum ZoeChallengeRejection {
    /// The challenge has been rejected. The connection will be closed. A message
    /// might be provided to explain why and what to do before trying again.
    #[discriminant(30)]
    GenericRejection(Option<String>),

    /// The client failed to respond to the challenge with the valied answer.
    #[discriminant(31)]
    ChallengeFailed,

    /// The client failed to respond to the challenge on time
    #[discriminant(32)]
    ChallengeExpired,

    /// The client failed to respond to the challenge with a valid complete answer.
    #[discriminant(33)]
    ChallengeIncomplete,

    /// The client has been blocked by the server. Any fruther requests from this
    /// client will be rejected.
    #[discriminant(40)]
    Blocked,

    /// Unknown rejection type for forward compatibility
    Unknown { discriminant: u32, data: Vec<u8> },
}

#[derive(Debug, Clone, ForwardCompatibleEnum)]
pub enum ZoeChallengeWarning {
    /// The connection is throttled, a reason to display to the user might be provided
    #[discriminant(20)]
    ConnectionThrottled(String),

    /// The user is close to the quota limit on this server. This should probably be
    /// prompted to the user as uploading data might be rejected by the server.
    #[discriminant(33)]
    QuotaLimitInReach(String),

    /// The user is out of quota on this server. This should probably be prompted
    /// to the user as uploading data might be rejected by the server.
    #[discriminant(34)]
    OutOfQuota(String),

    /// The user has reached the limits (other than quota) on this server. This
    /// should probably be prompted to the user as uploading data might be rejected
    /// by the server.
    #[discriminant(35)]
    UserHasLimitedResource(String),

    /// The best ngeotiated protocol version is deprecated. The client should
    /// prompt the user about this server soon not supporting this client anymore
    /// and thus the client needs to be upgraded.
    #[discriminant(41)]
    ProtocolVersionDeprecated(String),

    /// The server has detected a potential breach of the security policy. The
    /// client should prompt the user about this and ask for confirmation to
    /// continue.
    #[discriminant(42)]
    PotentialBreach(String),

    /// The server has detected a weak crypto algorithm being used. The client
    /// should prompt the user about this and ask for confirmation to continue.
    #[discriminant(43)]
    WeakCrypto(String),

    /// THe server has a generic warning to show to the user. This should be shown
    /// to the user by the client.
    #[discriminant(50)]
    GenericServerWarning(String),

    /// The server is under high load and might be throttling the user connection,
    #[discriminant(51)]
    ServerUnderLoad(String),

    /// Unknown warning type for forward compatibility
    Unknown { discriminant: u32, data: Vec<u8> },
}

/// Forward-compatible result system for challenge verification
///
/// After verifying challenge responses, the server sends back results indicating
/// which proofs succeeded or failed. This allows clients to understand the
/// verification status and take appropriate action.
#[derive(Debug, Clone, ForwardCompatibleEnum)]
pub enum ZoeChallengeResult {
    /// The challanges have been accepted.
    #[discriminant(20)]
    Accepted,

    /// This challenge has been accepted, but there is another challenge to come
    /// and perform. After this read for another ZoeChallenge.
    #[discriminant(30)]
    Next,

    /// The server wants to warn the client and potentially the user about a
    /// possible issue. The client should consider the warning and maybe
    /// inform the user about it. But other than that it should continue by
    /// waiting for another result. Can be sent multiple times for different
    /// problems.
    #[discriminant(31)]
    Warning(ZoeChallengeWarning),

    /// The challenge has been rejected The connection will be closed. A message
    /// might be provided to explain why and what to do before trying again.
    #[discriminant(40)]
    Rejected(ZoeChallengeRejection),

    /// An error occured (on the server side). The connection will be closed.
    /// The error should probably be shown to the user to allow them to figure
    /// out what went wrong before trying again.
    #[discriminant(50)]
    Error(String),

    /// Unknown result type for forward compatibility
    Unknown { discriminant: u32, data: Vec<u8> },
}

/// Challenge for proving possession of multiple ML-DSA private keys
///
/// This challenge is sent by the server immediately after QUIC connection
/// establishment. The client must respond by proving possession of one or
/// more ML-DSA private keys.
///
/// ## Security Considerations
///
/// - **Nonce**: Must be cryptographically random and unique per challenge
/// - **Server Key**: Binds the signature to this specific server
/// - **Expiration**: Prevents replay attacks and limits challenge lifetime
/// - **Key Encoding**: ML-DSA public keys should use the standard encoding
///
/// ## Wire Size
///
/// Approximate serialized size: ~80 bytes
/// - nonce: 32 bytes
/// - server_public_key: ~1312 bytes (ML-DSA-44)  
/// - expires_at: 8 bytes
/// - overhead: ~8 bytes (postcard encoding)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeyChallenge {
    /// Cryptographically random nonce that must be included in signatures
    ///
    /// This 32-byte nonce provides replay protection by ensuring each challenge
    /// is unique. Clients must include this exact nonce when constructing
    /// their signature data.
    pub nonce: [u8; 32],

    /// Server's ML-DSA-44 public key that must be included in signatures
    ///
    /// Including the server's public key in the signature data prevents
    /// signature replay attacks across different servers. This should be
    /// the same ML-DSA-44 key used in the server's TLS certificate.
    pub signature: Signature,

    /// Unix timestamp when this challenge expires
    ///
    /// Challenges have a limited lifetime (typically 30-60 seconds) to prevent
    /// replay attacks. Clients must respond before this timestamp or the
    /// challenge will be rejected.
    pub expires_at: u64,
}

/// Response containing proofs of ML-DSA private key possession
///
/// The client responds to an `KeyChallenge` by providing one or more
/// key proofs. Each proof demonstrates possession of a specific ML-DSA private key.
///
/// ## Proof Construction
///
/// For each key the client wishes to prove:
///
/// 1. Construct signature data: `nonce || server_public_key`
/// 2. Sign the data using the ML-DSA private key
/// 3. Create a `KeyProof` with the public key and signature
///
/// ## Wire Size
///
/// Approximate serialized size per key proof: ~2500 bytes
/// - public_key: ~1312 bytes (ML-DSA-65 public key)
/// - signature: ~2420 bytes (ML-DSA-65 signature)
/// - overhead: ~8 bytes (postcard encoding)
///
/// Total message size scales linearly with number of keys being proven.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeyResponse {
    /// List of key proofs - one for each ML-DSA key being proven
    ///
    /// The client can prove multiple keys in a single response. Each proof
    /// is verified independently by the server. At least one proof must
    /// succeed for the handshake to complete successfully.
    ///
    /// ## Ordering
    ///
    /// The order of proofs in this vector corresponds to the indices used
    /// in `KeyResult.failed_indices` for error reporting.
    pub key_proofs: Vec<KeyProof>,
}

/// Cryptographic proof of ML-DSA private key possession
///
/// Each proof consists of a public key and a signature that demonstrates
/// the client possesses the corresponding private key. The signature is
/// computed over challenge-specific data to prevent replay attacks.
///
/// ## Verification Process
///
/// The server verifies each proof by:
///
/// 1. Decoding the ML-DSA public key from `public_key`
/// 2. Reconstructing signature data: `nonce || server_public_key`  
/// 3. Verifying the signature using the public key and signature data
/// 4. Adding successfully verified keys to the connection's verified set
///
/// ## Key Encoding
///
/// ML-DSA public keys must be encoded using the standard ML-DSA encoding:
/// - ML-DSA-44: 1312 bytes
/// - ML-DSA-65: 1952 bytes  
/// - ML-DSA-87: 2592 bytes
///
/// This implementation uses ML-DSA-65 (security level 3, ~192-bit security).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeyProof {
    /// Encoded ML-DSA public key being proven
    ///
    /// This should be the result of calling `verifying_key.encode()` on
    /// an ML-DSA verifying key. The encoding includes all necessary
    /// information to reconstruct the public key for verification.
    pub public_key: VerifyingKey,

    /// ML-DSA signature over (nonce || server_public_key)
    ///
    /// This signature proves possession of the private key corresponding
    /// to `public_key`. It must be computed over the exact concatenation
    /// of the challenge nonce and server public key.
    ///
    /// Signature sizes:
    /// - ML-DSA-44: ~2420 bytes
    /// - ML-DSA-65: ~3309 bytes
    /// - ML-DSA-87: ~4627 bytes
    pub signature: Signature,
}

/// Result of ML-DSA multi-key challenge verification
///
/// After verifying all key proofs in a response, the server sends back
/// this result indicating which proofs succeeded or failed. This allows
/// the client to understand their verification status.
///
/// ## Success Criteria
///
/// The handshake is considered successful if at least one key proof is valid.
/// Even if some proofs fail, the connection can continue with the successfully
/// verified keys.
///
/// ## Error Handling
///
/// If all proofs fail, the server should close the connection. Clients can
/// use the failure information to:
/// - Log which specific keys were rejected
/// - Retry the connection with different keys
/// - Debug key or signature generation issues
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum KeyResult {
    /// All key proofs were successfully verified
    ///
    /// This is the ideal case where every key the client attempted to prove
    /// was successfully verified. All provided keys are now available for
    /// use in message authentication on this connection.
    AllValid,

    /// Some key proofs failed verification
    ///
    /// Contains the indices (into the original `key_proofs` vector) of proofs
    /// that failed verification. The connection continues with the successfully
    /// verified keys.
    ///
    /// Common failure reasons:
    /// - Invalid signature (wrong private key used)
    /// - Malformed public key encoding
    /// - Signature over wrong data (incorrect nonce/server key)
    /// - Expired challenge (client took too long to respond)
    PartialFailure {
        /// Zero-based indices of failed key proofs
        ///
        /// These indices correspond to positions in the original
        /// `KeyResponse.key_proofs` vector that failed verification.
        failed_indices: Vec<usize>,
    },

    /// All key proofs failed verification
    ///
    /// No keys were successfully verified. The server will typically close
    /// the connection after sending this result. Clients should not attempt
    /// to establish service streams after receiving this result.
    AllFailed,
}

impl KeyResult {
    /// Check if the handshake was successful (at least one key verified)
    ///
    /// Returns `true` if at least one key was successfully verified,
    /// `false` if all keys failed verification.
    ///
    /// # Example
    ///
    /// ```rust
    /// use zoe_wire_protocol::KeyResult;
    ///
    /// let result = KeyResult::PartialFailure {
    ///     failed_indices: vec![1, 3]
    /// };
    /// assert!(result.is_successful());
    ///
    /// let result = KeyResult::AllFailed;
    /// assert!(!result.is_successful());
    /// ```
    pub fn is_successful(&self) -> bool {
        !matches!(self, KeyResult::AllFailed)
    }

    /// Get the number of failed key proofs
    ///
    /// Returns the count of key proofs that failed verification.
    /// For `AllValid`, returns 0. For `AllFailed`, the count depends
    /// on how many keys were originally submitted.
    ///
    /// # Parameters
    ///
    /// * `total_keys` - Total number of keys that were submitted (needed for AllFailed case)
    ///
    /// # Example
    ///
    /// ```rust
    /// use zoe_wire_protocol::KeyResult;
    ///
    /// let result = KeyResult::PartialFailure {
    ///     failed_indices: vec![1, 3]
    /// };
    /// assert_eq!(result.failed_count(5), 2);
    ///
    /// let result = KeyResult::AllFailed;
    /// assert_eq!(result.failed_count(3), 3);
    /// ```
    pub fn failed_count(&self, total_keys: usize) -> usize {
        match self {
            KeyResult::AllValid => 0,
            KeyResult::PartialFailure { failed_indices } => failed_indices.len(),
            KeyResult::AllFailed => total_keys,
        }
    }

    /// Get the number of successfully verified keys
    ///
    /// Returns the count of key proofs that passed verification.
    ///
    /// # Parameters
    ///
    /// * `total_keys` - Total number of keys that were submitted
    ///
    /// # Example
    ///
    /// ```rust
    /// use zoe_wire_protocol::KeyResult;
    ///
    /// let result = KeyResult::PartialFailure {
    ///     failed_indices: vec![1]
    /// };
    /// assert_eq!(result.success_count(3), 2);
    ///
    /// let result = KeyResult::AllValid;
    /// assert_eq!(result.success_count(5), 5);
    /// ```
    pub fn success_count(&self, total_keys: usize) -> usize {
        total_keys - self.failed_count(total_keys)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ConnectionInfo;
    use std::collections::HashSet;

    #[test]
    fn test_ml_dsa_result_success_check() {
        assert!(KeyResult::AllValid.is_successful());
        assert!(KeyResult::PartialFailure {
            failed_indices: vec![1]
        }
        .is_successful());
        assert!(!KeyResult::AllFailed.is_successful());
    }

    #[test]
    fn test_ml_dsa_result_counts() {
        let result = KeyResult::PartialFailure {
            failed_indices: vec![0, 2],
        };
        assert_eq!(result.failed_count(5), 2);
        assert_eq!(result.success_count(5), 3);

        let result = KeyResult::AllValid;
        assert_eq!(result.failed_count(3), 0);
        assert_eq!(result.success_count(3), 3);

        let result = KeyResult::AllFailed;
        assert_eq!(result.failed_count(4), 4);
        assert_eq!(result.success_count(4), 0);
    }

    #[test]
    fn test_connection_info_key_verification() {
        let keypair1 = {
            use crate::KeyPair;
            KeyPair::generate_ed25519(&mut rand::thread_rng())
        };
        let keypair2 = {
            use crate::KeyPair;
            KeyPair::generate_ml_dsa65(&mut rand::thread_rng())
        };

        let mut verified_keys = HashSet::new();
        verified_keys.insert(keypair1.public_key());
        verified_keys.insert(keypair2.public_key());

        let transport_key = {
            use crate::KeyPair;
            KeyPair::generate(&mut rand::thread_rng()).public_key()
        };
        let connection_info = ConnectionInfo::with_verified_keys(
            transport_key,
            verified_keys,
            "127.0.0.1:8080".parse().unwrap(),
        );

        assert!(connection_info.has_verified_key(&keypair1.public_key()));
        assert!(connection_info.has_verified_key(&keypair2.public_key()));

        assert_eq!(connection_info.verified_key_count(), 2);
    }
}