improvement: add option to get device id from token
This commit is contained in:
		
							parent
							
								
									ee0d6940bd
								
							
						
					
					
						commit
						1dbde0e1c1
					
				
					 4 changed files with 36 additions and 24 deletions
				
			
		|  | @ -105,7 +105,7 @@ pub fn register_route( | ||||||
|                 stages: vec!["m.login.dummy".to_owned()], |                 stages: vec!["m.login.dummy".to_owned()], | ||||||
|             }], |             }], | ||||||
|             completed: vec![], |             completed: vec![], | ||||||
|             params: RawValue::from_string("".to_owned()).unwrap(), |             params: RawValue::from_string("{}".to_owned()).unwrap(), | ||||||
|             session: Some(utils::random_string(SESSION_ID_LENGTH)), |             session: Some(utils::random_string(SESSION_ID_LENGTH)), | ||||||
|             auth_error: None, |             auth_error: None, | ||||||
|         }))); |         }))); | ||||||
|  |  | ||||||
|  | @ -44,11 +44,11 @@ impl Database { | ||||||
|                 userid_displayname: db.open_tree("userid_displayname").unwrap(), |                 userid_displayname: db.open_tree("userid_displayname").unwrap(), | ||||||
|                 userid_avatarurl: db.open_tree("userid_avatarurl").unwrap(), |                 userid_avatarurl: db.open_tree("userid_avatarurl").unwrap(), | ||||||
|                 userdeviceid_token: db.open_tree("userdeviceid_token").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 { |             rooms: rooms::Rooms { | ||||||
|                 edus: rooms::RoomEdus { |                 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
 |                     roomlatestid_roomlatest: db.open_tree("roomlatestid_roomlatest").unwrap(), // Read receipts
 | ||||||
|                     roomactiveid_roomactive: db.open_tree("roomactiveid_roomactive").unwrap(), // Typing notifs
 |                     roomactiveid_roomactive: db.open_tree("roomactiveid_roomactive").unwrap(), // Typing notifs
 | ||||||
|                 }, |                 }, | ||||||
|  |  | ||||||
|  | @ -8,7 +8,7 @@ pub struct Users { | ||||||
|     pub(super) userid_avatarurl: sled::Tree, |     pub(super) userid_avatarurl: sled::Tree, | ||||||
|     pub(super) userdeviceids: sled::Tree, |     pub(super) userdeviceids: sled::Tree, | ||||||
|     pub(super) userdeviceid_token: sled::Tree, |     pub(super) userdeviceid_token: sled::Tree, | ||||||
|     pub(super) token_userid: sled::Tree, |     pub(super) token_userdeviceid: sled::Tree, | ||||||
| } | } | ||||||
| 
 | 
 | ||||||
| impl Users { | impl Users { | ||||||
|  | @ -24,12 +24,23 @@ impl Users { | ||||||
|     } |     } | ||||||
| 
 | 
 | ||||||
|     /// Find out which user an access token belongs to.
 |     /// Find out which user an access token belongs to.
 | ||||||
|     pub fn find_from_token(&self, token: &str) -> Result<Option<UserId>> { |     pub fn find_from_token(&self, token: &str) -> Result<Option<(UserId, String)>> { | ||||||
|         self.token_userid.get(token)?.map_or(Ok(None), |bytes| { |         self.token_userdeviceid | ||||||
|             utils::string_from_bytes(&bytes) |             .get(token)? | ||||||
|                 .and_then(|string| Ok(UserId::try_from(string)?)) |             .map_or(Ok(None), |bytes| { | ||||||
|                 .map(Some) |                 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.
 |     /// Returns an iterator over all users on this homeserver.
 | ||||||
|  | @ -105,27 +116,25 @@ impl Users { | ||||||
| 
 | 
 | ||||||
|     /// Replaces the access token of one device.
 |     /// Replaces the access token of one device.
 | ||||||
|     pub fn set_token(&self, user_id: &UserId, device_id: &str, token: &str) -> Result<()> { |     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(); |         let mut userdeviceid = user_id.to_string().as_bytes().to_vec(); | ||||||
|         key.push(0xff); |         userdeviceid.push(0xff); | ||||||
|         key.extend_from_slice(device_id.as_bytes()); |         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( |             return Err(Error::BadRequest( | ||||||
|                 "Tried to set token for nonexistent device", |                 "Tried to set token for nonexistent device", | ||||||
|             )); |             )); | ||||||
|         } |         } | ||||||
| 
 | 
 | ||||||
|         // Remove old token
 |         // Remove old token
 | ||||||
|         if let Some(old_token) = self.userdeviceid_token.get(&key)? { |         if let Some(old_token) = self.userdeviceid_token.get(&userdeviceid)? { | ||||||
|             self.token_userid.remove(old_token)?; |             self.token_userdeviceid.remove(old_token)?; | ||||||
|             // It will be removed from userdeviceid_token by the insert later
 |             // It will be removed from userdeviceid_token by the insert later
 | ||||||
|         } |         } | ||||||
| 
 | 
 | ||||||
|         // Assign token to device_id
 |         // Assign token to user device combination
 | ||||||
|         self.userdeviceid_token.insert(key, &*token)?; |         self.userdeviceid_token.insert(&userdeviceid, &*token)?; | ||||||
| 
 |         self.token_userdeviceid.insert(token, userdeviceid)?; | ||||||
|         // Assign token to user
 |  | ||||||
|         self.token_userid.insert(token, &*user_id.to_string())?; |  | ||||||
| 
 | 
 | ||||||
|         Ok(()) |         Ok(()) | ||||||
|     } |     } | ||||||
|  |  | ||||||
|  | @ -18,6 +18,7 @@ const MESSAGE_LIMIT: u64 = 65535; | ||||||
| pub struct Ruma<T> { | pub struct Ruma<T> { | ||||||
|     body: T, |     body: T, | ||||||
|     pub user_id: Option<UserId>, |     pub user_id: Option<UserId>, | ||||||
|  |     pub device_id: Option<String>, | ||||||
|     pub json_body: serde_json::Value, |     pub json_body: serde_json::Value, | ||||||
| } | } | ||||||
| 
 | 
 | ||||||
|  | @ -40,7 +41,7 @@ impl<'a, T: Endpoint> FromData<'a> for Ruma<T> { | ||||||
|         Box::pin(async move { |         Box::pin(async move { | ||||||
|             let data = rocket::try_outcome!(outcome.owned()); |             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(); |                 let db = request.guard::<State<'_, crate::Database>>().await.unwrap(); | ||||||
| 
 | 
 | ||||||
|                 // Get token from header or query value
 |                 // Get token from header or query value
 | ||||||
|  | @ -59,10 +60,11 @@ impl<'a, T: Endpoint> FromData<'a> for Ruma<T> { | ||||||
|                 match db.users.find_from_token(&token).unwrap() { |                 match db.users.find_from_token(&token).unwrap() { | ||||||
|                     // TODO: M_UNKNOWN_TOKEN
 |                     // TODO: M_UNKNOWN_TOKEN
 | ||||||
|                     None => return Failure((Status::Unauthorized, ())), |                     None => return Failure((Status::Unauthorized, ())), | ||||||
|                     Some(user_id) => Some(user_id), |                     Some((user_id, device_id)) => (Some(user_id), Some(device_id)), | ||||||
|                 } |                 } | ||||||
|  | 
 | ||||||
|             } else { |             } else { | ||||||
|                 None |                 (None, None) | ||||||
|             }; |             }; | ||||||
| 
 | 
 | ||||||
|             let mut http_request = http::Request::builder() |             let mut http_request = http::Request::builder() | ||||||
|  | @ -83,6 +85,7 @@ impl<'a, T: Endpoint> FromData<'a> for Ruma<T> { | ||||||
|                 Ok(t) => Success(Ruma { |                 Ok(t) => Success(Ruma { | ||||||
|                     body: t, |                     body: t, | ||||||
|                     user_id, |                     user_id, | ||||||
|  |                     device_id, | ||||||
|                     // TODO: Can we avoid parsing it again?
 |                     // TODO: Can we avoid parsing it again?
 | ||||||
|                     json_body: if !body.is_empty() { |                     json_body: if !body.is_empty() { | ||||||
|                         serde_json::from_slice(&body).expect("Ruma already parsed it successfully") |                         serde_json::from_slice(&body).expect("Ruma already parsed it successfully") | ||||||
|  |  | ||||||
		Loading…
	
		Reference in a new issue