xmtp_mls_common/
tls_set.rs

1#![deny(missing_docs)]
2
3//! A deterministic, sorted set with TLS codec serialization.
4//!
5//! [`TlsSet`] is backed by a [`TlsMap<K, ()>`](crate::tls_map::TlsMap), inheriting
6//! deterministic serialization, O(log n) lookup, and O(n) insert/remove.
7//!
8//! Since `()` serializes as zero bytes in TLS codec, the wire format is
9//! effectively `vlen(total_byte_length) || key[0] || key[1] || ...`.
10
11use std::io::{Read, Write};
12
13use tls_codec::{Deserialize, Serialize, Size};
14
15use crate::tls_map::{TlsMap, TlsMapDelta, TlsMapError, TlsMapMutation};
16
17/// A SHA-256 hash of a key's TLS-serialized form, used as the lookup token
18/// for [`TlsSetMutation::RemoveByHash`].
19///
20/// This is a newtype around `[u8; 32]` to prevent accidental mixing with
21/// other 32-byte SHA-256 hashes elsewhere in the codebase (commit hashes,
22/// installation IDs, etc.).
23///
24/// The only safe way to construct one is via [`TlsKeyHash::of`], which
25/// guarantees the hash always uses the canonical TLS serialization of the
26/// key. Raw byte construction is reserved for wire-format deserialization.
27#[derive(Clone, Copy, PartialEq, Eq, Hash)]
28pub struct TlsKeyHash([u8; 32]);
29
30impl TlsKeyHash {
31    /// Compute the hash of a key by SHA-256-ing its TLS-serialized form.
32    /// This is the only public way to construct a `TlsKeyHash`, ensuring
33    /// every hash on the wire was produced by the canonical encoding.
34    #[inline]
35    pub fn of<K: Serialize + Size>(key: &K) -> Result<Self, tls_codec::Error> {
36        let bytes = key.tls_serialize_detached()?;
37        Ok(Self(xmtp_common::sha256_array(&bytes)))
38    }
39
40    /// Construct from raw bytes. Used by TLS deserialization to reconstruct
41    /// a hash that was produced by an earlier `TlsKeyHash::of` on the sender.
42    #[inline]
43    pub(crate) const fn from_bytes(bytes: [u8; 32]) -> Self {
44        Self(bytes)
45    }
46
47    /// Borrow the underlying bytes.
48    #[inline]
49    pub const fn as_bytes(&self) -> &[u8; 32] {
50        &self.0
51    }
52}
53
54impl std::fmt::Debug for TlsKeyHash {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        f.write_str("TlsKeyHash(")?;
57        for byte in &self.0 {
58            write!(f, "{byte:02x}")?;
59        }
60        f.write_str(")")
61    }
62}
63
64/// Error type for [`TlsSet`] operations.
65#[derive(Debug, thiserror::Error)]
66pub enum TlsSetError {
67    /// Two or more keys in the set produce the same SHA-256 digest.
68    ///
69    /// Returned by [`TlsSet::apply_delta`] when building the hash index for
70    /// `RemoveByHash` lookups. SHA-256 collisions are cryptographically
71    /// infeasible to find by chance, so in practice this signals either:
72    /// - A bug in `Serialize` for `K` producing identical bytes for distinct
73    ///   logically-different keys (e.g., a non-deterministic encoding), or
74    /// - A deliberate adversarial attempt to seed the set with crafted entries.
75    ///
76    /// We surface it as a defense-in-depth guarantee: `RemoveByHash` resolves
77    /// to exactly one key, never silently picking one of several candidates.
78    #[error("duplicate hash in set")]
79    DuplicateHash,
80    /// An underlying [`TlsMap`] error.
81    #[error(transparent)]
82    Map(#[from] TlsMapError),
83    /// A TLS codec serialization error (from hashing a key).
84    #[error("tls codec error: {0}")]
85    Codec(#[from] tls_codec::Error),
86}
87
88/// A sorted set with deterministic TLS codec serialization.
89///
90/// # Examples
91///
92/// ```
93/// use xmtp_mls_common::tls_set::TlsSet;
94/// use tls_codec::{Serialize, Deserialize};
95///
96/// let mut set = TlsSet::<u16>::new();
97/// set.insert(3).unwrap();
98/// set.insert(1).unwrap();
99/// set.insert(2).unwrap();
100///
101/// let bytes = set.tls_serialize_detached().unwrap();
102/// let deserialized = TlsSet::<u16>::tls_deserialize_exact(&bytes).unwrap();
103/// assert_eq!(set, deserialized);
104/// ```
105#[derive(Debug, Clone, PartialEq, Eq)]
106pub struct TlsSet<K> {
107    inner: TlsMap<K, ()>,
108}
109
110impl<K> Default for TlsSet<K> {
111    #[inline]
112    fn default() -> Self {
113        Self {
114            inner: TlsMap::default(),
115        }
116    }
117}
118
119impl<K: Ord + Eq> TlsSet<K> {
120    /// Create an empty set.
121    #[inline]
122    pub fn new() -> Self {
123        Self::default()
124    }
125
126    /// Create a set from an iterator of keys. Duplicates are silently ignored.
127    ///
128    /// ```
129    /// use xmtp_mls_common::tls_set::TlsSet;
130    ///
131    /// let set = TlsSet::from_keys([3, 1, 2, 1]);
132    /// assert_eq!(set.len(), 3);
133    /// ```
134    #[inline]
135    pub fn from_keys(iter: impl IntoIterator<Item = K>) -> Self {
136        Self {
137            inner: TlsMap::from_pairs(iter.into_iter().map(|k| (k, ()))),
138        }
139    }
140
141    /// Insert a key. Returns an error if the key already exists.
142    ///
143    /// ```
144    /// use xmtp_mls_common::tls_set::TlsSet;
145    ///
146    /// let mut set = TlsSet::<u8>::new();
147    /// assert!(set.insert(1).is_ok());
148    /// assert!(set.insert(1).is_err()); // duplicate
149    /// ```
150    #[inline]
151    pub fn insert(&mut self, key: K) -> Result<(), TlsSetError> {
152        Ok(self.inner.insert(key, ())?)
153    }
154
155    /// Remove a key. Returns an error if the key doesn't exist.
156    ///
157    /// ```
158    /// use xmtp_mls_common::tls_set::TlsSet;
159    ///
160    /// let mut set = TlsSet::<u8>::new();
161    /// set.insert(1).unwrap();
162    /// assert!(set.remove(&1).is_ok());
163    /// assert!(set.remove(&1).is_err()); // already removed
164    /// ```
165    #[inline]
166    pub fn remove(&mut self, key: &K) -> Result<(), TlsSetError> {
167        self.inner.remove(key).map(|_| ())?;
168        Ok(())
169    }
170
171    /// Returns true if the set contains the key.
172    #[inline]
173    pub fn contains(&self, key: &K) -> bool {
174        self.inner.contains_key(key)
175    }
176
177    /// Returns the number of elements in the set.
178    #[inline]
179    pub fn len(&self) -> usize {
180        self.inner.len()
181    }
182
183    /// Returns true if the set is empty.
184    #[inline]
185    pub fn is_empty(&self) -> bool {
186        self.inner.is_empty()
187    }
188
189    /// Iterate over keys in sorted order.
190    ///
191    /// ```
192    /// use xmtp_mls_common::tls_set::TlsSet;
193    ///
194    /// let set = TlsSet::from_keys([3, 1, 2]);
195    /// let keys: Vec<_> = set.iter().collect();
196    /// assert_eq!(keys, vec![&1, &2, &3]);
197    /// ```
198    #[inline]
199    pub fn iter(&self) -> impl Iterator<Item = &K> {
200        self.inner.keys()
201    }
202}
203
204impl<K: Ord + Eq + Clone + Serialize + Size> TlsSet<K> {
205    /// Apply a delta atomically. If any mutation fails, the set is unchanged.
206    ///
207    /// If the delta contains any `RemoveByHash` mutations, a hash index of all
208    /// existing keys is built once (O(n)) and used for lookups (O(1) each),
209    /// avoiding O(n*m) behavior. Returns [`TlsSetError::DuplicateHash`] if two
210    /// existing keys produce the same SHA-256 hash.
211    ///
212    /// # Performance
213    ///
214    /// - **Time:** O(n + m) when the delta contains any `RemoveByHash`
215    ///   (one O(n) index build, then O(1) per lookup), or O(m) otherwise,
216    ///   where `n` is the current set size and `m` is the mutation count.
217    /// - **Memory:** O(m) — all mutations are resolved into a temporary
218    ///   `Vec<TlsMapMutation>` before being applied. This allocation is
219    ///   required to preserve atomicity: resolution must complete before
220    ///   any actual mutation happens, so a partial delta cannot leave the
221    ///   set in an inconsistent state. When `RemoveByHash` is present, an
222    ///   additional O(n) `HashMap` is allocated for the hash index, holding
223    ///   borrowed key references (no key cloning during index build).
224    ///
225    /// ```
226    /// use xmtp_mls_common::tls_set::{TlsSet, TlsSetDelta};
227    ///
228    /// let mut set = TlsSet::<u8>::new();
229    /// set.insert(1).unwrap();
230    ///
231    /// let delta = TlsSetDelta::new().insert(2).remove(1);
232    /// set.apply_delta(delta).unwrap();
233    /// assert!(set.contains(&2));
234    /// assert!(!set.contains(&1));
235    /// ```
236    pub fn apply_delta(&mut self, delta: TlsSetDelta<K>) -> Result<(), TlsSetError> {
237        let has_hash_removal = delta
238            .mutations
239            .iter()
240            .any(|m| matches!(m, TlsSetMutation::RemoveByHash(_)));
241
242        // Build hash → &K index only when needed — O(n) once instead of
243        // O(n) per removal. Stores borrowed references, no key cloning.
244        let hash_index: Option<std::collections::HashMap<TlsKeyHash, &K>> = if has_hash_removal {
245            let mut idx = std::collections::HashMap::with_capacity(self.inner.len());
246            for key in self.inner.keys() {
247                if idx.insert(TlsKeyHash::of(key)?, key).is_some() {
248                    return Err(TlsSetError::DuplicateHash);
249                }
250            }
251            Some(idx)
252        } else {
253            None
254        };
255
256        let resolved: Vec<TlsMapMutation<K, ()>> = delta
257            .mutations
258            .into_iter()
259            .map(|m| -> Result<TlsMapMutation<K, ()>, TlsSetError> {
260                match m {
261                    TlsSetMutation::Insert(key) => Ok(TlsMapMutation::Insert { key, value: () }),
262                    TlsSetMutation::Remove(key) => Ok(TlsMapMutation::Delete { key }),
263                    TlsSetMutation::RemoveByHash(target_hash) => {
264                        let idx = hash_index
265                            .as_ref()
266                            .expect("hash index must be built for RemoveByHash");
267                        let key = idx.get(&target_hash).ok_or(TlsMapError::KeyNotFound)?;
268                        Ok(TlsMapMutation::Delete {
269                            key: (*key).clone(),
270                        })
271                    }
272                }
273            })
274            .collect::<Result<Vec<_>, _>>()?;
275        let map_delta = TlsMapDelta {
276            mutations: resolved,
277        };
278        self.inner.apply_delta(map_delta)?;
279        Ok(())
280    }
281}
282
283// -- TLS codec --
284
285impl<K: Size> Size for TlsSet<K> {
286    #[inline]
287    fn tls_serialized_len(&self) -> usize {
288        self.inner.tls_serialized_len()
289    }
290}
291
292impl<K: Serialize + Size + std::fmt::Debug> Serialize for TlsSet<K> {
293    #[inline]
294    fn tls_serialize<W: Write>(&self, writer: &mut W) -> Result<usize, tls_codec::Error> {
295        self.inner.tls_serialize(writer)
296    }
297}
298
299impl<K> Deserialize for TlsSet<K>
300where
301    K: Deserialize + Size + Ord + Eq,
302{
303    #[inline]
304    fn tls_deserialize<R: Read>(bytes: &mut R) -> Result<Self, tls_codec::Error>
305    where
306        Self: Sized,
307    {
308        let inner = TlsMap::tls_deserialize(bytes)?;
309        Ok(Self { inner })
310    }
311}
312
313impl<K: Ord + Eq> FromIterator<K> for TlsSet<K> {
314    #[inline]
315    fn from_iter<T: IntoIterator<Item = K>>(iter: T) -> Self {
316        Self::from_keys(iter)
317    }
318}
319
320impl<K> IntoIterator for TlsSet<K> {
321    type Item = K;
322    type IntoIter = std::iter::Map<<TlsMap<K, ()> as IntoIterator>::IntoIter, fn((K, ())) -> K>;
323
324    #[inline]
325    fn into_iter(self) -> Self::IntoIter {
326        self.inner.into_iter().map(|(k, _)| k)
327    }
328}
329
330// ============================================================================
331// TlsSetDelta — atomic batch mutations for TlsSet
332// ============================================================================
333
334/// A mutation to apply to a [`TlsSet`].
335///
336/// Wire format: `u8(tag) || payload` where:
337/// - tag 0 = Insert: `K(tls)`
338/// - tag 1 = Remove: `K(tls)`
339/// - tag 2 = RemoveByHash: `[u8; 32]` (SHA-256 of the key's TLS serialization)
340#[derive(Debug, Clone, PartialEq, Eq)]
341pub enum TlsSetMutation<K> {
342    /// Insert a key. Fails if the key already exists.
343    Insert(K),
344    /// Remove a key by value. Fails if the key doesn't exist.
345    Remove(K),
346    /// Remove a key by the SHA-256 hash of its TLS serialization.
347    /// Avoids sending large keys over the wire for removal.
348    /// Fails if no key with a matching hash exists.
349    RemoveByHash(TlsKeyHash),
350}
351
352/// A batch of mutations to apply atomically to a [`TlsSet`].
353#[derive(Debug, Clone, PartialEq, Eq)]
354pub struct TlsSetDelta<K> {
355    /// The ordered list of mutations to apply.
356    pub mutations: Vec<TlsSetMutation<K>>,
357}
358
359impl<K> TlsSetDelta<K> {
360    /// Create an empty delta.
361    #[inline]
362    pub fn new() -> Self {
363        Self {
364            mutations: Vec::new(),
365        }
366    }
367
368    /// Append an insert mutation.
369    #[inline]
370    pub fn insert(mut self, key: K) -> Self {
371        self.mutations.push(TlsSetMutation::Insert(key));
372        self
373    }
374
375    /// Append a remove mutation.
376    #[inline]
377    pub fn remove(mut self, key: K) -> Self {
378        self.mutations.push(TlsSetMutation::Remove(key));
379        self
380    }
381
382    /// Append a remove-by-hash mutation.
383    #[inline]
384    pub fn remove_by_hash(mut self, hash: TlsKeyHash) -> Self {
385        self.mutations.push(TlsSetMutation::RemoveByHash(hash));
386        self
387    }
388}
389
390impl<K> Default for TlsSetDelta<K> {
391    #[inline]
392    fn default() -> Self {
393        Self::new()
394    }
395}
396
397// -- TLS codec for delta types --
398
399impl<K: Size> Size for TlsSetMutation<K> {
400    #[inline]
401    fn tls_serialized_len(&self) -> usize {
402        1 + match self {
403            Self::Insert(key) | Self::Remove(key) => key.tls_serialized_len(),
404            Self::RemoveByHash(_) => 32,
405        }
406    }
407}
408
409impl<K: Serialize + Size> Serialize for TlsSetMutation<K> {
410    #[inline]
411    fn tls_serialize<W: Write>(&self, writer: &mut W) -> Result<usize, tls_codec::Error> {
412        match self {
413            Self::Insert(key) => {
414                let mut written = 0_u8.tls_serialize(writer)?;
415                written += key.tls_serialize(writer)?;
416                Ok(written)
417            }
418            Self::Remove(key) => {
419                let mut written = 1_u8.tls_serialize(writer)?;
420                written += key.tls_serialize(writer)?;
421                Ok(written)
422            }
423            Self::RemoveByHash(hash) => {
424                let written = 2_u8.tls_serialize(writer)?;
425                let bytes = hash.as_bytes();
426                writer
427                    .write_all(bytes)
428                    .map_err(|e| tls_codec::Error::EncodingError(e.to_string()))?;
429                Ok(written + bytes.len())
430            }
431        }
432    }
433}
434
435impl<K: Deserialize + Size> Deserialize for TlsSetMutation<K> {
436    #[inline]
437    fn tls_deserialize<R: Read>(bytes: &mut R) -> Result<Self, tls_codec::Error>
438    where
439        Self: Sized,
440    {
441        let tag = u8::tls_deserialize(bytes)?;
442        match tag {
443            0 => Ok(Self::Insert(K::tls_deserialize(bytes)?)),
444            1 => Ok(Self::Remove(K::tls_deserialize(bytes)?)),
445            2 => {
446                let mut buf = [0_u8; 32];
447                bytes
448                    .read_exact(&mut buf)
449                    .map_err(|e| tls_codec::Error::DecodingError(e.to_string()))?;
450                Ok(Self::RemoveByHash(TlsKeyHash::from_bytes(buf)))
451            }
452            _ => Err(tls_codec::Error::DecodingError(format!(
453                "unknown TlsSetMutation tag: {tag}"
454            ))),
455        }
456    }
457}
458
459impl<K: Size> Size for TlsSetDelta<K> {
460    #[inline]
461    fn tls_serialized_len(&self) -> usize {
462        self.mutations.tls_serialized_len()
463    }
464}
465
466impl<K: Serialize + Size + std::fmt::Debug> Serialize for TlsSetDelta<K> {
467    #[inline]
468    fn tls_serialize<W: Write>(&self, writer: &mut W) -> Result<usize, tls_codec::Error> {
469        self.mutations.tls_serialize(writer)
470    }
471}
472
473impl<K: Deserialize + Size> Deserialize for TlsSetDelta<K> {
474    #[inline]
475    fn tls_deserialize<R: Read>(bytes: &mut R) -> Result<Self, tls_codec::Error>
476    where
477        Self: Sized,
478    {
479        let mutations = Vec::<TlsSetMutation<K>>::tls_deserialize(bytes)?;
480        Ok(Self { mutations })
481    }
482}
483
484#[cfg(test)]
485mod tests {
486    use super::*;
487    use tls_codec::{Deserialize, Serialize};
488
489    #[xmtp_common::test]
490    fn test_insert_and_contains() {
491        let mut set = TlsSet::<u16>::new();
492        set.insert(42).unwrap();
493        assert!(set.contains(&42));
494        assert!(!set.contains(&99));
495    }
496
497    #[xmtp_common::test]
498    fn test_insert_duplicate_fails() {
499        let mut set = TlsSet::<u8>::new();
500        set.insert(1).unwrap();
501        assert!(set.insert(1).is_err());
502    }
503
504    #[xmtp_common::test]
505    fn test_remove() {
506        let mut set = TlsSet::<u8>::new();
507        set.insert(1).unwrap();
508        set.remove(&1).unwrap();
509        assert!(!set.contains(&1));
510        assert!(set.is_empty());
511    }
512
513    #[xmtp_common::test]
514    fn test_remove_missing_fails() {
515        let mut set = TlsSet::<u8>::new();
516        assert!(set.remove(&1).is_err());
517    }
518
519    #[xmtp_common::test]
520    fn test_from_keys_deduplicates() {
521        let set = TlsSet::from_keys([1_u8, 2, 3, 1, 2]);
522        assert_eq!(set.len(), 3);
523    }
524
525    #[xmtp_common::test]
526    fn test_iter_sorted() {
527        let set = TlsSet::from_keys([3_u8, 1, 2]);
528        let keys: Vec<_> = set.iter().copied().collect();
529        assert_eq!(keys, vec![1, 2, 3]);
530    }
531
532    #[xmtp_common::test]
533    fn test_tls_round_trip() {
534        let set = TlsSet::from_keys([10_u16, 20, 30]);
535        let bytes = set.tls_serialize_detached().unwrap();
536        let restored = TlsSet::<u16>::tls_deserialize_exact(&bytes).unwrap();
537        assert_eq!(set, restored);
538    }
539
540    #[xmtp_common::test]
541    fn test_empty_round_trip() {
542        let set = TlsSet::<u8>::new();
543        let bytes = set.tls_serialize_detached().unwrap();
544        let restored = TlsSet::<u8>::tls_deserialize_exact(&bytes).unwrap();
545        assert_eq!(set, restored);
546        assert!(restored.is_empty());
547    }
548
549    #[xmtp_common::test]
550    fn test_into_iter() {
551        let set = TlsSet::from_keys([3_u8, 1, 2]);
552        let keys: Vec<_> = set.into_iter().collect();
553        assert_eq!(keys, vec![1, 2, 3]);
554    }
555
556    #[xmtp_common::test]
557    fn test_collect() {
558        let set: TlsSet<u8> = [3, 1, 2].into_iter().collect();
559        assert_eq!(set.len(), 3);
560        assert!(set.contains(&1));
561    }
562
563    #[xmtp_common::test]
564    fn test_apply_delta() {
565        let mut set = TlsSet::from_keys([1_u8, 2, 3]);
566        let delta = TlsSetDelta::new().insert(4).remove(2);
567        set.apply_delta(delta).unwrap();
568        assert!(set.contains(&1));
569        assert!(!set.contains(&2));
570        assert!(set.contains(&3));
571        assert!(set.contains(&4));
572    }
573
574    #[xmtp_common::test]
575    fn test_apply_delta_rollback_on_failure() {
576        let mut set = TlsSet::from_keys([1_u8, 2]);
577        // Try to add 3 then add 1 (duplicate) — should rollback
578        let delta = TlsSetDelta::new().insert(3).insert(1);
579        assert!(set.apply_delta(delta).is_err());
580        // Set should be unchanged
581        assert_eq!(set.len(), 2);
582        assert!(!set.contains(&3));
583    }
584
585    #[xmtp_common::test]
586    fn test_remove_by_hash() {
587        let mut set = TlsSet::from_keys([10_u16, 20, 30]);
588        let hash = TlsKeyHash::of(&20_u16).unwrap();
589        let delta = TlsSetDelta::new().remove_by_hash(hash);
590        set.apply_delta(delta).unwrap();
591        assert!(set.contains(&10));
592        assert!(!set.contains(&20));
593        assert!(set.contains(&30));
594    }
595
596    #[xmtp_common::test]
597    fn test_remove_by_hash_not_found() {
598        let mut set = TlsSet::from_keys([10_u16, 20]);
599        let hash = TlsKeyHash::of(&99_u16).unwrap(); // not in set
600        let delta = TlsSetDelta::new().remove_by_hash(hash);
601        assert!(set.apply_delta(delta).is_err());
602    }
603
604    #[xmtp_common::test]
605    fn test_remove_by_hash_matches_remove_by_value() {
606        let mut set_a = TlsSet::from_keys([1_u16, 2, 3]);
607        let mut set_b = set_a.clone();
608
609        // Remove by value
610        let delta_a = TlsSetDelta::new().remove(2);
611        set_a.apply_delta(delta_a).unwrap();
612
613        // Remove by hash
614        let hash = TlsKeyHash::of(&2_u16).unwrap();
615        let delta_b = TlsSetDelta::new().remove_by_hash(hash);
616        set_b.apply_delta(delta_b).unwrap();
617
618        assert_eq!(set_a, set_b);
619    }
620
621    #[xmtp_common::test]
622    fn test_delta_tls_round_trip() {
623        let delta = TlsSetDelta::<u16>::new().insert(10).remove(20).insert(30);
624        let bytes = delta.tls_serialize_detached().unwrap();
625        let restored = TlsSetDelta::<u16>::tls_deserialize_exact(&bytes).unwrap();
626        assert_eq!(delta, restored);
627    }
628
629    #[xmtp_common::test]
630    fn test_remove_by_hash_mutation_tls_round_trip() {
631        let hash = TlsKeyHash::of(&42_u16).unwrap();
632        let mutation = TlsSetMutation::<u16>::RemoveByHash(hash);
633        let bytes = mutation.tls_serialize_detached().unwrap();
634        let restored = TlsSetMutation::<u16>::tls_deserialize_exact(&bytes).unwrap();
635        assert_eq!(mutation, restored);
636    }
637
638    #[xmtp_common::test]
639    fn test_mutation_tls_round_trip() {
640        let add = TlsSetMutation::Insert(42_u16);
641        let bytes = add.tls_serialize_detached().unwrap();
642        let restored = TlsSetMutation::<u16>::tls_deserialize_exact(&bytes).unwrap();
643        assert_eq!(add, restored);
644
645        let remove = TlsSetMutation::Remove(99_u16);
646        let bytes = remove.tls_serialize_detached().unwrap();
647        let restored = TlsSetMutation::<u16>::tls_deserialize_exact(&bytes).unwrap();
648        assert_eq!(remove, restored);
649    }
650
651    #[xmtp_common::test]
652    fn test_deserialize_unknown_tag() {
653        // Tag 3 is not a valid TlsSetMutation variant (only 0/1/2 are defined).
654        let bytes = [3_u8, 0, 42];
655        let result = TlsSetMutation::<u16>::tls_deserialize_exact(bytes);
656        assert!(matches!(result, Err(tls_codec::Error::DecodingError(_))));
657    }
658
659    /// A key type with a deliberately buggy `Serialize` impl that produces
660    /// identical bytes for distinct logical values. Used to exercise the
661    /// duplicate-hash detection path in `apply_delta`, since SHA-256 collisions
662    /// cannot be found by chance with well-formed inputs.
663    #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
664    struct CollidingKey(u16);
665
666    impl Size for CollidingKey {
667        fn tls_serialized_len(&self) -> usize {
668            1
669        }
670    }
671
672    impl Serialize for CollidingKey {
673        fn tls_serialize<W: std::io::Write>(
674            &self,
675            writer: &mut W,
676        ) -> Result<usize, tls_codec::Error> {
677            // Buggy: always serialize as the same byte regardless of value.
678            0_u8.tls_serialize(writer)
679        }
680    }
681
682    impl Deserialize for CollidingKey {
683        fn tls_deserialize<R: std::io::Read>(bytes: &mut R) -> Result<Self, tls_codec::Error>
684        where
685            Self: Sized,
686        {
687            let _ = u8::tls_deserialize(bytes)?;
688            Ok(CollidingKey(0))
689        }
690    }
691
692    #[xmtp_common::test]
693    fn test_apply_delta_duplicate_hash() {
694        let mut set = TlsSet::<CollidingKey>::new();
695        // Distinct logical keys but identical TLS serialization → identical hashes.
696        set.insert(CollidingKey(1)).unwrap();
697        set.insert(CollidingKey(2)).unwrap();
698
699        let any_hash = TlsKeyHash::of(&CollidingKey(0)).unwrap();
700        let delta = TlsSetDelta::new().remove_by_hash(any_hash);
701        let result = set.apply_delta(delta);
702        assert!(matches!(result, Err(TlsSetError::DuplicateHash)));
703    }
704
705    #[xmtp_common::test]
706    fn test_apply_delta_remove_by_hash_not_found_in_index() {
707        // Build the hash index successfully (no collisions), then look up a
708        // hash that doesn't match any key. This exercises the
709        // `idx.get(...).ok_or(KeyNotFound)?` path inside the resolve loop.
710        let mut set = TlsSet::<u16>::new();
711        set.insert(10).unwrap();
712        set.insert(20).unwrap();
713
714        let missing_hash = TlsKeyHash::of(&999_u16).unwrap();
715        let delta = TlsSetDelta::new().remove_by_hash(missing_hash);
716        let result = set.apply_delta(delta);
717        assert!(matches!(
718            result,
719            Err(TlsSetError::Map(TlsMapError::KeyNotFound))
720        ));
721        // Set is unchanged.
722        assert!(set.contains(&10));
723        assert!(set.contains(&20));
724    }
725}
xmtp_mls_common/tls_set.rs

xmtp_mls_common/
tls_set.rs
