refactor ProtocolVersion to enum

webrtc-rs · Aug 30, 2023 · 5a518de · 5a518de
1 parent ef06da9
commit 5a518de
Show file tree

Hide file tree

Showing 8 changed files with 44 additions and 21 deletions.
diff --git a/rmls/src/framing/mod.rs b/rmls/src/framing/mod.rs
@@ -14,9 +14,32 @@ use crate::serde::*;
 use crate::tree::math::LeafIndex;
 use crate::tree::secret::RatchetSecret;
 
-pub(crate) type ProtocolVersion = u16;
+/// [RFC9420 Sec.6](https://www.rfc-editor.org/rfc/rfc9420.html#section-6) ProtocolVersion
+#[derive(Default, Debug, Copy, Clone, Eq, PartialEq, Hash)]
+#[repr(u16)]
+pub enum ProtocolVersion {
+    #[default]
+    MLS10 = 1,
+    Unsupported(u16),
+}
 
-pub(crate) const PROTOCOL_VERSION_MLS10: ProtocolVersion = 1;
+impl From<u16> for ProtocolVersion {
+    fn from(v: u16) -> Self {
+        match v {
+            1 => ProtocolVersion::MLS10,
+            _ => ProtocolVersion::Unsupported(v),
+        }
+    }
+}
+
+impl From<ProtocolVersion> for u16 {
+    fn from(val: ProtocolVersion) -> u16 {
+        match val {
+            ProtocolVersion::MLS10 => 1,
+            ProtocolVersion::Unsupported(v) => v,
+        }
+    }
+}
 
 #[derive(Default, Debug, Copy, Clone, Eq, PartialEq)]
 #[repr(u8)]
