feat: Implement oauth auth code and refresh token flows

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| {