update psq implementation

This commit is contained in:
Georgio Nicolas
2025-08-08 11:07:35 +02:00
parent 481ad79376
commit 84fcf81852
5 changed files with 104 additions and 72 deletions
+6 -14
View File
@@ -6,25 +6,17 @@ edition = "2021"
license.workspace = true
[dependencies]
arc-swap = { workspace = true }
bytes = { workspace = true }
futures = { workspace = true }
tracing = { workspace = true }
pin-project = { workspace = true }
sha2 = { workspace = true }
snow = { workspace = true }
strum = { workspace = true, features = ["derive"] }
thiserror = { workspace = true }
tokio = { workspace = true, features = ["full"] }
tokio-util = { workspace = true, features = ["codec"] }
# internal
nym-crypto = { path = "../crypto", features = ["asymmetric", "serde"]}
libcrux-traits = { git = "https://github.com/cryspen/libcrux" }
libcrux-kem = { git = "https://github.com/cryspen/libcrux" }
libcrux-psq = { git = "https://github.com/cryspen/libcrux", features = ["test-utils"] }
rand = "0.9.1"
libcrux-traits = { git = "https://github.com/georgio/libcrux" }
libcrux-kem = { git = "https://github.com/georgio/libcrux" }
libcrux-psq = { git = "https://github.com/georgio/libcrux", features = ["test-utils"] }
libcrux-ed25519 = { git = "https://github.com/georgio/libcrux" }
rand = "0.9.2"
tls_codec = { version = "0.4.2", features = ["derive"] }
[dev-dependencies]
criterion = {workspace = true}
+11 -3
View File
@@ -15,6 +15,9 @@ pub enum PSQError {
#[error("encountered a PSQ error")]
PSQError,
#[error("encountered a Serialization/Deserialization error")]
SerializationError,
#[error("encountered an IO error: {0}")]
IoError(#[from] io::Error),
@@ -26,9 +29,6 @@ pub enum PSQError {
#[error("Unknown PSQ version")]
UnknownVersion,
#[error("Handshake timeout")]
HandshakeTimeout(#[from] tokio::time::error::Elapsed),
}
impl From<libcrux_kem::Error> for PSQError {
@@ -48,3 +48,11 @@ impl From<libcrux_psq::Error> for PSQError {
}
}
}
impl From<tls_codec::Error> for PSQError {
fn from(err: tls_codec::Error) -> Self {
match err {
err => PSQError::SerializationError,
}
}
}
-1
View File
@@ -3,4 +3,3 @@
pub mod error;
pub mod psq;
pub mod stream;
+87 -53
View File
@@ -3,15 +3,18 @@
use std::{marker::PhantomData, time::Duration};
use libcrux_ed25519::VerificationKey;
use libcrux_psq::{
cred::Ed25519,
psk_registration::{Initiator, InitiatorMsg, Responder, ResponderMsg},
traits::{Decode, Encode, PSQ},
traits::PSQ,
};
use tls_codec::{Deserialize, Serialize};
use libcrux_traits::kem::KEM;
use nym_crypto::asymmetric::ed25519;
use rand::CryptoRng;
use rand::{CryptoRng, RngCore};
use crate::error::PSQError;
@@ -19,44 +22,53 @@ pub const PSK_LENGTH: usize = 32;
pub const CONTEXT_LEN: usize = 16;
pub const PSK_HANDLE_LEN: usize = 16;
pub struct PSQInitiator<'a, T: PSQ> {
signing_keypair: &'a ed25519::KeyPair,
pub struct PSQInitiator<T: PSQ> {
signing_key: [u8; 32],
verification_key: VerificationKey,
state: Option<Initiator>,
_t: PhantomData<T>,
}
impl<'a, T: PSQ> PSQInitiator<'a, T> {
pub fn init(signing_keypair: &'a ed25519::KeyPair) -> Self {
impl<'a, T: PSQ> PSQInitiator<T> {
pub fn init(signing_key: impl AsRef<[u8]>, verification_key: impl AsRef<[u8]>) -> Self {
let mut sig_key: [u8; 32] = [0u8; 32];
sig_key.clone_from_slice(signing_key.as_ref());
let mut verif_key: [u8; 32] = [0u8; 32];
verif_key.clone_from_slice(verification_key.as_ref());
Self {
signing_keypair,
signing_key: sig_key,
verification_key: VerificationKey::from_bytes(verif_key),
state: None,
_t: PhantomData,
}
}
pub fn compute_initiator_message(
pub fn compute_initiator_message<R>(
&mut self,
rng: &mut impl CryptoRng,
responder_kem_public_key: &'a <T::InnerKEM as KEM>::EncapsulationKey,
rng: &mut R,
responder_kem_public_key: &<T::InnerKEM as KEM>::EncapsulationKey,
context: &[u8; CONTEXT_LEN],
psk_ttl: Duration,
) -> Result<InitiatorMsg<T::InnerKEM>, PSQError> {
) -> Result<Vec<u8>, PSQError>
where
<T as PSQ>::InnerKEM: KEM,
R: RngCore + CryptoRng,
{
let (state, message) = Initiator::send_initial_message::<Ed25519, T>(
context,
psk_ttl,
&responder_kem_public_key,
&self.signing_keypair.private_key().to_bytes(),
&self.signing_keypair.public_key().to_bytes(),
responder_kem_public_key,
&self.signing_key,
&self.verification_key,
rng,
)?;
self.state = Some(state);
let initiator_msg_bytes: Vec<u8> = message.encode();
let message_bytes = message.tls_serialize_detached()?;
// This only works for MlKem768
// let decoded_msg: InitiatorMsg<T::InnerKEM> = InitiatorMsg::decode(&initiator_msg_bytes)?.0;
Ok(message)
Ok(message_bytes)
}
// The initator generates the PSK themselves while computing the first message.
@@ -70,30 +82,42 @@ impl<'a, T: PSQ> PSQInitiator<'a, T> {
// We need to ask Cryspen to make Initiator.k_pq to be accessible for this.
// pub fn get_psk(&self) -> Option<[u8; 32]> {
// match self.state {
// Some(initiator_state) => initiator_state.k_pq,
// None => None,
// }
// }
pub fn finalize(&self, responder_message: &ResponderMsg) -> Result<[u8; PSK_LENGTH], PSQError> {
pub fn get_psk(&self) -> Option<[u8; 32]> {
match &self.state {
Some(state) => match state.complete_handshake(responder_message) {
Ok(psk) => Ok(psk.psk),
Some(initiator_state) => Some(initiator_state.derive_unregistered_psk().unwrap()),
None => None,
}
}
pub fn finalize(&self, responder_message: &[u8]) -> Result<[u8; PSK_LENGTH], PSQError> {
match &self.state {
Some(state) => match ResponderMsg::tls_deserialize_exact(responder_message) {
Ok(deserialized_responder_message) => {
match state.complete_handshake(&deserialized_responder_message) {
Ok(psk) => Ok(psk.psk),
Err(err) => Err(err.into()),
}
}
Err(err) => Err(err.into()),
},
None => Err(PSQError::IncorrectStateError),
}
}
}
pub struct PSQResponder<'a, T: PSQ> {
pub struct PSQResponder<'a, T: PSQ>
where
<T as PSQ>::InnerKEM: KEM,
{
kem_private_key: &'a <T::InnerKEM as KEM>::DecapsulationKey,
kem_public_key: &'a <T::InnerKEM as KEM>::EncapsulationKey,
_t: PhantomData<T>,
}
impl<'a, T: PSQ> PSQResponder<'a, T> {
impl<'a, T: PSQ> PSQResponder<'a, T>
where
<T as PSQ>::InnerKEM: KEM,
{
pub fn init(
kem_private_key: &'a <T::InnerKEM as KEM>::DecapsulationKey,
kem_public_key: &'a <T::InnerKEM as KEM>::EncapsulationKey,
@@ -107,23 +131,27 @@ impl<'a, T: PSQ> PSQResponder<'a, T> {
pub fn compute_responder_message(
&self,
initiator_verification_key: &'a ed25519::PublicKey,
initiator_message: &InitiatorMsg<T::InnerKEM>,
initiator_verification_key: &VerificationKey,
initiator_message: &mut [u8],
context: &[u8; CONTEXT_LEN],
psk_ttl: Duration,
psk_handle: &[u8; PSK_HANDLE_LEN],
) -> Result<([u8; PSK_LENGTH], ResponderMsg), PSQError> {
) -> Result<([u8; PSK_LENGTH], Vec<u8>), PSQError> {
let deserialized_initiator_message =
InitiatorMsg::tls_deserialize_exact(initiator_message)?;
let (registered_psk, responder_msg) = Responder::send::<Ed25519, T>(
psk_handle,
psk_ttl,
context,
self.kem_public_key,
self.kem_private_key,
&initiator_verification_key.to_bytes(),
initiator_message,
&initiator_verification_key,
&deserialized_initiator_message,
)?;
Ok((registered_psk.psk, responder_msg))
let responder_bytes = responder_msg.tls_serialize_detached()?;
Ok((registered_psk.psk, responder_bytes))
}
}
@@ -131,9 +159,9 @@ impl<'a, T: PSQ> PSQResponder<'a, T> {
mod test {
use std::time::Duration;
use libcrux_ed25519::VerificationKey;
use libcrux_psq::impls::{MlKem768, XWingKemDraft06, X25519};
use libcrux_psq::psk_registration::{Initiator, InitiatorMsg, Responder};
use libcrux_psq::traits::{Decode, Encode, PSQ};
use libcrux_psq::traits::PSQ;
use libcrux_traits::kem::KEM;
use nym_crypto::asymmetric::ed25519;
use rand::prelude::*;
@@ -142,11 +170,13 @@ mod test {
use super::{PSQInitiator, PSQResponder};
fn test_helper<T: PSQ>(
rng: &mut impl CryptoRng,
fn test_helper<T: PSQ, R>(
rng: &mut R,
responder_kem_private_key: <T::InnerKEM as KEM>::DecapsulationKey,
responder_kem_public_key: <T::InnerKEM as KEM>::EncapsulationKey,
) {
) where
R: CryptoRng + rand::RngCore,
{
// set PSQ TTL
let psk_ttl: Duration = Duration::from_secs(3600);
@@ -163,23 +193,26 @@ mod test {
rng.fill_bytes(&mut secret_initiator);
let initiator_ed25519_keypair = ed25519::KeyPair::from_secret(secret_initiator, 0);
let mut initiator: PSQInitiator<T> = PSQInitiator::init(&initiator_ed25519_keypair);
let mut initiator: PSQInitiator<T> = PSQInitiator::init(
initiator_ed25519_keypair.private_key().to_bytes(),
initiator_ed25519_keypair.public_key().to_bytes(),
);
let responder: PSQResponder<T> =
PSQResponder::init(&responder_kem_private_key, &responder_kem_public_key);
let initiator_msg = initiator
let mut initiator_msg = initiator
.compute_initiator_message(rng, &responder_kem_public_key, &context, psk_ttl)
.unwrap();
let initiator_msg_bytes = initiator_msg.encode();
let pre_response_psk = initiator.get_psk().unwrap();
// let decoded_initiator_msg: (InitiatorMsg<T::InnerKEM>, usize) =
// InitiatorMsg::decode(&initiator_msg_bytes).unwrap();
let initiator_public_key =
VerificationKey::from_bytes(initiator_ed25519_keypair.public_key().to_bytes());
let (responder_psk, responder_msg) = responder
.compute_responder_message(
initiator_ed25519_keypair.public_key(),
&initiator_msg,
&initiator_public_key,
&mut initiator_msg,
&context,
psk_ttl,
&psk_handle,
@@ -188,6 +221,7 @@ mod test {
let initiator_psk = initiator.finalize(&responder_msg).unwrap();
assert_eq!(initiator_psk, pre_response_psk);
assert_eq!(initiator_psk, responder_psk);
}
@@ -199,7 +233,7 @@ mod test {
let (responder_kem_private_key, responder_kem_public_key) =
libcrux_kem::key_gen(libcrux_kem::Algorithm::MlKem768, &mut rng).unwrap();
test_helper::<MlKem768>(
test_helper::<MlKem768, _>(
&mut rng,
responder_kem_private_key,
responder_kem_public_key,
@@ -215,7 +249,7 @@ mod test {
libcrux_kem::key_gen(libcrux_kem::Algorithm::XWingKemDraft06, &mut rand::rng())
.unwrap();
test_helper::<XWingKemDraft06>(
test_helper::<XWingKemDraft06, _>(
&mut rng,
responder_kem_private_key,
responder_kem_public_key,
@@ -230,7 +264,7 @@ mod test {
let (responder_kem_private_key, responder_kem_public_key) =
libcrux_kem::key_gen(libcrux_kem::Algorithm::X25519, &mut rand::rng()).unwrap();
test_helper::<X25519>(
test_helper::<X25519, _>(
&mut rng,
responder_kem_private_key,
responder_kem_public_key,
-1
View File
@@ -1 +0,0 @@