feat: add header integrity check and replay protection to control messages

data-plane/core/auth/src/utils.rs

@@ -31,6 +31,34 @@ pub fn generate_mls_signature_keys() -> Result<(Vec<u8>, Vec<u8>), crate::errors
     ))
 }
 
+/// Sign the header AAD bytes using an Ed25519 private key.
+pub fn sign_header_aad(
+    aad_bytes: &[u8],
+    private_key_bytes: &[u8],
+) -> Result<Vec<u8>, crate::errors::AuthError> {
+    use aws_lc_rs::signature::Ed25519KeyPair;
+    let key_pair = Ed25519KeyPair::from_seed_and_public_key(
+        &private_key_bytes[..32],
+        &private_key_bytes[32..],
+    )
+    .map_err(|_| crate::errors::AuthError::MlsKeyGenerationFailed)?;
+    let signature = key_pair.sign(aad_bytes);
+    Ok(signature.as_ref().to_vec())
+}
+
+/// Verify the header AAD signature using an Ed25519 public key.
+pub fn verify_header_aad(
+    aad_bytes: &[u8],
+    signature_bytes: &[u8],
+    public_key_bytes: &[u8],
+) -> Result<(), crate::errors::AuthError> {
+    use aws_lc_rs::signature::{ED25519, UnparsedPublicKey};
+    let public_key = UnparsedPublicKey::new(&ED25519, public_key_bytes);
+    public_key
+        .verify(aad_bytes, signature_bytes)
+        .map_err(|_| crate::errors::AuthError::TokenInvalid)
+}
+
 /// Convert arbitrary bytes into a PEM-formatted string with the provided header/footer.
 /// The body is wrapped at 64 character lines per RFC 7468 guidance.
 /// Header/footer should include the BEGIN/END lines with trailing/leading newlines as desired.
@@ -58,6 +86,14 @@ pub fn bytes_to_pem(key_bytes: &[u8], header: &str, footer: &str) -> String {
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_sign_header_aad_with_generated_keys() {
+        let (secret, public) = generate_mls_signature_keys().unwrap();
+        let msg = b"hello";
+        let sig = sign_header_aad(msg, &secret).unwrap();
+        verify_header_aad(msg, &sig, &public).unwrap();
+    }
+
     #[test]
     fn test_bytes_to_pem_basic() {
         let data = b"hello world"; // base64: aGVsbG8gd29ybGQ=

data-plane/core/datapath/benches/name_benchmark.rs

@@ -39,6 +39,8 @@ fn make_slim_header() -> SlimHeader {
         incoming_conn: None,
         error: None,
         header_mac: None,
+        e2e_header_sig: None,
+        sequence_number: None,
         ttl: DEFAULT_TTL,
     }
 }

data-plane/core/datapath/src/api/gen/dataplane.proto.v1.rs

@@ -96,6 +96,10 @@ pub struct SlimHeader {
     pub error: ::core::option::Option<bool>,
     #[prost(bytes = "vec", optional, tag = "11")]
     pub header_mac: ::core::option::Option<::prost::alloc::vec::Vec<u8>>,
+    #[prost(uint64, optional, tag = "12")]
+    pub sequence_number: ::core::option::Option<u64>,
+    #[prost(bytes = "vec", optional, tag = "13")]
+    pub e2e_header_sig: ::core::option::Option<::prost::alloc::vec::Vec<u8>>,
 }
 #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)]
 pub struct SessionHeader {
@@ -311,6 +315,8 @@ pub struct HeaderIntegrityAad {
     pub message_id: u32,
     #[prost(string, tag = "9")]
     pub payload_type: ::prost::alloc::string::String,
+    #[prost(uint64, optional, tag = "10")]
+    pub sequence_number: ::core::option::Option<u64>,
 }
 /// Group Add Payload
 /// sent when a new participant is added

data-plane/core/datapath/src/header_mac.rs

@@ -321,6 +321,8 @@ mod tests {
             incoming_conn: Some(999),
             error: Some(false),
             header_mac: None,
+            sequence_number: None,
+            e2e_header_sig: None,
             ttl: DEFAULT_TTL,
         }
     }

data-plane/core/datapath/src/link_ecdh.rs

@@ -78,6 +78,8 @@ mod tests {
             incoming_conn: None,
             error: None,
             header_mac: None,
+            sequence_number: None,
+            e2e_header_sig: None,
             ttl: DEFAULT_TTL,
         };
         a.sign_slim_header(&mut h, &lid).unwrap();

data-plane/core/datapath/src/messages/utils.rs

@@ -299,6 +299,8 @@ impl SlimHeader {
             incoming_conn: flags.incoming_conn,
             error: flags.error,
             header_mac: None,
+            sequence_number: None,
+            e2e_header_sig: None,
             ttl: flags.ttl,
         }
     }
