[Updt] Mtp 0.3.0

src/transport/connection.rs

@@ -2,24 +2,94 @@ pub(crate) use super::omikron_connection::{OmikronConnection, OmikronResult};
 
 use mtp::codec::{CommunicationValue, DataValue};
 
-pub(crate) trait MtpValueCompat {
-    fn get_id(&self) -> u32;
-    fn get_sender(&self) -> u64;
-    fn get_receiver(&self) -> u64;
+pub(crate) trait RequiredMtpFields {
+    /// Require a nonzero request correlation ID. `None` and `Some(0)` are
+    /// distinct MTP wire states, but both are invalid for Omega requests.
+    fn require_id(&self) -> OmikronResult<u32>;
+    /// Require a nonzero authenticated peer identity.
+    fn require_sender(&self) -> OmikronResult<u64>;
+    /// Require a nonzero application routing target.
+    fn require_receiver(&self) -> OmikronResult<u64>;
+    fn require_sender_i64(&self) -> OmikronResult<i64>;
+    fn require_receiver_i64(&self) -> OmikronResult<i64>;
 }
 
-impl MtpValueCompat for CommunicationValue {
-    fn get_id(&self) -> u32 {
-        self.id().unwrap_or_default()
+impl RequiredMtpFields for CommunicationValue {
+    fn require_id(&self) -> OmikronResult<u32> {
+        self.id().filter(|id| *id != 0).ok_or_else(|| {
+            crate::OmegaError::Validation("MTP message is missing request id".into())
+        })
     }
 
-    fn get_sender(&self) -> u64 {
-        self.sender().unwrap_or_default()
+    fn require_sender(&self) -> OmikronResult<u64> {
+        self.sender()
+            .filter(|sender| *sender != 0)
+            .ok_or_else(|| crate::OmegaError::Validation("MTP message is missing sender".into()))
     }
 
-    fn get_receiver(&self) -> u64 {
-        self.receiver().unwrap_or_default()
+    fn require_receiver(&self) -> OmikronResult<u64> {
+        self.receiver()
+            .filter(|receiver| *receiver != 0)
+            .ok_or_else(|| crate::OmegaError::Validation("MTP message is missing receiver".into()))
     }
+
+    fn require_sender_i64(&self) -> OmikronResult<i64> {
+        positive_i64(self.require_sender()?, "sender")
+    }
+
+    fn require_receiver_i64(&self) -> OmikronResult<i64> {
+        positive_i64(self.require_receiver()?, "receiver")
+    }
+}
+
+pub(crate) fn positive_i64(value: impl TryInto<i128>, field: &str) -> OmikronResult<i64> {
+    let value = value
+        .try_into()
+        .map_err(|_| crate::OmegaError::Validation(format!("invalid {field}")))?;
+    let value = i64::try_from(value)
+        .map_err(|_| crate::OmegaError::Validation(format!("invalid {field}")))?;
+    if value <= 0 {
+        return Err(crate::OmegaError::Validation(format!("invalid {field}")));
+    }
+    Ok(value)
+}
+
+pub(crate) fn validate_dispatch_fields(value: &CommunicationValue) -> OmikronResult<()> {
+    let Some(message_type) = value.get_comm_type_enum() else {
+        return Err(crate::OmegaError::Validation(
+            "MTP message has an unknown communication type".into(),
+        ));
+    };
+
+    if !matches!(
+        message_type,
+        mtp::codec::CommunicationType::ClientChanged
+            | mtp::codec::CommunicationType::PushNotification
+    ) {
+        value.require_id()?;
+    }
+
+    if matches!(
+        message_type,
+        mtp::codec::CommunicationType::GetUserData
+            | mtp::codec::CommunicationType::ChangeUserData
+            | mtp::codec::CommunicationType::ChangeIotaData
+            | mtp::codec::CommunicationType::DeleteUser
+            | mtp::codec::CommunicationType::AttachUserBegin
+            | mtp::codec::CommunicationType::AttachUserComplete
+            | mtp::codec::CommunicationType::DeleteUserCredentialBegin
+            | mtp::codec::CommunicationType::DeleteUserCredentialComplete
+            | mtp::codec::CommunicationType::EraseHostedUserDataAck
+            | mtp::codec::CommunicationType::ReleaseUserFromIota
+            | mtp::codec::CommunicationType::DeleteIota
+            | mtp::codec::CommunicationType::GetNotifications
+            | mtp::codec::CommunicationType::ReadNotification
+            | mtp::codec::CommunicationType::StateSubscribe
+    ) {
+        value.require_sender()?;
+    }
+
+    Ok(())
 }
 
 pub(crate) trait OptionalDataValueCompat {
@@ -46,3 +116,65 @@ impl OptionalDataValueCompat for Option<&DataValue> {
         self.and_then(|value| value.as_bytes())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{RequiredMtpFields, validate_dispatch_fields};
+    use mtp::codec::{CommunicationType, CommunicationValue};
+
+    #[test]
+    fn missing_request_id_is_rejected_without_a_zero_fallback() {
+        let value = CommunicationValue::new(CommunicationType::GetUserData).without_id();
+        assert!(value.require_id().is_err());
+        assert!(validate_dispatch_fields(&value).is_err());
+    }
+
+    #[test]
+    fn zero_routing_values_are_rejected_by_omega_contract() {
+        let value = CommunicationValue::new(CommunicationType::GetUserData)
+            .with_id(0)
+            .with_sender(0)
+            .with_receiver(0);
+
+        assert!(value.require_id().is_err());
+        assert!(value.require_sender().is_err());
+        assert!(value.require_receiver().is_err());
+        assert!(validate_dispatch_fields(&value).is_err());
+    }
+
+    #[test]
+    fn security_sensitive_messages_require_sender() {
+        for message_type in [
+            CommunicationType::GetUserData,
+            CommunicationType::ChangeUserData,
+            CommunicationType::ChangeIotaData,
+            CommunicationType::DeleteUser,
+            CommunicationType::AttachUserBegin,
+            CommunicationType::AttachUserComplete,
+            CommunicationType::DeleteUserCredentialBegin,
+            CommunicationType::DeleteUserCredentialComplete,
+            CommunicationType::EraseHostedUserDataAck,
+            CommunicationType::ReleaseUserFromIota,
+            CommunicationType::DeleteIota,
+            CommunicationType::GetNotifications,
+            CommunicationType::ReadNotification,
+            CommunicationType::StateSubscribe,
+        ] {
+            let value = CommunicationValue::new(message_type).with_id(1);
+            assert!(value.require_sender().is_err());
+            assert!(validate_dispatch_fields(&value).is_err());
+        }
+    }
+
+    #[test]
+    fn optional_client_changed_id_remains_optional() {
+        let value = CommunicationValue::new(CommunicationType::ClientChanged).without_id();
+        assert!(validate_dispatch_fields(&value).is_ok());
+    }
+
+    #[test]
+    fn push_notification_uses_its_logical_sender_field() {
+        let value = CommunicationValue::new(CommunicationType::PushNotification);
+        assert!(validate_dispatch_fields(&value).is_ok());
+    }
+}