From 1dbde0e1c17bbedd8631a210cad1b6f211acb7b2 Mon Sep 17 00:00:00 2001
From: timokoesters <timo@koesters.xyz>
Date: Sun, 10 May 2020 18:30:12 +0200
Subject: [PATCH] improvement: add option to get device id from token

---
 src/client_server.rs  |  2 +-
 src/database.rs       |  4 ++--
 src/database/users.rs | 45 ++++++++++++++++++++++++++-----------------
 src/ruma_wrapper.rs   |  9 ++++++---
 4 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/src/client_server.rs b/src/client_server.rs
index d740ab50e..183abbfe4 100644
--- a/src/client_server.rs
+++ b/src/client_server.rs
@@ -105,7 +105,7 @@ pub fn register_route(
                 stages: vec!["m.login.dummy".to_owned()],
             }],
             completed: vec![],
-            params: RawValue::from_string("".to_owned()).unwrap(),
+            params: RawValue::from_string("{}".to_owned()).unwrap(),
             session: Some(utils::random_string(SESSION_ID_LENGTH)),
             auth_error: None,
         })));
diff --git a/src/database.rs b/src/database.rs
index 3b8f9279c..4e6ac57e1 100644
--- a/src/database.rs
+++ b/src/database.rs
@@ -44,11 +44,11 @@ pub fn load_or_create(hostname: &str) -> Self {
                 userid_displayname: db.open_tree("userid_displayname").unwrap(),
                 userid_avatarurl: db.open_tree("userid_avatarurl").unwrap(),
                 userdeviceid_token: db.open_tree("userdeviceid_token").unwrap(),
-                token_userid: db.open_tree("token_userid").unwrap(),
+                token_userdeviceid: db.open_tree("token_userdeviceid").unwrap(),
             },
             rooms: rooms::Rooms {
                 edus: rooms::RoomEdus {
-                    roomuserid_lastread: db.open_tree("roomuserid_lastread").unwrap(),
+                    roomuserid_lastread: db.open_tree("roomuserid_lastread").unwrap(), // "Private" read receipt
                     roomlatestid_roomlatest: db.open_tree("roomlatestid_roomlatest").unwrap(), // Read receipts
                     roomactiveid_roomactive: db.open_tree("roomactiveid_roomactive").unwrap(), // Typing notifs
                 },
diff --git a/src/database/users.rs b/src/database/users.rs
index 529da91d1..5a8b0aa63 100644
--- a/src/database/users.rs
+++ b/src/database/users.rs
@@ -8,7 +8,7 @@ pub struct Users {
     pub(super) userid_avatarurl: sled::Tree,
     pub(super) userdeviceids: sled::Tree,
     pub(super) userdeviceid_token: sled::Tree,
-    pub(super) token_userid: sled::Tree,
+    pub(super) token_userdeviceid: sled::Tree,
 }
 
 impl Users {
@@ -24,12 +24,23 @@ pub fn create(&self, user_id: &UserId, hash: &str) -> Result<()> {
     }
 
     /// Find out which user an access token belongs to.
-    pub fn find_from_token(&self, token: &str) -> Result<Option<UserId>> {
-        self.token_userid.get(token)?.map_or(Ok(None), |bytes| {
-            utils::string_from_bytes(&bytes)
-                .and_then(|string| Ok(UserId::try_from(string)?))
-                .map(Some)
-        })
+    pub fn find_from_token(&self, token: &str) -> Result<Option<(UserId, String)>> {
+        self.token_userdeviceid
+            .get(token)?
+            .map_or(Ok(None), |bytes| {
+                let mut parts = bytes.split(|&b| b == 0xff);
+                let user_bytes = parts
+                    .next()
+                    .ok_or(Error::BadDatabase("token_userdeviceid value invalid"))?;
+                let device_bytes = parts
+                    .next()
+                    .ok_or(Error::BadDatabase("token_userdeviceid value invalid"))?;
+
+                Ok(Some((
+                    UserId::try_from(utils::string_from_bytes(&user_bytes)?)?,
+                    utils::string_from_bytes(&device_bytes)?,
+                )))
+            })
     }
 
     /// Returns an iterator over all users on this homeserver.
@@ -105,27 +116,25 @@ pub fn create_device(&self, user_id: &UserId, device_id: &str, token: &str) -> R
 
     /// Replaces the access token of one device.
     pub fn set_token(&self, user_id: &UserId, device_id: &str, token: &str) -> Result<()> {
-        let mut key = user_id.to_string().as_bytes().to_vec();
-        key.push(0xff);
-        key.extend_from_slice(device_id.as_bytes());
+        let mut userdeviceid = user_id.to_string().as_bytes().to_vec();
+        userdeviceid.push(0xff);
+        userdeviceid.extend_from_slice(device_id.as_bytes());
 
-        if self.userdeviceids.get(&key)?.is_none() {
+        if self.userdeviceids.get(&userdeviceid)?.is_none() {
             return Err(Error::BadRequest(
                 "Tried to set token for nonexistent device",
             ));
         }
 
         // Remove old token
-        if let Some(old_token) = self.userdeviceid_token.get(&key)? {
-            self.token_userid.remove(old_token)?;
+        if let Some(old_token) = self.userdeviceid_token.get(&userdeviceid)? {
+            self.token_userdeviceid.remove(old_token)?;
             // It will be removed from userdeviceid_token by the insert later
         }
 
-        // Assign token to device_id
-        self.userdeviceid_token.insert(key, &*token)?;
-
-        // Assign token to user
-        self.token_userid.insert(token, &*user_id.to_string())?;
+        // Assign token to user device combination
+        self.userdeviceid_token.insert(&userdeviceid, &*token)?;
+        self.token_userdeviceid.insert(token, userdeviceid)?;
 
         Ok(())
     }
diff --git a/src/ruma_wrapper.rs b/src/ruma_wrapper.rs
index 7568573dc..7c904dbd3 100644
--- a/src/ruma_wrapper.rs
+++ b/src/ruma_wrapper.rs
@@ -18,6 +18,7 @@
 pub struct Ruma<T> {
     body: T,
     pub user_id: Option<UserId>,
+    pub device_id: Option<String>,
     pub json_body: serde_json::Value,
 }
 
@@ -40,7 +41,7 @@ fn from_data(
         Box::pin(async move {
             let data = rocket::try_outcome!(outcome.owned());
 
-            let user_id = if T::METADATA.requires_authentication {
+            let (user_id, device_id) = if T::METADATA.requires_authentication {
                 let db = request.guard::<State<'_, crate::Database>>().await.unwrap();
 
                 // Get token from header or query value
@@ -59,10 +60,11 @@ fn from_data(
                 match db.users.find_from_token(&token).unwrap() {
                     // TODO: M_UNKNOWN_TOKEN
                     None => return Failure((Status::Unauthorized, ())),
-                    Some(user_id) => Some(user_id),
+                    Some((user_id, device_id)) => (Some(user_id), Some(device_id)),
                 }
+
             } else {
-                None
+                (None, None)
             };
 
             let mut http_request = http::Request::builder()
@@ -83,6 +85,7 @@ fn from_data(
                 Ok(t) => Success(Ruma {
                     body: t,
                     user_id,
+                    device_id,
                     // TODO: Can we avoid parsing it again?
                     json_body: if !body.is_empty() {
                         serde_json::from_slice(&body).expect("Ruma already parsed it successfully")
-- 
GitLab