Fix bug where proposal got different psk id than key schedule

awslabs · Jan 15, 2025 · 38433c1 · 38433c1
1 parent 27465a0
commit 38433c1
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 7 deletions.
diff --git a/mls-rs/src/group/commit/builder.rs b/mls-rs/src/group/commit/builder.rs
@@ -33,7 +33,9 @@ use crate::WireFormat;
 
 #[cfg(feature = "psk")]
 use crate::{
-    group::{JustPreSharedKeyID, PskGroupId, ResumptionPSKUsage, ResumptionPsk},
+    group::{
+        JustPreSharedKeyID, PreSharedKeyProposal, PskGroupId, ResumptionPSKUsage, ResumptionPsk,
+    },
     psk::{
         secret::{PskSecret, PskSecretInput},
         ExternalPskId,
@@ -243,14 +245,17 @@ where
         use crate::psk::PreSharedKeyID;
 
         let key_id = JustPreSharedKeyID::External(id);
-        let proposal = self.group.psk_proposal(key_id.clone())?;
-        self.proposals.push(proposal);
+        let id = PreSharedKeyID::new(key_id, &self.group.cipher_suite_provider)?;
 
         self.psks.push(PskSecretInput {
-            id: PreSharedKeyID::new(key_id, &self.group.cipher_suite_provider)?,
+            id: id.clone(),
             psk,
         });
 
+        let proposal = Proposal::Psk(PreSharedKeyProposal { psk: id });
+
+        self.proposals.push(proposal);
+
         Ok(self)
     }
 
@@ -286,14 +291,17 @@ where
         };
 
         let key_id = JustPreSharedKeyID::Resumption(psk_id);
-        let proposal = self.group.psk_proposal(key_id.clone())?;
-        self.proposals.push(proposal);
+        let id = PreSharedKeyID::new(key_id, &self.group.cipher_suite_provider)?;
 
         self.psks.push(PskSecretInput {
-            id: PreSharedKeyID::new(key_id, &self.group.cipher_suite_provider)?,
+            id: id.clone(),
             psk,
         });
 
+        let proposal = Proposal::Psk(PreSharedKeyProposal { psk: id });
+
+        self.proposals.push(proposal);
+
         Ok(self)
     }
 

diff --git a/mls-rs/src/group/mod.rs b/mls-rs/src/group/mod.rs
@@ -3679,6 +3679,44 @@ mod tests {
         .unwrap();
     }
 
+    #[cfg(feature = "psk")]
+    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
+    async fn can_process_with_psk() {
+        let mut alice = test_group(TEST_PROTOCOL_VERSION, TEST_CIPHER_SUITE).await;
+        let (mut bob, _) = alice.join("bob");
+
+        let psk_id_external = ExternalPskId::new(vec![0]);
+        let psk_external = PreSharedKey::from(vec![1]);
+        let psk_epoch = bob.context().epoch;
+
+        let psk_id_resumption = ResumptionPsk {
+            usage: ResumptionPSKUsage::Application,
+            psk_group_id: bob.context().group_id.clone().into(),
+            psk_epoch,
+        };
+
+        let psk_resumption = bob.resumption_secret(psk_epoch).unwrap();
+
+        let commit = alice
+            .commit_builder()
+            .add_external_psk(psk_id_external.clone(), psk_external.clone())
+            .unwrap()
+            .add_resumption_psk(psk_epoch, psk_resumption.clone())
+            .unwrap()
+            .build()
+            .await
+            .unwrap();
+
+        bob.commit_processor(commit.commit_message)
+            .await
+            .unwrap()
+            .with_external_psk(psk_id_external, psk_external)
+            .with_resumption_psk(psk_id_resumption, psk_resumption)
+            .process()
+            .await
+            .unwrap();
+    }
+
     #[cfg(feature = "by_ref_proposal")]
     #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
     async fn invalid_update_does_not_prevent_other_updates() {

diff --git a/mls-rs/src/psk.rs b/mls-rs/src/psk.rs
@@ -66,6 +66,12 @@ impl Debug for PskGroupId {
     }
 }
 
+impl From<Vec<u8>> for PskGroupId {
+    fn from(value: Vec<u8>) -> Self {
+        Self(value)
+    }
+}
+
 #[derive(Clone, Eq, Hash, PartialEq, PartialOrd, Ord, MlsSize, MlsEncode, MlsDecode)]
 #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]