feat: Implement oauth auth code and refresh token flows

src/service/Cargo.toml

@@ -119,6 +119,7 @@ recaptcha-verify = { version = "0.2.0", default-features = false }
 reqwest_recaptcha = { package = "reqwest", version = "0.12.28", default-features = false, features = ["rustls-tls-native-roots-no-provider"]  } # As long as recaptcha-verify's reqwest is outdated
 yansi.workspace = true
 lettre.workspace = true
+serde_urlencoded.workspace = true
 
 [target.'cfg(all(unix, target_os = "linux"))'.dependencies]
 sd-notify.workspace = true

src/service/appservice/mod.rs

@@ -67,7 +67,7 @@ impl crate::Service for Service {
 		for (id, registration) in appservices {
 			// During startup, resolve any token collisions in favour of appservices
 			// by logging out conflicting user devices
-			if let Ok((user_id, device_id)) = self
+			if let Some((user_id, device_id)) = self
 				.services
 				.users
 				.find_from_token(&registration.as_token)
@@ -158,7 +158,7 @@ impl Service {
 			.users
 			.find_from_token(&registration.as_token)
 			.await
-			.is_ok()
+			.is_some()
 		{
 			return Err(err!(Request(InvalidParam(
 				"Cannot register appservice: The provided token is already in use by a user \

src/service/oauth/client_metadata.rs

@@ -38,7 +38,7 @@ pub struct ClientMetadata {
 }
 
 impl ClientMetadata {
-	const ACCEPTABLE_LOCALHOSTS: [&str; 3] = ["localhost", "127.0.0.1", "[::1]"];
+	pub(super) const ACCEPTABLE_LOCALHOSTS: [&str; 3] = ["localhost", "127.0.0.1", "[::1]"];
 
 	pub(super) fn validate(&self) -> Result<(), &'static str> {
 		let Some(client_domain) = self.client_uri.domain() else {
@@ -137,6 +137,7 @@ pub enum GrantType {
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
+#[non_exhaustive]
 pub enum ResponseType {
 	Code,
 }

src/service/oauth/grant.rs

@@ -0,0 +1,150 @@
+use std::{
+	collections::{BTreeSet, HashSet},
+	fmt::Debug,
+	hash::Hash,
+	mem::discriminant,
+};
+
+use regex::Regex;
+use ruma::OwnedDeviceId;
+use serde::{Deserialize, Serialize};
+use url::Url;
+
+use super::client_metadata::ResponseType;
+
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct AuthorizationCodeQuery {
+	pub response_type: ResponseType,
+	pub client_id: String,
+	pub redirect_uri: Url,
+	pub scope: RawScopes,
+	pub state: String,
+	#[serde(default)]
+	pub response_mode: ResponseMode,
+	pub code_challenge: String,
+	pub code_challenge_method: CodeChallengeMethod,
+}
+
+#[derive(Debug, Clone, Default, Deserialize, Serialize)]
+#[serde(rename_all = "snake_case")]
+#[non_exhaustive]
+pub enum ResponseMode {
+	#[default]
+	// default for `code` response type, see https://openid.net/specs/oauth-v2-multiple-response-types-1_0.html#:~:text=Client%2E-,For,encoding%2E,-See
+	Query,
+	Fragment,
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize)]
+#[non_exhaustive]
+pub enum CodeChallengeMethod {
+	S256,
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize, PartialOrd, Ord)]
+pub enum Scope {
+	Device(OwnedDeviceId),
+	ClientApi,
+}
+
+impl PartialEq for Scope {
+	fn eq(&self, other: &Self) -> bool { discriminant(self) == discriminant(other) }
+}
+
+impl Eq for Scope {}
+
+impl Hash for Scope {
+	fn hash<H: std::hash::Hasher>(&self, state: &mut H) { discriminant(self).hash(state); }
+}
+
+impl std::fmt::Display for Scope {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		let urn = match self {
+			| Self::ClientApi => "urn:matrix:client:api:*".to_owned(),
+			| Self::Device(device_id) => format!("urn:matrix:client:device:{device_id}"),
+		};
+
+		f.write_str(&urn)
+	}
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct RawScopes(String);
+
+impl RawScopes {
+	pub fn to_scopes(&self) -> Result<BTreeSet<Scope>, String> {
+		let client_api_token_regex =
+			Regex::new(r"urn:matrix:(client|org.matrix.msc2967.client):api:\*").unwrap();
+		let device_token_regex = Regex::new(
+			r"urn:matrix:(client|org.matrix.msc2967.client):device:([a-zA-Z0-9-._~]{5,})",
+		)
+		.unwrap();
+
+		let mut scopes = BTreeSet::new();
+
+		for token in self.0.split(' ') {
+			let scope_was_new = {
+				if client_api_token_regex.is_match(token) {
+					scopes.insert(Scope::ClientApi)
+				} else if let Some(captures) = device_token_regex.captures(token) {
+					scopes.insert(Scope::Device(captures.get(2).unwrap().as_str().into()))
+				} else if token == "openid" {
+					// TODO(unspecced): Element sets this scope but doesn't use it for anything
+					true
+				} else {
+					return Err(format!("Invalid scope: {token}"));
+				}
+			};
+
+			if !scope_was_new {
+				return Err("Scope was specified more than once".to_owned());
+			}
+		}
+
+		Ok(scopes)
+	}
+}
+
+#[derive(Serialize)]
+pub struct AuthorizationCodeResponse {
+	pub state: String,
+	pub code: String,
+}
+
+#[derive(Deserialize)]
+#[serde(tag = "grant_type", rename_all = "snake_case")]
+pub enum TokenRequest {
+	AuthorizationCode {
+		code: String,
+		redirect_uri: Url,
+		client_id: String,
+		code_verifier: String,
+	},
+	RefreshToken {
+		client_id: String,
+		refresh_token: String,
+	},
+}
+
+impl TokenRequest {
+	pub fn client_id(&self) -> &str {
+		match self {
+			| Self::AuthorizationCode { client_id, .. }
+			| Self::RefreshToken { client_id, .. } => client_id,
+		}
+	}
+}
+
+#[derive(Serialize)]
+pub struct TokenResponse {
+	pub access_token: String,
+	pub token_type: TokenType,
+	pub expires_in: u64,
+	pub refresh_token: String,
+	pub scope: String,
+}
+
+#[derive(Serialize)]
+pub enum TokenType {
+	Bearer,
+}

src/service/oauth/mod.rs

@@ -1,30 +1,93 @@
 use std::{
-	collections::HashMap,
+	collections::{BTreeMap, BTreeSet, HashMap},
 	sync::{Arc, Mutex},
 	time::{Duration, SystemTime},
 };
 
 use base64::Engine;
-use conduwuit::{Result, utils::hash::sha256};
+use conduwuit::{
+	Err, Result, err, info,
+	utils::{self, hash::sha256},
+};
 use database::{Deserialized, Json, Map};
-use ruma::DeviceId;
+use itertools::Itertools;
+use ruma::{DeviceId, OwnedDeviceId, OwnedUserId, UserId};
+use serde::{Deserialize, Serialize};
+use url::Url;
 
-use crate::{Dep, config, oauth::client_metadata::ClientMetadata};
+use crate::{
+	Dep, config,
+	oauth::{
+		client_metadata::{ApplicationType, ClientMetadata, ResponseType},
+		grant::{
+			AuthorizationCodeQuery, AuthorizationCodeResponse, CodeChallengeMethod, ResponseMode,
+			Scope, TokenRequest, TokenResponse, TokenType,
+		},
+	},
+	users,
+};
 
 pub mod client_metadata;
+pub mod grant;
 
 pub struct Service {
 	services: Services,
 	db: Data,
 	tickets: Mutex<HashMap<String, HashMap<OAuthTicket, SystemTime>>>,
+	pending_code_grants: tokio::sync::Mutex<HashMap<String, PendingCodeGrant>>,
 }
 
 struct Data {
 	clientid_clientmetadata: Arc<Map>,
+	userdeviceid_oauthsessioninfo: Arc<Map>,
+	refreshtoken_refreshtokeninfo: Arc<Map>,
 }
 
 struct Services {
 	config: Dep<config::Service>,
+	users: Dep<users::Service>,
+}
+
+#[derive(Deserialize, Serialize)]
+struct SessionInfo {
+	client_id: String,
+	current_refresh_token: String,
+	scopes: BTreeSet<Scope>,
+}
+
+#[derive(Deserialize, Serialize)]
+struct RefreshTokenInfo {
+	client_id: String,
+	user_id: OwnedUserId,
+	device_id: OwnedDeviceId,
+}
+
+struct PendingCodeGrant {
+	authorizing_user: OwnedUserId,
+	requested_scopes: BTreeSet<Scope>,
+	client_name: Option<String>,
+	expected_client_id: String,
+	expected_redirect_uri: Url,
+	code_challenge: String,
+	requested_at: SystemTime,
+}
+
+impl PendingCodeGrant {
+	const MAX_AGE: Duration = Duration::from_mins(1);
+	const RANDOM_CODE_LENGTH: usize = 32;
+
+	#[must_use]
+	pub fn generate_code() -> String { utils::random_string(Self::RANDOM_CODE_LENGTH) }
+
+	#[must_use]
+	pub fn is_valid_for(&self, client_id: &str) -> bool {
+		let now = SystemTime::now();
+
+		self.expected_client_id == client_id
+			&& now
+				.duration_since(self.requested_at)
+				.is_ok_and(|age| age < Self::MAX_AGE)
+	}
 }
 
 /// A time-limited grant for a client to perform some sensitive action.
@@ -49,11 +112,15 @@ impl crate::Service for Service {
 		Ok(Arc::new(Self {
 			services: Services {
 				config: args.depend::<config::Service>("config"),
+				users: args.depend::<users::Service>("users"),
 			},
 			db: Data {
 				clientid_clientmetadata: args.db["clientid_clientmetadata"].clone(),
+				userdeviceid_oauthsessioninfo: args.db["userdeviceid_oauthsessioninfo"].clone(),
+				refreshtoken_refreshtokeninfo: args.db["refreshtoken_refreshtokeninfo"].clone(),
 			},
 			tickets: Mutex::default(),
+			pending_code_grants: tokio::sync::Mutex::default(),
 		}))
 	}
 
@@ -61,6 +128,11 @@ impl crate::Service for Service {
 }
 
 impl Service {
+	const ACCESS_TOKEN_MAX_AGE: Duration = Duration::from_hours(1);
+	const RANDOM_TOKEN_LENGTH: usize = 32;
+
+	fn generate_token() -> String { utils::random_string(Self::RANDOM_TOKEN_LENGTH) }
+
 	pub async fn register_client(
 		&self,
 		metadata: &ClientMetadata,
@@ -85,7 +157,7 @@ impl Service {
 		Ok(client_id)
 	}
 
-	pub async fn get_client_registration(&self, client_id: &str) -> Option<ClientMetadata> {
+	pub async fn get_client_metadata(&self, client_id: &str) -> Option<ClientMetadata> {
 		self.db
 			.clientid_clientmetadata
 			.get(client_id)
@@ -94,15 +166,304 @@ impl Service {
 			.ok()
 	}
 
-	pub async fn get_client_id_for_device(&self, _device_id: &DeviceId) -> Option<String> {
-		None // TODO
+	pub async fn get_client_id_for_device(
+		&self,
+		user_id: &UserId,
+		device_id: &DeviceId,
+	) -> Option<String> {
+		self.db
+			.userdeviceid_oauthsessioninfo
+			.qry(&(user_id, device_id))
+			.await
+			.deserialized::<SessionInfo>()
+			.ok()
+			.map(|session| session.client_id)
+	}
+
+	pub async fn request_authorization_code(
+		&self,
+		authorizing_user: OwnedUserId,
+		query: AuthorizationCodeQuery,
+	) -> Result<String, String> {
+		let Some(client_metadata) = self.get_client_metadata(&query.client_id).await else {
+			return Err("Invalid client ID".to_owned());
+		};
+
+		if !(client_metadata
+			.response_types
+			.contains(&query.response_type)
+			&& matches!(query.response_type, ResponseType::Code))
+		{
+			return Err("Invalid response type".to_owned());
+		}
+
+		if !matches!(query.code_challenge_method, CodeChallengeMethod::S256) {
+			return Err("Invalid code challenge type".to_owned());
+		}
+
+		{
+			let mut stripped_uri = query.redirect_uri.clone();
+
+			if client_metadata.application_type == ApplicationType::Native
+				&& query
+					.redirect_uri
+					.host_str()
+					.is_some_and(|host| ClientMetadata::ACCEPTABLE_LOCALHOSTS.contains(&host))
+			{
+				// Remove the port from localhost redirect URIs for native applications when
+				// checking if it's valid
+				stripped_uri.set_port(None).unwrap();
+			}
+
+			if !client_metadata.redirect_uris.contains(&stripped_uri) {
+				return Err("Invalid redirect URI".to_owned());
+			}
+		}
+
+		let requested_scopes = query.scope.to_scopes()?;
+
+		let redirect_uri_query_separator = match query.response_mode {
+			| ResponseMode::Fragment => '#',
+			| ResponseMode::Query => '?',
+		};
+
+		let code = PendingCodeGrant::generate_code();
+
+		info!(
+			client_id = &query.client_id,
+			client_name = &client_metadata.client_name,
+			?requested_scopes,
+			?authorizing_user,
+			"Issuing oauth authorization code"
+		);
+
+		let redirect_uri = format!(
+			"{}{}{}",
+			query.redirect_uri,
+			redirect_uri_query_separator,
+			serde_urlencoded::to_string(AuthorizationCodeResponse {
+				state: query.state,
+				code: code.clone(),
+			})
+			.unwrap(),
+		);
+
+		let pending_grant = PendingCodeGrant {
+			authorizing_user,
+			requested_scopes,
+			client_name: client_metadata.client_name,
+			expected_client_id: query.client_id,
+			expected_redirect_uri: query.redirect_uri,
+			code_challenge: query.code_challenge,
+			requested_at: SystemTime::now(),
+		};
+
+		self.pending_code_grants
+			.lock()
+			.await
+			.insert(code, pending_grant);
+
+		Ok(redirect_uri)
+	}
+
+	pub async fn issue_token(&self, request: TokenRequest) -> Result<TokenResponse> {
+		match request {
+			| TokenRequest::AuthorizationCode {
+				code,
+				redirect_uri,
+				client_id,
+				code_verifier,
+			} => {
+				let mut pending_grants = self.pending_code_grants.lock().await;
+
+				let Some(pending_grant) = pending_grants
+					.remove(&code)
+					.filter(|grant| grant.is_valid_for(&client_id))
+				else {
+					return Err!("Invalid code");
+				};
+
+				if redirect_uri != pending_grant.expected_redirect_uri {
+					return Err!("Unexpected redirect uri");
+				}
+
+				let expected_code_challenge =
+					base64::prelude::BASE64_URL_SAFE_NO_PAD.encode(sha256::hash(&code_verifier));
+				if expected_code_challenge != pending_grant.code_challenge {
+					return Err!("Invalid code challenge");
+				}
+
+				self.create_session(
+					pending_grant.authorizing_user,
+					pending_grant.requested_scopes,
+					pending_grant.client_name,
+					client_id,
+				)
+				.await
+			},
+			| TokenRequest::RefreshToken { client_id, refresh_token } =>
+				self.refresh_session(client_id, refresh_token).await,
+		}
+	}
+
+	async fn create_session(
+		&self,
+		authorizing_user: OwnedUserId,
+		requested_scopes: BTreeSet<Scope>,
+		client_name: Option<String>,
+		client_id: String,
+	) -> Result<TokenResponse> {
+		let access_token = Self::generate_token();
+		let refresh_token = Self::generate_token();
+
+		let device_id = requested_scopes
+			.iter()
+			.find_map(|scope| {
+				if let Scope::Device(device_id) = scope {
+					Some(device_id)
+				} else {
+					None
+				}
+			})
+			.ok_or_else(|| err!("No device ID scope supplied"))?;
+
+		self.services
+			.users
+			.create_device(
+				&authorizing_user,
+				device_id,
+				&access_token,
+				Some(Self::ACCESS_TOKEN_MAX_AGE),
+				client_name,
+				None,
+			)
+			.await?;
+
+		self.db.userdeviceid_oauthsessioninfo.put(
+			(&authorizing_user, device_id),
+			Json(SessionInfo {
+				client_id: client_id.clone(),
+				current_refresh_token: refresh_token.clone(),
+				scopes: requested_scopes.clone(),
+			}),
+		);
+
+		self.db.refreshtoken_refreshtokeninfo.raw_put(
+			&refresh_token,
+			Json(RefreshTokenInfo {
+				client_id: client_id.clone(),
+				user_id: authorizing_user.clone(),
+				device_id: device_id.to_owned(),
+			}),
+		);
+
+		info!(
+			?client_id,
+			?authorizing_user,
+			?device_id,
+			?requested_scopes,
+			"Created new oauth session"
+		);
+
+		Ok(TokenResponse {
+			access_token,
+			token_type: TokenType::Bearer,
+			expires_in: Self::ACCESS_TOKEN_MAX_AGE.as_secs(),
+			scope: requested_scopes.iter().join(" "),
+			refresh_token,
+		})
+	}
+
+	async fn refresh_session(
+		&self,
+		client_id: String,
+		refresh_token: String,
+	) -> Result<TokenResponse> {
+		let Some(refresh_token_info) = self
+			.db
+			.refreshtoken_refreshtokeninfo
+			.get(&refresh_token)
+			.await
+			.deserialized::<RefreshTokenInfo>()
+			.ok()
+		else {
+			return Err!("Invalid refresh token");
+		};
+
+		assert_eq!(&client_id, &refresh_token_info.client_id, "refresh token client id mismatch");
+
+		let mut session_info = self
+			.db
+			.userdeviceid_oauthsessioninfo
+			.qry(&(&refresh_token_info.user_id, &refresh_token_info.device_id))
+			.await
+			.deserialized::<SessionInfo>()
+			.expect("session info should exist");
+
+		assert_eq!(&client_id, &session_info.client_id, "session info client id mismatch");
+
+		let new_access_token = Self::generate_token();
+		let new_refresh_token = Self::generate_token();
+		let scope = session_info.scopes.iter().join(" ");
+		session_info
+			.current_refresh_token
+			.clone_from(&new_refresh_token);
+
+		self.services
+			.users
+			.set_token(
+				&refresh_token_info.user_id,
+				&refresh_token_info.device_id,
+				&new_access_token,
+				Some(Self::ACCESS_TOKEN_MAX_AGE),
+			)
+			.await?;
+
+		self.db.userdeviceid_oauthsessioninfo.put(
+			(&refresh_token_info.user_id, &refresh_token_info.device_id),
+			Json(session_info),
+		);
+
+		self.db.refreshtoken_refreshtokeninfo.remove(&refresh_token);
+		drop(refresh_token);
+		self.db
+			.refreshtoken_refreshtokeninfo
+			.raw_put(&new_refresh_token, Json(refresh_token_info));
+
+		Ok(TokenResponse {
+			access_token: new_access_token,
+			token_type: TokenType::Bearer,
+			expires_in: Self::ACCESS_TOKEN_MAX_AGE.as_secs(),
+			scope,
+			refresh_token: new_refresh_token,
+		})
+	}
+
+	pub async fn remove_session(&self, user_id: &UserId, device_id: &DeviceId) {
+		let session_info = self
+			.db
+			.userdeviceid_oauthsessioninfo
+			.qry(&(user_id, device_id))
+			.await
+			.deserialized::<SessionInfo>()
+			.ok();
+
+		if let Some(session_info) = session_info {
+			self.db
+				.refreshtoken_refreshtokeninfo
+				.remove(&session_info.current_refresh_token);
+			self.db
+				.userdeviceid_oauthsessioninfo
+				.del(&(user_id, device_id));
+			info!(?user_id, ?device_id, "Removed OAuth session");
+		}
 	}
 
 	/// Issue a ticket for `localpart` to perform some action.
 	pub fn issue_ticket(&self, localpart: String, ticket: OAuthTicket) {
 		self.tickets
 			.lock()
-			.expect("should be able to lock tickets")
+			.unwrap()
 			.entry(localpart)
 			.or_default()
 			.insert(ticket, SystemTime::now());
@@ -114,7 +475,7 @@ impl Service {
 
 		self.tickets
 			.lock()
-			.expect("should be able to lock tickets")
+			.unwrap()
 			.get_mut(localpart)
 			.and_then(|tickets| tickets.remove(&ticket))
 			.is_some_and(|issued| {

src/service/uiaa/mod.rs

@@ -253,17 +253,17 @@ impl Service {
 		let mut info = assign!(UiaaInfo::new(flows), { params: Some(params), session: Some(session_id.clone()) });
 
 		let session_metadata = if let Some(initiator) = initiator {
-			let is_oauth = OptionFuture::from(
-				initiator.device_id.map(async |device_id| {
-					self
-						.services
-						.oauth
-						.get_client_id_for_device(device_id)
-						.await
-				})
-			)
-			.await
-			.is_some();
+			let is_oauth = if let Some(device_id) = initiator.device_id {
+				self
+					.services
+					.oauth
+					.get_client_id_for_device(initiator.user_id, device_id)
+					.await
+					.is_some()
+			} else {
+				// Appservices never have oauth sessions
+				false
+			};
 
 			if is_oauth {
 				if let Some(oauth_ticket) = initiator.oauth_ticket {
@@ -279,7 +279,18 @@ impl Service {
 						.unwrap();
 
 					info.flows = vec![AuthFlow::new(vec![AuthType::OAuth])];
-					info.params = Some(to_raw_value(&json!({"url": ticket_url})).unwrap());
+					info.params = Some(
+						to_raw_value(&json!({
+							AuthType::OAuth.as_str(): {
+								"url": ticket_url,
+							},
+							// TODO(compat): This is necessary for older versions of matrix-rust-sdk
+							"org.matrix.cross_signing_reset": {
+								"url": ticket_url,
+							}
+						}))
+						.unwrap(),
+					);
 
 					UiaaSessionMetadata::OAuth {
 						localpart: initiator.user_id.localpart().to_owned(),

src/service/users/dehydrated_device.rs

@@ -54,6 +54,7 @@ pub async fn set_dehydrated_device(&self, user_id: &UserId, request: Request) ->
 		user_id,
 		&request.device_id,
 		"",
+		None,
 		request.initial_device_display_name.clone(),
 		None,
 	)

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(())
 	}