[Add] Ease of use functions

codec/src/communication_value.rs

@@ -3,7 +3,7 @@ use std::collections::BTreeMap;
 use std::fmt;
 use std::io::{Cursor, Read};
 
-use crate::data_value::DataValue;
+use crate::data_value::{DataKind, DataValue};
 use crate::rand_u32;
 use mtp_common::CodecError;
 use mtp_type_map::{
@@ -12,9 +12,7 @@ use mtp_type_map::{
 };
 
 #[cfg(feature = "crypto")]
-use mtp_crypto::SigAlgorithm;
-#[cfg(feature = "crypto")]
-use mtp_crypto::SignatureScheme;
+use mtp_crypto::{PublicKeyBundle, SigAlgorithm, SignatureScheme};
 
 const FLAG_HAS_SENDER: u8 = 0b0000_0001;
 const FLAG_HAS_RECEIVER: u8 = 0b0000_0010;
@@ -121,13 +119,147 @@ impl CommunicationValue {
         self
     }
 
-    pub fn get_data(&self, data_type: DataTypeId) -> &DataValue {
-        self.data.get(&data_type).unwrap_or(&DataValue::Null)
+    pub fn get_data(&self, data_type: DataType) -> &DataValue {
+        let tm_owned;
+        let tm = match &self.type_map {
+            Some(tm) => tm,
+            None => {
+                tm_owned = TypeMap::latest();
+                &tm_owned
+            }
+        };
+        match tm.data_id_enum(data_type) {
+            Some(raw_id) => self.data.get(&DataTypeId(raw_id)).unwrap_or(&DataValue::Null),
+            None => &DataValue::Null,
+        }
+    }
+
+    pub fn get_data_opt(&self, data_type: DataType) -> Option<&DataValue> {
+        let tm_owned;
+        let tm = match &self.type_map {
+            Some(tm) => tm,
+            None => {
+                tm_owned = TypeMap::latest();
+                &tm_owned
+            }
+        };
+        let raw_id = tm.data_id_enum(data_type)?;
+        self.data.get(&DataTypeId(raw_id))
+    }
+
+    pub fn has_data(&self, data_type: DataType) -> Option<DataKind> {
+        self.get_data_opt(data_type).map(|v| v.kind())
+    }
+
+    pub fn get_comm_type_enum(&self) -> Option<CommunicationType> {
+        let tm_owned;
+        let tm = match &self.type_map {
+            Some(tm) => tm,
+            None => {
+                tm_owned = TypeMap::latest();
+                &tm_owned
+            }
+        };
+        tm.comm_enum_id(self.comm_type.0)
     }
 
     pub fn data(&self) -> &BTreeMap<DataTypeId, DataValue> {
         &self.data
     }
+
+    pub fn data_len(&self) -> usize {
+        self.data.len()
+    }
+
+    // ── type checks ──────────────────────────────────────────────────────────
+
+    pub fn is_type(&self, comm_type: CommunicationType) -> bool {
+        self.get_comm_type_enum() == Some(comm_type)
+    }
+
+    pub fn get_type_name(&self) -> Option<&'static str> {
+        communication_type_name(self.comm_type.0)
+    }
+
+    // ── mutation ─────────────────────────────────────────────────────────────
+
+    pub fn set_data(&mut self, data_type: DataType, value: DataValue) {
+        let tm = self.type_map.clone().unwrap_or_else(TypeMap::latest);
+        if let Some(raw_id) = tm.data_id_enum(data_type) {
+            self.data.insert(DataTypeId(raw_id), value);
+        }
+    }
+
+    #[must_use]
+    pub fn with_data(mut self, data_type: DataType, value: DataValue) -> Self {
+        self.set_data(data_type, value);
+        self
+    }
+
+    pub fn remove_data(&mut self, data_type: DataType) -> Option<DataValue> {
+        let tm_owned;
+        let tm = match &self.type_map {
+            Some(tm) => tm,
+            None => {
+                tm_owned = TypeMap::latest();
+                &tm_owned
+            }
+        };
+        let raw_id = tm.data_id_enum(data_type)?;
+        self.data.remove(&DataTypeId(raw_id))
+    }
+
+    #[must_use]
+    pub fn reply_to(&self, comm_type: CommunicationType) -> Self {
+        Self::new(comm_type)
+            .with_sender(self.receiver)
+            .with_receiver(self.sender)
+    }
+
+    pub fn merge(&mut self, other: &CommunicationValue) {
+        for (id, value) in &other.data {
+            self.data.insert(*id, value.clone());
+        }
+    }
+
+    // ── typed iteration ──────────────────────────────────────────────────────
+
+    pub fn iter_typed_data(&self) -> impl Iterator<Item = (Option<DataType>, &DataValue)> + '_ {
+        let tm = self.type_map.clone().unwrap_or_else(TypeMap::latest);
+        self.data
+            .iter()
+            .map(move |(id, val)| (tm.data_enum_id(id.0), val))
+    }
+
+    // ── typed field accessors ─────────────────────────────────────────────────
+
+    pub fn get_bool(&self, data_type: DataType) -> Option<bool> {
+        self.get_data_opt(data_type)?.as_bool()
+    }
+
+    pub fn get_str(&self, data_type: DataType) -> Option<&str> {
+        self.get_data_opt(data_type)?.as_str()
+    }
+
+    pub fn get_u128(&self, data_type: DataType) -> Option<u128> {
+        self.get_data_opt(data_type)?.as_unsigned_number()
+    }
+
+    pub fn get_i128(&self, data_type: DataType) -> Option<i128> {
+        self.get_data_opt(data_type)?.as_signed_number()
+    }
+
+    pub fn get_float(&self, data_type: DataType) -> Option<(u8, u32)> {
+        self.get_data_opt(data_type)?.as_float()
+    }
+
+    pub fn get_bytes(&self, data_type: DataType) -> Option<&[u8]> {
+        self.get_data_opt(data_type)?.as_bytes_slice()
+    }
+
+    pub fn get_array(&self, data_type: DataType) -> Option<&[DataValue]> {
+        self.get_data_opt(data_type)?.as_array_slice()
+    }
 }
 
 impl CommunicationValue {
@@ -433,6 +565,64 @@ impl CommunicationValue {
         self.frame_signature.as_ref()
     }
 
+    /*
+     * Verify the frame signature using a `PublicKeyBundle`. Dispatches to
+     * Ed25519, ML-DSA-65, or both (DUAL) based on the stored algorithm byte.
+     * Returns `false` if the frame has no signature or verification fails.
+     */
+    #[cfg(feature = "crypto")]
+    pub fn validate_signature(&self, pk: &PublicKeyBundle) -> bool {
+        let Some((alg, _)) = &self.frame_signature else {
+            return false;
+        };
+        struct Ed25519Verifier<'a>(&'a mtp_crypto::SignaturePublicKey);
+        impl SignatureScheme for Ed25519Verifier<'_> {
+            fn sign(&self, _: &[u8]) -> Result<Vec<u8>, mtp_crypto::CryptoError> {
+                Err(mtp_crypto::CryptoError::SigningFailed)
+            }
+            fn verify(&self, msg: &[u8], sig: &[u8]) -> Result<(), mtp_crypto::CryptoError> {
+                mtp_crypto::verify_ed25519(self.0, msg, sig)
+            }
+        }
+        struct MlDsaVerifier<'a>(&'a mtp_crypto::SignaturePqPublicKey);
+        impl SignatureScheme for MlDsaVerifier<'_> {
+            fn sign(&self, _: &[u8]) -> Result<Vec<u8>, mtp_crypto::CryptoError> {
+                Err(mtp_crypto::CryptoError::SigningFailed)
+            }
+            fn verify(&self, msg: &[u8], sig: &[u8]) -> Result<(), mtp_crypto::CryptoError> {
+                mtp_crypto::verify_ml_dsa(self.0, msg, sig)
+            }
+        }
+        match *alg {
+            SigAlgorithm::ED25519 => {
+                self.verify_frame(&Ed25519Verifier(&pk.sig_cl_public_key)).is_ok()
+            }
+            SigAlgorithm::ML_DSA_65 => {
+                self.verify_frame(&MlDsaVerifier(&pk.sig_pq_public_key)).is_ok()
+            }
+            SigAlgorithm::DUAL => {
+                // For DUAL, verify_frame passes the full combined sig to the verifier.
+                // We wrap a verifier that splits and checks both halves.
+                struct DualVerifier<'a>(&'a mtp_crypto::SignaturePublicKey, &'a mtp_crypto::SignaturePqPublicKey);
+                impl SignatureScheme for DualVerifier<'_> {
+                    fn sign(&self, _: &[u8]) -> Result<Vec<u8>, mtp_crypto::CryptoError> {
+                        Err(mtp_crypto::CryptoError::SigningFailed)
+                    }
+                    fn verify(&self, msg: &[u8], sig: &[u8]) -> Result<(), mtp_crypto::CryptoError> {
+                        const ED_LEN: usize = 64;
+                        if sig.len() < ED_LEN {
+                            return Err(mtp_crypto::CryptoError::InvalidSignature);
+                        }
+                        mtp_crypto::verify_ed25519(self.0, msg, &sig[..ED_LEN])?;
+                        mtp_crypto::verify_ml_dsa(self.1, msg, &sig[ED_LEN..])
+                    }
+                }
+                self.verify_frame(&DualVerifier(&pk.sig_cl_public_key, &pk.sig_pq_public_key)).is_ok()
+            }
+            _ => false,
+        }
+    }
+
     #[cfg(feature = "registry")]
     pub fn migrate(&self, target_tm: &TypeMap) -> Result<Self, CodecError> {
         let comm_name = communication_type_name(self.comm_type.0)
@@ -544,15 +734,23 @@ impl fmt::Display for CommunicationValue {
             write!(f, ", R:{}{:X}{}", ORANGE, self.receiver, RESET)?;
         }
 
-        let name = communication_type_name(self.comm_type.0).unwrap_or("?");
+        let name = self
+            .get_comm_type_enum()
+            .map(|t| t.name())
+            .unwrap_or_else(|| communication_type_name(self.comm_type.0).unwrap_or("?"));
         write!(f, ", {}: ", name)?;
 
+        let tm = self.type_map.clone().unwrap_or_else(TypeMap::latest);
         write!(f, "{{")?;
-        for (i, (key, value)) in self.data.iter().enumerate() {
+        for (i, (raw_id, value)) in self.data.iter().enumerate() {
             if i > 0 {
                 write!(f, ", ")?;
             }
-            let dname = data_type_name(key.0).unwrap_or("?");
+            let dname = tm
+                .data_enum_id(raw_id.0)
+                .map(|t| t.name())
+                .or_else(|| data_type_name(raw_id.0))
+                .unwrap_or("?");
             write!(f, "{}: ", dname)?;
             fmt_data_value(value, f)?;
         }
@@ -645,11 +843,11 @@ mod tests {
         assert_eq!(decoded.get_receiver(), 222);
         assert_eq!(decoded.get_type(), CommunicationType::Disconnect.to_id(&tm));
         assert_eq!(
-            decoded.get_data(DataType::Id.to_id(&tm)),
+            decoded.get_data(DataType::Id),
             &DataValue::Str("alice".to_string())
         );
         assert_eq!(
-            decoded.get_data(DataType::ClientNonce.to_id(&tm)),
+            decoded.get_data(DataType::ClientNonce),
             &DataValue::SignedNumber(42)
         );
     }

codec/src/data_value.rs

@@ -39,6 +39,28 @@ pub enum DataKind {
     Null,
 }
 
+impl fmt::Display for DataKind {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            DataKind::Bool => f.write_str("Bool"),
+            DataKind::SignedNumber => f.write_str("SignedNumber"),
+            DataKind::UnsignedNumber => f.write_str("UnsignedNumber"),
+            DataKind::Float => f.write_str("Float"),
+            DataKind::Str => f.write_str("Str"),
+            DataKind::Bytes => f.write_str("Bytes"),
+            DataKind::Array(inner) => write!(f, "Array<{}>", inner),
+            DataKind::Container => f.write_str("Container"),
+            #[cfg(feature = "crypto")]
+            DataKind::EncryptedContainer => f.write_str("EncryptedContainer"),
+            #[cfg(feature = "crypto")]
+            DataKind::SignedContainer => f.write_str("SignedContainer"),
+            #[cfg(feature = "crypto")]
+            DataKind::SignedEncryptedContainer => f.write_str("SignedEncryptedContainer"),
+            DataKind::Null => f.write_str("Null"),
+        }
+    }
+}
+
 #[derive(Debug, Clone, Eq)]
 pub enum DataValue {
     BoolTrue,
@@ -259,6 +281,77 @@ impl DataValue {
         }
     }
 