@@ -2081,6 +2083,8 @@ mod tests {
             incoming_conn: None,
             error: None,
             header_mac: None,
+            sequence_number: None,
+            e2e_header_sig: None,
             ttl: DEFAULT_TTL,
         };
 

data-plane/core/session/Cargo.toml

@@ -14,6 +14,7 @@ agntcy-slim-datapath = { workspace = true }
 agntcy-slim-mls = { workspace = true }
 agntcy-slim-version = { workspace = true }
 async-trait = { workspace = true }
+base64 = { workspace = true }
 display-error-chain = { workspace = true }
 futures = { workspace = true }
 futures-timer = { workspace = true }

data-plane/core/session/src/errors.rs

@@ -79,7 +79,7 @@ pub enum SessionError {
     MlsOp(#[from] MlsError),
 
     // Authorization and roles
-    #[error("auth error")]
+    #[error("auth error: {0}")]
     Auth(#[from] AuthError),
 
     // Acknowledgements and routing
@@ -162,13 +162,13 @@ pub enum SessionError {
     ModeratorTaskUnsupportedPhase,
     #[error("unexpected timer id: {0}")]
     ModeratorTaskUnexpectedTimerId(u32),
-    #[error("failed to add participant to session")]
+    #[error("failed to add participant to session: {source}")]
     ModeratorTaskAddFailed { source: Box<SessionError> },
-    #[error("failed to remove participant from session")]
+    #[error("failed to remove participant from session: {source}")]
     ModeratorTaskRemoveFailed { source: Box<SessionError> },
-    #[error("failed to update session")]
+    #[error("failed to update session: {source}")]
     ModeratorTaskUpdateFailed { source: Box<SessionError> },
-    #[error("failed to close session")]
+    #[error("failed to close session: {source}")]
     ModeratorTaskCloseFailed { source: Box<SessionError> },
 }
 

data-plane/core/session/src/mls_state.rs

@@ -366,33 +366,39 @@ where
 
     /// Builds the Authenticated Data (AAD) for header integrity checks
     fn build_aad(&self, msg: &Message) -> Vec<u8> {
-        let slim_header = msg.get_slim_header();
-        let session_header = msg.get_session_header();
+        build_aad(msg)
+    }
+}
 
-        let payload_type = if let Some(payload) = msg.get_payload() {
-            if let Ok(app_payload) = payload.as_application_payload() {
-                app_payload.payload_type.clone()
-            } else {
-                String::new()
-            }
+/// Builds the Authenticated Data (AAD) for header integrity checks
+pub(crate) fn build_aad(msg: &Message) -> Vec<u8> {
+    let slim_header = msg.get_slim_header();
+    let session_header = msg.get_session_header();
+
+    let payload_type = if let Some(payload) = msg.get_payload() {
+        if let Ok(app_payload) = payload.as_application_payload() {
+            app_payload.payload_type.clone()
         } else {
             String::new()
-        };
-
-        let aad = HeaderIntegrityAad {
-            version: 1,
-            source: Some(slim_header.get_source().clone()),
-            destination: Some(slim_header.get_dst().clone()),
-            identity: slim_header.get_identity().to_string(),
-            session_type: session_header.session_type() as i32,
-            session_message_type: session_header.session_message_type() as i32,
-            session_id: session_header.get_session_id(),
-            message_id: session_header.get_message_id(),
-            payload_type,
-        };
-
-        aad.encode_to_vec()
-    }
+        }
+    } else {
+        String::new()
+    };
+
+    let aad = HeaderIntegrityAad {
+        version: 1,
+        source: Some(slim_header.get_source().clone()),
+        destination: Some(slim_header.get_dst().clone()),
+        identity: slim_header.get_identity().to_string(),
+        session_type: session_header.session_type() as i32,
+        session_message_type: session_header.session_message_type() as i32,
+        session_id: session_header.get_session_id(),
+        message_id: session_header.get_message_id(),
+        payload_type,
+        sequence_number: slim_header.sequence_number,
+    };
+
+    aad.encode_to_vec()
 }
 
 #[derive(Debug)]

data-plane/core/session/src/session_builder.rs

@@ -446,6 +446,7 @@ where
             graceful_shutdown_timeout: self.graceful_shutdown_timeout,
             subscription_manager,
             service_id: self.service_id.unwrap_or_default(),
+            seen_control_seqs: crate::session_settings::new_seen_control_seqs(),
         };
 
         let wrapper = wrapper_constructor(inner, settings.clone());