From c122549e0d587f1297f81ad5f432deda7d139913 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Damir=20Jeli=C4=87?= <poljar@termina.org.uk>
Date: Tue, 18 May 2021 08:29:10 +0200
Subject: [PATCH] base: Correctly get the user ids of all room members

---
 matrix_sdk_base/src/store/sled_store/mod.rs | 41 ++++++++++++++-------
 1 file changed, 28 insertions(+), 13 deletions(-)

diff --git a/matrix_sdk_base/src/store/sled_store/mod.rs b/matrix_sdk_base/src/store/sled_store/mod.rs
index 1991c610..c60534bc 100644
--- a/matrix_sdk_base/src/store/sled_store/mod.rs
+++ b/matrix_sdk_base/src/store/sled_store/mod.rs
@@ -83,8 +83,9 @@ impl From<SerializationError> for StoreError {
     }
 }
 
+const ENCODE_SEPARATOR: u8 = 0xff;
+
 trait EncodeKey {
-    const SEPARATOR: u8 = 0xff;
     fn encode(&self) -> Vec<u8>;
 }
 
@@ -102,13 +103,13 @@ impl EncodeKey for &RoomId {
 
 impl EncodeKey for &str {
     fn encode(&self) -> Vec<u8> {
-        [self.as_bytes(), &[Self::SEPARATOR]].concat()
+        [self.as_bytes(), &[ENCODE_SEPARATOR]].concat()
     }
 }
 
 impl EncodeKey for (&str, &str) {
     fn encode(&self) -> Vec<u8> {
-        [self.0.as_bytes(), &[Self::SEPARATOR], self.1.as_bytes(), &[Self::SEPARATOR]].concat()
+        [self.0.as_bytes(), &[ENCODE_SEPARATOR], self.1.as_bytes(), &[ENCODE_SEPARATOR]].concat()
     }
 }
 
@@ -116,11 +117,11 @@ impl EncodeKey for (&str, &str, &str) {
     fn encode(&self) -> Vec<u8> {
         [
             self.0.as_bytes(),
-            &[Self::SEPARATOR],
+            &[ENCODE_SEPARATOR],
             self.1.as_bytes(),
-            &[Self::SEPARATOR],
+            &[ENCODE_SEPARATOR],
             self.2.as_bytes(),
-            &[Self::SEPARATOR],
+            &[ENCODE_SEPARATOR],
         ]
         .concat()
     }
@@ -506,11 +507,22 @@ impl SledStore {
             .transpose()?)
     }
 
-    pub async fn get_user_ids(&self, room_id: &RoomId) -> impl Stream<Item = Result<UserId>> {
-        stream::iter(self.members.scan_prefix(room_id.encode()).map(|u| {
-            UserId::try_from(String::from_utf8_lossy(&u?.1).to_string())
-                .map_err(StoreError::Identifier)
-        }))
+    pub async fn get_user_ids_stream(
+        &self,
+        room_id: &RoomId,
+    ) -> impl Stream<Item = Result<UserId>> {
+        let decode = |key: &[u8]| -> Result<UserId> {
+            let mut iter = key.split(|c| c == &ENCODE_SEPARATOR);
+            // Our key is a the room id separated from the user id by a null
+            // byte, discard the first value of the split.
+            iter.next();
+
+            let user_id = iter.next().expect("User ids weren't properly encoded");
+
+            Ok(UserId::try_from(String::from_utf8_lossy(user_id).to_string())?)
+        };
+
+        stream::iter(self.members.scan_prefix(room_id.encode()).map(move |u| decode(&u?.0)))
     }
 
     pub async fn get_invited_user_ids(
@@ -636,7 +648,7 @@ impl StateStore for SledStore {
     }
 
     async fn get_user_ids(&self, room_id: &RoomId) -> Result<Vec<UserId>> {
-        self.get_user_ids(room_id).await.try_collect().await
+        self.get_user_ids_stream(room_id).await.try_collect().await
     }
 
     async fn get_invited_user_ids(&self, room_id: &RoomId) -> Result<Vec<UserId>> {
@@ -698,7 +710,7 @@ mod test {
     use serde_json::json;
 
     use super::{SledStore, StateChanges};
-    use crate::deserialized_responses::MemberEvent;
+    use crate::{deserialized_responses::MemberEvent, StateStore};
 
     fn user_id() -> UserId {
         user_id!("@example:localhost")
@@ -748,6 +760,9 @@ mod test {
 
         store.save_changes(&changes).await.unwrap();
         assert!(store.get_member_event(&room_id, &user_id).await.unwrap().is_some());
+
+        let members = store.get_user_ids(&room_id).await.unwrap();
+        assert!(!members.is_empty())
     }
 
     #[async_test]