shortcake 0.1.0-pre.4

// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed
// licenses.

//! Responder protocol implementation.
//!
//! The Responder responds to the Initiator's first message by encapsulating
//! to their public key and sending back a ciphertext and nonce.

use core::marker::PhantomData;

use digest::Output;
use rand_core::CryptoRng;
use zeroize::Zeroize;

use crate::ciphersuite::{CipherSuite, Kem};
use crate::commitment;
use crate::error::Error;
use crate::initiator::{MessageOne, MessageThree};
use crate::sas::{compute_sas, derive_session_key};
use crate::verification::ProtocolOutput;
use crate::Nonce;

/// The second protocol message (Responder -> Initiator).
#[derive(Clone)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Serialize, serde::Deserialize),
    serde(bound(
        serialize = "<CS::Kem as Kem>::Ciphertext: serde::Serialize",
        deserialize = "<CS::Kem as Kem>::Ciphertext: serde::Deserialize<'de>",
    ))
)]
pub struct MessageTwo<CS: CipherSuite> {
    /// The ciphertext from KEM encapsulation.
    pub(crate) ct: <CS::Kem as Kem>::Ciphertext,
    /// The Responder's nonce.
    pub(crate) responder_nonce: Nonce,
}

/// Responder state in the 3-move SAS protocol.
///
/// Created by [`Responder::start`] upon receiving the initiator's first
/// message. Call [`Responder::finish`] after receiving the initiator's
/// final message to obtain a [`ProtocolOutput`].
#[cfg_attr(
    feature = "serde",
    derive(serde::Serialize, serde::Deserialize),
    serde(bound(
        serialize = "<CS::Kem as Kem>::EncapsulationKey: serde::Serialize, <CS::Kem as Kem>::Ciphertext: serde::Serialize, <CS::Kem as Kem>::SharedSecret: serde::Serialize",
        deserialize = "<CS::Kem as Kem>::EncapsulationKey: serde::Deserialize<'de>, <CS::Kem as Kem>::Ciphertext: serde::Deserialize<'de>, <CS::Kem as Kem>::SharedSecret: serde::Deserialize<'de>",
    ))
)]
pub struct Responder<CS: CipherSuite> {
    ek: <CS::Kem as Kem>::EncapsulationKey,
    commitment: Output<CS::Hash>,
    responder_nonce: Nonce,
    ct: <CS::Kem as Kem>::Ciphertext,
    /// Wrapped in `Option` so the consuming method can `.take()` the value
    /// before `self` is dropped (Drop still zeroizes if present).
    #[cfg_attr(feature = "serde", serde(deserialize_with = "require_some"))]
    shared_secret: Option<<CS::Kem as Kem>::SharedSecret>,
    _marker: PhantomData<CS>,
}

impl<CS: CipherSuite> Drop for Responder<CS> {
    fn drop(&mut self) {
        self.responder_nonce.zeroize();
        self.ek.zeroize();
        self.ct.zeroize();
        self.commitment.as_mut_slice().zeroize();
        if let Some(ref mut ss) = self.shared_secret {
            ss.zeroize();
        }
    }
}

impl<CS: CipherSuite> Responder<CS> {
    /// Start the protocol as Responder upon receiving the Initiator's first message.
    ///
    /// # Arguments
    ///
    /// * `rng` - A cryptographically secure random number generator.
    /// * `msg1` - The first protocol message from the Initiator.
    ///
    /// # Returns
    ///
    /// A tuple of (responder_state, second_message) on success.
    pub fn start(
        rng: &mut impl CryptoRng,
        msg1: MessageOne<CS>,
    ) -> Result<(Self, MessageTwo<CS>), Error> {
        // Encapsulate to Initiator's public key
        let (ct, shared_secret) =
            CS::Kem::encaps(&msg1.ek, rng).map_err(|_| Error::EncapsulationFailed)?;

        // Generate Responder's nonce
        let mut responder_nonce = [0u8; 32];
        rng.fill_bytes(&mut responder_nonce);

        let state = Self {
            ek: msg1.ek,
            commitment: msg1.commitment,
            responder_nonce,
            ct: ct.clone(),
            shared_secret: Some(shared_secret),
            _marker: PhantomData,
        };

        let message = MessageTwo {
            ct,
            responder_nonce,
        };

        Ok((state, message))
    }

    /// Process the initiator's final message and produce the protocol output.
    ///
    /// This verifies the commitment and computes the SAS.
    ///
    /// # Arguments
    ///
    /// * `msg3` - The third protocol message from the Initiator.
    ///
    /// # Returns
    ///
    /// A [`ProtocolOutput`] on success.
    pub fn finish(mut self, msg3: MessageThree) -> Result<ProtocolOutput<CS>, Error> {
        // Verify commitment
        commitment::open::<CS::Hash>(self.ek.as_ref(), &msg3.initiator_nonce, &self.commitment)?;

        // Compute SAS
        let sas = compute_sas::<CS::Hash>(
            &self.responder_nonce,
            &msg3.initiator_nonce,
            self.ct.as_ref(),
        );

        let mut kem_ss = self
            .shared_secret
            .take()
            .expect("shared_secret should always be Some");

        let session_key = derive_session_key::<CS::Hash>(
            self.ek.as_ref(),
            self.ct.as_ref(),
            &self.responder_nonce,
            &msg3.initiator_nonce,
            kem_ss.as_ref(),
        );
        kem_ss.zeroize();

        Ok(ProtocolOutput {
            sas,
            session_key,
            _marker: PhantomData,
        })
    }
}

#[cfg(feature = "serde")]
fn require_some<'de, D, T>(deserializer: D) -> Result<Option<T>, D::Error>
where
    D: serde::Deserializer<'de>,
    T: serde::Deserialize<'de>,
{
    use serde::Deserialize;
    let val = Option::<T>::deserialize(deserializer)?;
    if val.is_none() {
        return Err(serde::de::Error::custom("shared_secret must not be None"));
    }
    Ok(val)
}