feat: Implement oauth auth code and refresh token flows

src/service/users/mod.rs

@@ -1,6 +1,12 @@
 pub(super) mod dehydrated_device;
 
-use std::{collections::BTreeMap, mem, net::IpAddr, sync::Arc};
+use std::{
+	collections::BTreeMap,
+	mem,
+	net::IpAddr,
+	sync::Arc,
+	time::{Duration, SystemTime},
+};
 
 use conduwuit::{
 	Err, Error, Result, Server, debug_error, debug_warn, err, trace,
@@ -26,7 +32,7 @@ use ruminuwuity::invite_permission_config::{FilterLevel, InvitePermissionConfigE
 use serde::{Deserialize, Serialize};
 use serde_json::json;
 
-use crate::{Dep, account_data, admin, appservice, globals, rooms};
+use crate::{Dep, account_data, admin, appservice, globals, oauth, rooms};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct UserSuspension {
@@ -62,6 +68,7 @@ struct Services {
 	admin: Dep<admin::Service>,
 	appservice: Dep<appservice::Service>,
 	globals: Dep<globals::Service>,
+	oauth: Dep<oauth::Service>,
 	state_accessor: Dep<rooms::state_accessor::Service>,
 	state_cache: Dep<rooms::state_cache::Service>,
 }
@@ -75,6 +82,7 @@ struct Data {
 	logintoken_expiresatuserid: Arc<Map>,
 	todeviceid_events: Arc<Map>,
 	token_userdeviceid: Arc<Map>,
+	userdeviceid_tokenexpires: Arc<Map>,
 	userdeviceid_metadata: Arc<Map>,
 	userdeviceid_token: Arc<Map>,
 	userfilterid_filter: Arc<Map>,
@@ -102,6 +110,7 @@ impl crate::Service for Service {
 				admin: args.depend::<admin::Service>("admin"),
 				appservice: args.depend::<appservice::Service>("appservice"),
 				globals: args.depend::<globals::Service>("globals"),
+				oauth: args.depend::<oauth::Service>("oauth"),
 				state_accessor: args
 					.depend::<rooms::state_accessor::Service>("rooms::state_accessor"),
 				state_cache: args.depend::<rooms::state_cache::Service>("rooms::state_cache"),
@@ -131,6 +140,7 @@ impl crate::Service for Service {
 				userid_selfsigningkeyid: args.db["userid_selfsigningkeyid"].clone(),
 				userid_usersigningkeyid: args.db["userid_usersigningkeyid"].clone(),
 				useridprofilekey_value: args.db["useridprofilekey_value"].clone(),
+				userdeviceid_tokenexpires: args.db["userdeviceid_tokenexpires"].clone(),
 			},
 		}))
 	}
@@ -337,8 +347,37 @@ impl Service {
 	pub async fn count(&self) -> usize { self.db.userid_password.count().await }
 
 	/// Find out which user an access token belongs to.
-	pub async fn find_from_token(&self, token: &str) -> Result<(OwnedUserId, OwnedDeviceId)> {
-		self.db.token_userdeviceid.get(token).await.deserialized()
+	pub async fn find_from_token(&self, token: &str) -> Option<(OwnedUserId, OwnedDeviceId)> {
+		let user = self
+			.db
+			.token_userdeviceid
+			.get(token)
+			.await
+			.deserialized()
+			.ok();
+
+		// Check if the token has expired
+		if let Some(user) = &user {
+			if let Some(expires) = self
+				.db
+				.userdeviceid_tokenexpires
+				.qry(user)
+				.await
+				.deserialized::<u64>()
+				.ok()
+				.map(Duration::from_secs)
+			{
+				let expires_at = SystemTime::UNIX_EPOCH
+					.checked_add(expires)
+					.expect("expiry time should not overflow SystemTime");
+
+				if SystemTime::now() > expires_at {
+					return None;
+				}
+			}
+		}
+
+		user
 	}
 
 	/// Returns an iterator over all users on this homeserver.
@@ -434,6 +473,7 @@ impl Service {
 		user_id: &UserId,
 		device_id: &DeviceId,
 		token: &str,
+		token_max_age: Option<Duration>,
 		initial_device_display_name: Option<String>,
 		client_ip: Option<String>,
 	) -> Result<()> {
@@ -451,7 +491,8 @@ impl Service {
 
 		increment(&self.db.userid_devicelistversion, user_id.as_bytes());
 		self.db.userdeviceid_metadata.put(key, Json(device));
-		self.set_token(user_id, device_id, token).await
+		self.set_token(user_id, device_id, token, token_max_age)
+			.await
 	}
 
 	/// Removes a device from a user.
@@ -467,6 +508,7 @@ impl Service {
 		if let Ok(old_token) = self.db.userdeviceid_token.qry(&userdeviceid).await {
 			self.db.userdeviceid_token.del(userdeviceid);
 			self.db.token_userdeviceid.remove(&old_token);
+			self.db.userdeviceid_tokenexpires.del(userdeviceid);
 		}
 
 		// Remove todevice events
@@ -480,6 +522,9 @@ impl Service {
 
 		// TODO: Remove onetimekeys
 
+		// Remove OAuth session information
+		self.services.oauth.remove_session(user_id, device_id).await;
+
 		increment(&self.db.userid_devicelistversion, user_id.as_bytes());
 
 		self.db.userdeviceid_metadata.del(userdeviceid);
@@ -535,6 +580,7 @@ impl Service {
 		user_id: &UserId,
 		device_id: &DeviceId,
 		token: &str,
+		token_max_age: Option<Duration>,
 	) -> Result<()> {
 		let key = (user_id, device_id);
 		if self.db.userdeviceid_metadata.qry(&key).await.is_err() {
@@ -561,6 +607,7 @@ impl Service {
 		// Remove old token
 		if let Ok(old_token) = self.db.userdeviceid_token.qry(&key).await {
 			self.db.token_userdeviceid.remove(&old_token);
+			self.db.userdeviceid_tokenexpires.remove(&old_token);
 			// It will be removed from userdeviceid_token by the insert later
 		}
 
@@ -568,6 +615,18 @@ impl Service {
 		self.db.userdeviceid_token.put_raw(key, token);
 		self.db.token_userdeviceid.raw_put(token, key);
 
+		if let Some(max_age) = token_max_age {
+			let expires = SystemTime::now()
+				.duration_since(SystemTime::UNIX_EPOCH)
+				.expect("system time should not be before the epoch")
+				.saturating_add(max_age)
+				.as_secs();
+
+			self.db.userdeviceid_tokenexpires.put(key, expires);
+		} else {
+			self.db.userdeviceid_tokenexpires.del(key);
+		}
+
 		Ok(())
 	}