+    pub fn as_number(&self) -> Option<i128> {
+        match self {
+            DataValue::SignedNumber(n) => Some(*n),
+            DataValue::UnsignedNumber(n) => Some(*n as i128),
+            _ => None,
+        }
+    }
+
+    pub fn is_null(&self) -> bool {
+        matches!(self, DataValue::Null)
+    }
+
+    pub fn is_truthy(&self) -> bool {
+        match self {
+            DataValue::BoolTrue | DataValue::Bool(true) => true,
+            DataValue::BoolFalse | DataValue::Bool(false) | DataValue::Null => false,
+            DataValue::UnsignedNumber(0) | DataValue::SignedNumber(0) => false,
+            _ => true,
+        }
+    }
+
+    pub fn get_field(&self, key: DataTypeId) -> Option<&DataValue> {
+        match self {
+            DataValue::Container(entries) => {
+                entries.iter().find(|(k, _)| *k == key).map(|(_, v)| v)
+            }
+            _ => None,
+        }
+    }
+
+    pub fn as_container_map(&self) -> Option<BTreeMap<DataTypeId, DataValue>> {
+        match self {
+            DataValue::Container(entries) => Some(entries.iter().cloned().collect()),
+            _ => None,
+        }
+    }
+
+    pub fn as_bytes_slice(&self) -> Option<&[u8]> {
+        match self {
+            DataValue::Bytes(b) => Some(b),
+            _ => None,
+        }
+    }
+
+    pub fn as_array_slice(&self) -> Option<&[DataValue]> {
+        match self {
+            DataValue::Array(a) => Some(a),
+            _ => None,
+        }
+    }
+
+    pub fn type_name(&self) -> &'static str {
+        match self {
+            DataValue::Bool(_) | DataValue::BoolTrue | DataValue::BoolFalse => "Bool",
+            DataValue::SignedNumber(_) => "SignedNumber",
+            DataValue::UnsignedNumber(_) => "UnsignedNumber",
+            DataValue::Float(_, _) => "Float",
+            DataValue::Str(_) => "Str",
+            DataValue::Bytes(_) => "Bytes",
+            DataValue::Array(_) => "Array",
+            DataValue::Container(_) => "Container",
+            #[cfg(feature = "crypto")]
+            DataValue::EncryptedContainer(_) => "EncryptedContainer",
+            #[cfg(feature = "crypto")]
+            DataValue::SignedContainer(_) => "SignedContainer",
+            #[cfg(feature = "crypto")]
+            DataValue::SignedEncryptedContainer(_) => "SignedEncryptedContainer",
+            DataValue::Null => "Null",
+        }
+    }
+
     #[cfg(feature = "crypto")]
     pub fn as_encrypted_container(&self) -> Option<Vec<u8>> {
         match self {
@@ -374,6 +467,60 @@ impl DataValue {
         Some(())
     }
 
+    /*
+     * Verify a `SignedContainer` without mutating self. Dispatches to
+     * Ed25519, ML-DSA-65, or both (DUAL) based on the algorithm byte
+     * embedded in the blob. Returns `false` for any other variant.
+     */
+    #[cfg(feature = "crypto")]
+    pub fn validate_signature(&self, pk: &PublicKeyBundle) -> bool {
+        let blob = match self {
+            DataValue::SignedContainer(b) => b,
+            _ => return false,
+        };
+        if blob.is_empty() {
+            return false;
+        }
+        let alg = blob[0];
+        let sig_len = match SigAlgorithm::length(alg) {
+            Some(n) => n,
+            None => return false,
+        };
+        if blob.len() < 1 + sig_len + 2 {
+            return false;
+        }
+        let signature = &blob[1..1 + sig_len];
+        let container_bytes = &blob[1 + sig_len..];
+        match alg {
+            SigAlgorithm::ED25519 => {
+                mtp_crypto::verify_ed25519(&pk.sig_cl_public_key, container_bytes, signature).is_ok()
+            }
+            SigAlgorithm::ML_DSA_65 => {
+                mtp_crypto::verify_ml_dsa(&pk.sig_pq_public_key, container_bytes, signature).is_ok()
+            }
+            SigAlgorithm::DUAL => {
+                const ED_LEN: usize = 64;
+                if signature.len() < ED_LEN {
+                    return false;
+                }
+                let ed_ok = mtp_crypto::verify_ed25519(
+                    &pk.sig_cl_public_key,
+                    container_bytes,
+                    &signature[..ED_LEN],
+                )
+                .is_ok();
+                let ml_ok = mtp_crypto::verify_ml_dsa(
+                    &pk.sig_pq_public_key,
+                    container_bytes,
+                    &signature[ED_LEN..],
+                )
+                .is_ok();
+                ed_ok && ml_ok
+            }
+            _ => false,
+        }
+    }
+
     /*
      * Encrypt a `Container` into a `SignedEncryptedContainer` in-place.
      * The container is first signed (with `algorithm`/`signer`), then the signed
@@ -895,6 +1042,133 @@ impl Hash for DataValue {
     }
 }
 
+/* ================================ FROM / TRY-FROM ================================ */
+
+impl From<bool> for DataValue {
+    fn from(v: bool) -> Self {
+        if v { DataValue::BoolTrue } else { DataValue::BoolFalse }
+    }
+}
+
+impl From<&str> for DataValue {
+    fn from(s: &str) -> Self {
+        DataValue::Str(s.to_string())
+    }
+}
+
+impl From<String> for DataValue {
+    fn from(s: String) -> Self {
+        DataValue::Str(s)
+    }
+}
+
+impl From<i64> for DataValue {
+    fn from(n: i64) -> Self {
+        DataValue::SignedNumber(n as i128)
+    }
+}
+
+impl From<i128> for DataValue {
+    fn from(n: i128) -> Self {
+        DataValue::SignedNumber(n)
+    }
+}
+
+impl From<u64> for DataValue {
+    fn from(n: u64) -> Self {
+        DataValue::UnsignedNumber(n as u128)
+    }
+}
+
+impl From<u128> for DataValue {
+    fn from(n: u128) -> Self {
+        DataValue::UnsignedNumber(n)
+    }
+}
+
+impl From<Vec<u8>> for DataValue {
+    fn from(b: Vec<u8>) -> Self {
+        DataValue::Bytes(b)
+    }
+}
+
+impl From<&[u8]> for DataValue {
+    fn from(b: &[u8]) -> Self {
+        DataValue::Bytes(b.to_vec())
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct DataValueTypeMismatch {
+    pub expected: &'static str,
+    pub got: &'static str,
+}
+
+impl fmt::Display for DataValueTypeMismatch {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "expected {}, got {}", self.expected, self.got)
+    }
+}
+
+impl std::error::Error for DataValueTypeMismatch {}
+
+impl TryFrom<DataValue> for bool {
+    type Error = DataValueTypeMismatch;
+    fn try_from(v: DataValue) -> Result<Self, Self::Error> {
+        v.as_bool().ok_or(DataValueTypeMismatch { expected: "Bool", got: v.type_name() })
+    }
+}
+
+impl TryFrom<DataValue> for String {
+    type Error = DataValueTypeMismatch;
+    fn try_from(v: DataValue) -> Result<Self, Self::Error> {
+        match v {
+            DataValue::Str(s) => Ok(s),
+            other => Err(DataValueTypeMismatch { expected: "Str", got: other.type_name() }),
+        }
+    }
+}
+
+impl TryFrom<DataValue> for i128 {
+    type Error = DataValueTypeMismatch;
+    fn try_from(v: DataValue) -> Result<Self, Self::Error> {
+        v.as_signed_number().ok_or(DataValueTypeMismatch { expected: "SignedNumber", got: v.type_name() })
+    }
+}
+
+impl TryFrom<DataValue> for i64 {
+    type Error = DataValueTypeMismatch;
+    fn try_from(v: DataValue) -> Result<Self, Self::Error> {
+        let n = v.as_signed_number().ok_or(DataValueTypeMismatch { expected: "SignedNumber", got: v.type_name() })?;
+        Ok(n as i64)
+    }
+}
+
+impl TryFrom<DataValue> for u128 {
+    type Error = DataValueTypeMismatch;
+    fn try_from(v: DataValue) -> Result<Self, Self::Error> {
+        v.as_unsigned_number().ok_or(DataValueTypeMismatch { expected: "UnsignedNumber", got: v.type_name() })
+    }
+}
+
+impl TryFrom<DataValue> for u64 {
+    type Error = DataValueTypeMismatch;
+    fn try_from(v: DataValue) -> Result<Self, Self::Error> {
+        let n = v.as_unsigned_number().ok_or(DataValueTypeMismatch { expected: "UnsignedNumber", got: v.type_name() })?;
+        Ok(n as u64)
+    }
+}
+
+impl TryFrom<DataValue> for Vec<u8> {
+    type Error = DataValueTypeMismatch;
+    fn try_from(v: DataValue) -> Result<Self, Self::Error> {
+        match v {
+            DataValue::Bytes(b) => Ok(b),
+            other => Err(DataValueTypeMismatch { expected: "Bytes", got: other.type_name() }),
+        }
+    }
+}
+
 /* ================================ TESTS ================================ */
 #[cfg(test)]
 mod tests {
@@ -1265,6 +1539,41 @@ mod tests {
         assert!(s.contains("6:"));
     }
 
+    #[test]
+    fn test_from_primitives() {
+        assert_eq!(DataValue::from(true), DataValue::BoolTrue);
+        assert_eq!(DataValue::from(false), DataValue::BoolFalse);
+        assert_eq!(DataValue::from("hello"), DataValue::Str("hello".to_string()));
+        assert_eq!(DataValue::from("hello".to_string()), DataValue::Str("hello".to_string()));
+        assert_eq!(DataValue::from(42i64), DataValue::SignedNumber(42));
+        assert_eq!(DataValue::from(42i128), DataValue::SignedNumber(42));
+        assert_eq!(DataValue::from(42u64), DataValue::UnsignedNumber(42));
+        assert_eq!(DataValue::from(42u128), DataValue::UnsignedNumber(42));
+        assert_eq!(DataValue::from(vec![1u8, 2, 3]), DataValue::Bytes(vec![1, 2, 3]));
+        assert_eq!(DataValue::from([1u8, 2, 3].as_ref()), DataValue::Bytes(vec![1, 2, 3]));
+    }
+
+    #[test]
+    fn test_try_from_ok() {
+        assert_eq!(bool::try_from(DataValue::BoolTrue).unwrap(), true);
+        assert_eq!(bool::try_from(DataValue::BoolFalse).unwrap(), false);
+        assert_eq!(String::try_from(DataValue::Str("hi".to_string())).unwrap(), "hi");
+        assert_eq!(i128::try_from(DataValue::SignedNumber(-1)).unwrap(), -1i128);
+        assert_eq!(i64::try_from(DataValue::SignedNumber(10)).unwrap(), 10i64);
+        assert_eq!(u128::try_from(DataValue::UnsignedNumber(99)).unwrap(), 99u128);
+        assert_eq!(u64::try_from(DataValue::UnsignedNumber(7)).unwrap(), 7u64);
+        assert_eq!(Vec::<u8>::try_from(DataValue::Bytes(vec![0xAB])).unwrap(), vec![0xABu8]);
+    }
+
+    #[test]
+    fn test_try_from_err() {
+        assert!(bool::try_from(DataValue::Null).is_err());
+        assert!(String::try_from(DataValue::SignedNumber(1)).is_err());
+        assert!(i128::try_from(DataValue::BoolTrue).is_err());
+        assert!(u128::try_from(DataValue::Str("x".to_string())).is_err());
+        assert!(Vec::<u8>::try_from(DataValue::Null).is_err());
+    }
+
     #[test]
     fn test_array_display() {
         let dv = DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]);