From 67ecf50949d29d388b490cdfdbac34d74bef2d3f Mon Sep 17 00:00:00 2001
From: Andrew Plaza <github.tech@liquidthink.net>
Date: Thu, 8 Aug 2024 19:19:34 -0400
Subject: [PATCH] skip empty key packages (#950)

skip empty keypackages to ensure TLS Deserialization does not fail
---
 xmtp_mls/src/api/mls.rs | 49 ++++++++++++++++++++++++++++++++++++++---
 1 file changed, 46 insertions(+), 3 deletions(-)

diff --git a/xmtp_mls/src/api/mls.rs b/xmtp_mls/src/api/mls.rs
index eaf3af6b4..7520eb403 100644
--- a/xmtp_mls/src/api/mls.rs
+++ b/xmtp_mls/src/api/mls.rs
@@ -236,11 +236,19 @@ where
             .key_packages
             .into_iter()
             .enumerate()
-            .map(|(idx, key_package)| {
-                (
+            .filter_map(|(idx, key_package)| {
+                if key_package.key_package_tls_serialized.is_empty() {
+                    log::warn!(
+                        "installation key 0x{} has empty key package",
+                        hex::encode(&installation_keys[idx])
+                    );
+                    return None;
+                }
+
+                Some((
                     installation_keys[idx].to_vec(),
                     key_package.key_package_tls_serialized,
-                )
+                ))
             })
             .collect();
 
@@ -683,4 +691,39 @@ pub mod tests {
             .unwrap();
         assert_eq!(result.len(), 50);
     }
+
+    #[tokio::test]
+    async fn test_fetch_key_packages_skips_empty() {
+        let mut mock_api = MockApiClient::new();
+        let installation_keys: Vec<Vec<u8>> = vec![vec![1, 2, 3], vec![4, 5, 6], vec![1, 1, 1]];
+        mock_api.expect_fetch_key_packages().returning(move |_| {
+            Ok(FetchKeyPackagesResponse {
+                key_packages: vec![
+                    KeyPackage {
+                        key_package_tls_serialized: vec![7, 8, 9],
+                    },
+                    KeyPackage {
+                        key_package_tls_serialized: vec![],
+                    },
+                    KeyPackage {
+                        key_package_tls_serialized: vec![10, 11, 12],
+                    },
+                ],
+            })
+        });
+        let wrapper = ApiClientWrapper::new(mock_api, Retry::default());
+        let result = wrapper
+            .fetch_key_packages(installation_keys.clone())
+            .await
+            .unwrap();
+        assert_eq!(result.len(), 2);
+
+        for (k, v) in result {
+            if k.eq(&installation_keys[0]) {
+                assert_eq!(v, vec![7, 8, 9]);
+            } else {
+                assert_eq!(v, vec![10, 11, 12]);
+            }
+        }
+    }
 }