@@ -367,7 +390,7 @@ impl AuthenticatedContent {
 
     fn framed_content_tbs(&self, ctx: &GroupContext) -> FramedContentTBS {
         FramedContentTBS {
-            version: PROTOCOL_VERSION_MLS10,
+            version: ProtocolVersion::MLS10,
             wire_format: self.wire_format,
             content: self.content.clone(),
             context: Some(ctx.clone()),
@@ -491,7 +514,7 @@ impl Deserializer for FramedContentTBS {
         if buf.remaining() < 2 {
             return Err(Error::BufferTooSmall);
         }
-        let version = buf.get_u16();
+        let version = buf.get_u16().into();
         let wire_format = WireFormat::deserialize(buf)?;
         let content = FramedContent::deserialize(buf)?;
         let context = match &content.sender {
@@ -514,7 +537,7 @@ impl Serializer for FramedContentTBS {
         Self: Sized,
         B: BufMut,
     {
-        buf.put_u16(self.version);
+        buf.put_u16(self.version.into());
         self.wire_format.serialize(buf)?;
         self.content.serialize(buf)?;
         match &self.content.sender {

diff --git a/rmls/src/group/group_test.rs b/rmls/src/group/group_test.rs
@@ -308,7 +308,7 @@ fn message_protection_test(
     tc: &MessageProtectionTest,
 ) -> Result<()> {
     let ctx = GroupContext {
-        version: PROTOCOL_VERSION_MLS10,
+        version: ProtocolVersion::MLS10,
         cipher_suite,
         group_id: tc.group_id.clone().into(),
         epoch: tc.epoch,

diff --git a/rmls/src/group/mod.rs b/rmls/src/group/mod.rs
@@ -359,10 +359,10 @@ impl Deserializer for Message {
         if buf.remaining() < 2 {
             return Err(Error::BufferTooSmall);
         }
-        let version = buf.get_u16();
+        let version: ProtocolVersion = buf.get_u16().into();
 
-        if version != PROTOCOL_VERSION_MLS10 {
-            return Err(Error::InvalidProtocolVersion(version));
+        if version != ProtocolVersion::MLS10 {
+            return Err(Error::InvalidProtocolVersion(version.into()));
         }
 
         let wire_format = WireFormat::deserialize(buf)?;
@@ -392,7 +392,7 @@ impl Serializer for Message {
         Self: Sized,
         B: BufMut,
     {
-        buf.put_u16(self.version);
+        buf.put_u16(self.version.into());
         self.wire_format.serialize(buf)?;
         match &self.message {
             WireFormatMessage::PublicMessage(message) => {

diff --git a/rmls/src/group/proposal.rs b/rmls/src/group/proposal.rs
@@ -296,7 +296,7 @@ impl Deserializer for ReInitProposal {
         if buf.remaining() < 4 {
             return Err(Error::BufferTooSmall);
         }
-        let version = buf.get_u16();
+        let version = buf.get_u16().into();
         let cipher_suite = buf.get_u16().try_into()?;
 
         let extensions = deserialize_extensions(buf)?;
@@ -316,7 +316,7 @@ impl Serializer for ReInitProposal {
         B: BufMut,
     {
         serialize_opaque_vec(&self.group_id, buf)?;
-        buf.put_u16(self.version);
+        buf.put_u16(self.version.into());
         buf.put_u16(self.cipher_suite as u16);
         serialize_extensions(&self.extensions, buf)
     }

diff --git a/rmls/src/key/key_test.rs b/rmls/src/key/key_test.rs
@@ -159,7 +159,7 @@ fn key_schedule_test(
         println!("epoch {}", i);
 
         let ctx = GroupContext {
-            version: PROTOCOL_VERSION_MLS10,
+            version: ProtocolVersion::MLS10,
             cipher_suite,
             group_id: tc.group_id.clone().into(),
             epoch: i as u64,

diff --git a/rmls/src/key/package.rs b/rmls/src/key/package.rs
@@ -29,7 +29,7 @@ impl Deserializer for KeyPackage {
             return Err(Error::BufferTooSmall);
         }
 
-        let version = buf.get_u16();
+        let version = buf.get_u16().into();
         let cipher_suite = buf.get_u16().try_into()?;
         let init_key = deserialize_opaque_vec(buf)?;
         let leaf_node = LeafNode::deserialize(buf)?;
@@ -64,7 +64,7 @@ impl KeyPackage {
         Self: Sized,
         B: BufMut,
     {
-        buf.put_u16(self.version);
+        buf.put_u16(self.version.into());
         buf.put_u16(self.cipher_suite as u16);
         serialize_opaque_vec(&self.init_key, buf)?;
         self.leaf_node.serialize(buf)?;

diff --git a/rmls/src/key/schedule.rs b/rmls/src/key/schedule.rs
@@ -27,7 +27,7 @@ impl Deserializer for GroupContext {
             return Err(Error::BufferTooSmall);
         }
 
-        let version = buf.get_u16();
+        let version: ProtocolVersion = buf.get_u16().into();
         let cipher_suite = buf.get_u16().try_into()?;
         let group_id = deserialize_opaque_vec(buf)?;
         if buf.remaining() < 8 {
@@ -37,8 +37,8 @@ impl Deserializer for GroupContext {
         let tree_hash = deserialize_opaque_vec(buf)?;
         let confirmed_transcript_hash = deserialize_opaque_vec(buf)?;
 
-        if version != PROTOCOL_VERSION_MLS10 {
-            return Err(Error::InvalidProposalTypeValue(version));
+        if version != ProtocolVersion::MLS10 {
+            return Err(Error::InvalidProposalTypeValue(version.into()));
         }
 
         let extensions = deserialize_extensions(buf)?;
@@ -60,7 +60,7 @@ impl Serializer for GroupContext {
         Self: Sized,
         B: BufMut,
     {
-        buf.put_u16(self.version);
+        buf.put_u16(self.version.into());
         buf.put_u16(self.cipher_suite as u16);
         serialize_opaque_vec(&self.group_id, buf)?;
         buf.put_u64(self.epoch);

diff --git a/rmls/src/tree/mod.rs b/rmls/src/tree/mod.rs
@@ -172,7 +172,7 @@ impl Deserializer for Capabilities {
             if b.remaining() < 2 {
                 return Err(Error::BufferTooSmall);
             }
-            let ver: ProtocolVersion = b.get_u16();
+            let ver: ProtocolVersion = b.get_u16().into();
             versions.push(ver);
             Ok(())
         })?;
@@ -236,7 +236,7 @@ impl Serializer for Capabilities {
             self.versions.len(),
             buf,
             |i: usize, b: &mut BytesMut| -> Result<()> {
-                b.put_u16(self.versions[i]);
+                b.put_u16(self.versions[i].into());
                 Ok(())
             },
         )?;