Split the service in multiple crates

crates/core/src/handlers/oauth2/authorization.rs

@@ -0,0 +1,459 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::{
+    collections::{HashMap, HashSet},
+    convert::TryFrom,
+};
+
+use chrono::Duration;
+use hyper::{
+    header::LOCATION,
+    http::uri::{Parts, PathAndQuery, Uri},
+    StatusCode,
+};
+use itertools::Itertools;
+use oauth2_types::{
+    errors::{ErrorResponse, InvalidRequest, OAuth2Error},
+    pkce,
+    requests::{
+        AccessTokenResponse, AuthorizationRequest, AuthorizationResponse, ResponseMode,
+        ResponseType,
+    },
+};
+use rand::{distributions::Alphanumeric, thread_rng, Rng};
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use sqlx::{PgPool, Postgres, Transaction};
+use url::Url;
+use warp::{
+    redirect::see_other,
+    reject::InvalidQuery,
+    reply::{html, with_header},
+    Filter, Rejection, Reply,
+};
+
+use crate::{
+    config::{CookiesConfig, OAuth2ClientConfig, OAuth2Config},
+    errors::WrapError,
+    filters::{
+        database::with_transaction,
+        session::{with_optional_session, with_session},
+        with_templates,
+    },
+    handlers::views::LoginRequest,
+    storage::{
+        oauth2::{
+            access_token::add_access_token,
+            refresh_token::add_refresh_token,
+            session::{get_session_by_id, start_session},
+        },
+        SessionInfo,
+    },
+    templates::{FormPostContext, Templates},
+    tokens,
+};
+
+#[derive(Deserialize)]
+struct PartialParams {
+    client_id: Option<String>,
+    redirect_uri: Option<String>,
+    /*
+    response_type: Option<String>,
+    response_mode: Option<String>,
+    */
+}
+
+enum ReplyOrBackToClient {
+    Reply(Box<dyn Reply>),
+    BackToClient {
+        params: Value,
+        redirect_uri: Url,
+        response_mode: ResponseMode,
+    },
+    Error(Box<dyn OAuth2Error>),
+}
+
+fn back_to_client<T>(
+    mut redirect_uri: Url,
+    response_mode: ResponseMode,
+    params: T,
+    templates: &Templates,
+) -> anyhow::Result<Box<dyn Reply>>
+where
+    T: Serialize,
+{
+    #[derive(Serialize)]
+    struct AllParams<'s, T> {
+        #[serde(flatten, skip_serializing_if = "Option::is_none")]
+        existing: Option<HashMap<&'s str, &'s str>>,
+
+        #[serde(flatten)]
+        params: T,
+    }
+
+    match response_mode {
+        ResponseMode::Query => {
+            let existing: Option<HashMap<&str, &str>> = redirect_uri
+                .query()
+                .map(|qs| serde_urlencoded::from_str(qs))
+                .transpose()?;
+
+            let merged = AllParams { existing, params };
+
+            let new_qs = serde_urlencoded::to_string(merged)?;
+
+            redirect_uri.set_query(Some(&new_qs));
+
+            Ok(Box::new(with_header(
+                StatusCode::SEE_OTHER,
+                LOCATION,
+                redirect_uri.as_str(),
+            )))
+        }
+        ResponseMode::Fragment => {
+            let existing: Option<HashMap<&str, &str>> = redirect_uri
+                .fragment()
+                .map(|qs| serde_urlencoded::from_str(qs))
+                .transpose()?;
+
+            let merged = AllParams { existing, params };
+
+            let new_qs = serde_urlencoded::to_string(merged)?;
+
+            redirect_uri.set_fragment(Some(&new_qs));
+
+            Ok(Box::new(with_header(
+                StatusCode::SEE_OTHER,
+                LOCATION,
+                redirect_uri.as_str(),
+            )))
+        }
+        ResponseMode::FormPost => {
+            let ctx = FormPostContext::new(redirect_uri, params);
+            let rendered = templates.render_form_post(&ctx)?;
+            Ok(Box::new(html(rendered)))
+        }
+    }
+}
+
+#[derive(Deserialize)]
+struct Params {
+    #[serde(flatten)]
+    auth: AuthorizationRequest,
+
+    #[serde(flatten)]
+    pkce: Option<pkce::Request>,
+}
+
+/// Given a list of response types and an optional user-defined response mode,
+/// figure out what response mode must be used, and emit an error if the
+/// suggested response mode isn't allowed for the given response types.
+fn resolve_response_mode(
+    response_type: &HashSet<ResponseType>,
+    suggested_response_mode: Option<ResponseMode>,
+) -> anyhow::Result<ResponseMode> {
+    use ResponseMode as M;
+    use ResponseType as T;
+
+    // If the response type includes either "token" or "id_token", the default
+    // response mode is "fragment" and the response mode "query" must not be
+    // used
+    if response_type.contains(&T::Token) || response_type.contains(&T::IdToken) {
+        match suggested_response_mode {
+            None => Ok(M::Fragment),
+            Some(M::Query) => Err(anyhow::anyhow!("invalid response mode")),
+            Some(mode) => Ok(mode),
+        }
+    } else {
+        // In other cases, all response modes are allowed, defaulting to "query"
+        Ok(suggested_response_mode.unwrap_or(M::Query))
+    }
+}
+
+pub fn filter(
+    pool: &PgPool,
+    templates: &Templates,
+    oauth2_config: &OAuth2Config,
+    cookies_config: &CookiesConfig,
+) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone + Send + Sync + 'static {
+    let clients = oauth2_config.clients.clone();
+    let authorize = warp::path!("oauth2" / "authorize")
+        .and(warp::get())
+        .map(move || clients.clone())
+        .and(warp::query())
+        .and(with_optional_session(pool, cookies_config))
+        .and(with_transaction(pool))
+        .and_then(get);
+
+    let step = warp::path!("oauth2" / "authorize" / "step")
+        .and(warp::get())
+        .and(warp::query().map(|s: StepRequest| s.id))
+        .and(with_session(pool, cookies_config))
+        .and(with_transaction(pool))
+        .and_then(step);
+
+    let clients = oauth2_config.clients.clone();
+    authorize
+        .or(step)
+        .unify()
+        .recover(recover)
+        .unify()
+        .and(warp::query())
+        .and(warp::any().map(move || clients.clone()))
+        .and(with_templates(templates))
+        .and_then(actually_reply)
+}
+
+async fn recover(rejection: Rejection) -> Result<ReplyOrBackToClient, Rejection> {
+    if rejection.find::<InvalidQuery>().is_some() {
+        Ok(ReplyOrBackToClient::Error(Box::new(InvalidRequest)))
+    } else {
+        Err(rejection)
+    }
+}
+
+async fn actually_reply(
+    rep: ReplyOrBackToClient,
+    q: PartialParams,
+    clients: Vec<OAuth2ClientConfig>,
+    templates: Templates,
+) -> Result<impl Reply, Rejection> {
+    let (redirect_uri, response_mode, params) = match rep {
+        ReplyOrBackToClient::Reply(r) => return Ok(r),
+        ReplyOrBackToClient::BackToClient {
+            redirect_uri,
+            response_mode,
+            params,
+        } => (redirect_uri, response_mode, params),
+        ReplyOrBackToClient::Error(error) => {
+            let PartialParams {
+                client_id,
+                redirect_uri,
+                ..
+            } = q;
+
+            // First, disover the client
+            let client = client_id.and_then(|client_id| {
+                clients
+                    .into_iter()
+                    .find(|client| client.client_id == client_id)
+            });
+
+            let client = match client {
+                Some(client) => client,
+                None => return Ok(Box::new(html(templates.render_error(&error.into())?))),
+            };
+
+            let redirect_uri: Result<Option<Url>, _> = redirect_uri.map(|r| r.parse()).transpose();
+            let redirect_uri = match redirect_uri {
+                Ok(r) => r,
+                Err(_) => return Ok(Box::new(html(templates.render_error(&error.into())?))),
+            };
+
+            let redirect_uri = client.resolve_redirect_uri(&redirect_uri);
+            let redirect_uri = match redirect_uri {
+                Ok(r) => r,
+                Err(_) => return Ok(Box::new(html(templates.render_error(&error.into())?))),
+            };
+
+            let reply: ErrorResponse = error.into();
+            let reply = serde_json::to_value(&reply).wrap_error()?;
+            // TODO: resolve response mode
+            (redirect_uri.clone(), ResponseMode::Query, reply)
+        }
+    };
+
+    // TODO: we should include the state param in errors
+    back_to_client(redirect_uri, response_mode, params, &templates).wrap_error()
+}
+
+async fn get(
+    clients: Vec<OAuth2ClientConfig>,
+    params: Params,
+    maybe_session: Option<SessionInfo>,
+    mut txn: Transaction<'_, Postgres>,
+) -> Result<ReplyOrBackToClient, Rejection> {
+    // First, find out what client it is
+    let client = clients
+        .into_iter()
+        .find(|client| client.client_id == params.auth.client_id)
+        .ok_or_else(|| anyhow::anyhow!("could not find client"))
+        .wrap_error()?;
+
+    let maybe_session_id = maybe_session.as_ref().map(SessionInfo::key);
+
+    let scope: String = {
+        let it = params.auth.scope.iter().map(ToString::to_string);
+        Itertools::intersperse(it, " ".to_string()).collect()
+    };
+
+    let redirect_uri = client
+        .resolve_redirect_uri(&params.auth.redirect_uri)
+        .wrap_error()?;
+    let response_type = &params.auth.response_type;
+    let response_mode =
+        resolve_response_mode(response_type, params.auth.response_mode).wrap_error()?;
+
+    let oauth2_session = start_session(
+        &mut txn,
+        maybe_session_id,
+        &client.client_id,
+        redirect_uri,
+        &scope,
+        params.auth.state.as_deref(),
+        params.auth.nonce.as_deref(),
+        params.auth.max_age,
+        response_type,
+        response_mode,
+    )
+    .await
+    .wrap_error()?;
+
+    // Generate the code at this stage, since we have the PKCE params ready
+    if response_type.contains(&ResponseType::Code) {
+        // 32 random alphanumeric characters, about 190bit of entropy
+        let code: String = thread_rng()
+            .sample_iter(&Alphanumeric)
+            .take(32)
+            .map(char::from)
+            .collect();
+
+        oauth2_session
+            .add_code(&mut txn, &code, &params.pkce)
+            .await
+            .wrap_error()?;
+    };
+
+    // Do we already have a user session for this oauth2 session?
+    let user_session = oauth2_session.fetch_session(&mut txn).await.wrap_error()?;
+
+    if let Some(user_session) = user_session {
+        step(oauth2_session.id, user_session, txn).await
+    } else {
+        // If not, redirect the user to the login page
+        txn.commit().await.wrap_error()?;
+
+        let next = StepRequest::new(oauth2_session.id)
+            .build_uri()
+            .wrap_error()?
+            .to_string();
+
+        let destination = LoginRequest::new(Some(next)).build_uri().wrap_error()?;
+        Ok(ReplyOrBackToClient::Reply(Box::new(see_other(destination))))
+    }
+}
+
+#[derive(Deserialize, Serialize)]
+struct StepRequest {
+    id: i64,
+}
+
+impl StepRequest {
+    fn new(id: i64) -> Self {
+        Self { id }
+    }
+
+    fn build_uri(&self) -> anyhow::Result<Uri> {
+        let qs = serde_urlencoded::to_string(self)?;
+        let path_and_query = PathAndQuery::try_from(format!("/oauth2/authorize/step?{}", qs))?;
+        let uri = Uri::from_parts({
+            let mut parts = Parts::default();
+            parts.path_and_query = Some(path_and_query);
+            parts
+        })?;
+        Ok(uri)
+    }
+}
+
+async fn step(
+    oauth2_session_id: i64,
+    user_session: SessionInfo,
+    mut txn: Transaction<'_, Postgres>,
+) -> Result<ReplyOrBackToClient, Rejection> {
+    let mut oauth2_session = get_session_by_id(&mut txn, oauth2_session_id)
+        .await
+        .wrap_error()?;
+
+    let user_session = oauth2_session
+        .match_or_set_session(&mut txn, user_session)
+        .await
+        .wrap_error()?;
+
+    let response_mode = oauth2_session.response_mode().wrap_error()?;
+    let response_type = oauth2_session.response_type().wrap_error()?;
+    let redirect_uri = oauth2_session.redirect_uri().wrap_error()?;
+
+    // Check if the active session is valid
+    let reply = if user_session.active
+        && user_session.last_authd_at >= oauth2_session.max_auth_time()
+    {
+        // Yep! Let's complete the auth now
+        let mut params = AuthorizationResponse {
+            state: oauth2_session.state.clone(),
+            ..AuthorizationResponse::default()
+        };
+
+        // Did they request an auth code?
+        if response_type.contains(&ResponseType::Code) {
+            params.code = Some(oauth2_session.fetch_code(&mut txn).await.wrap_error()?);
+        }
+
+        // Did they request an access token?
+        if response_type.contains(&ResponseType::Token) {
+            let ttl = Duration::minutes(5);
+            let (access_token, refresh_token) = {
+                let mut rng = thread_rng();
+                (
+                    tokens::generate(&mut rng, tokens::TokenType::AccessToken),
+                    tokens::generate(&mut rng, tokens::TokenType::RefreshToken),
+                )
+            };
+
+            let access_token = add_access_token(&mut txn, oauth2_session_id, &access_token, ttl)
+                .await
+                .wrap_error()?;
+
+            let refresh_token =
+                add_refresh_token(&mut txn, oauth2_session_id, access_token.id, &refresh_token)
+                    .await
+                    .wrap_error()?;
+
+            params.response = Some(
+                AccessTokenResponse::new(access_token.token)
+                    .with_expires_in(ttl)
+                    .with_refresh_token(refresh_token.token),
+            );
+        }
+
+        // Did they request an ID token?
+        if response_type.contains(&ResponseType::IdToken) {
+            todo!("id tokens are not implemented yet");
+        }
+
+        let params = serde_json::to_value(&params).unwrap();
+        ReplyOrBackToClient::BackToClient {
+            redirect_uri,
+            response_mode,
+            params,
+        }
+    } else {
+        // Ask for a reauth
+        // TODO: have the OAuth2 session ID in there
+        ReplyOrBackToClient::Reply(Box::new(see_other(Uri::from_static("/reauth"))))
+    };
+
+    txn.commit().await.wrap_error()?;
+    Ok(reply)
+}

crates/core/src/handlers/oauth2/discovery.rs

@@ -0,0 +1,87 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::HashSet;
+
+use oauth2_types::{
+    oidc::Metadata,
+    requests::{ClientAuthenticationMethod, GrantType, ResponseMode},
+};
+use warp::{Filter, Rejection, Reply};
+
+use crate::config::OAuth2Config;
+
+pub(super) fn filter(
+    config: &OAuth2Config,
+) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone + Send + Sync + 'static {
+    let base = config.issuer.clone();
+
+    let response_modes_supported = Some({
+        let mut s = HashSet::new();
+        s.insert(ResponseMode::FormPost);
+        s.insert(ResponseMode::Query);
+        s.insert(ResponseMode::Fragment);
+        s
+    });
+
+    let response_types_supported = Some({
+        let mut s = HashSet::new();
+        s.insert("code".to_string());
+        s.insert("token".to_string());
+        s.insert("id_token".to_string());
+        s.insert("code token".to_string());
+        s.insert("code id_token".to_string());
+        s.insert("token id_token".to_string());
+        s.insert("code token id_token".to_string());
+        s
+    });
+
+    let grant_types_supported = Some({
+        let mut s = HashSet::new();
+        s.insert(GrantType::AuthorizationCode);
+        s.insert(GrantType::RefreshToken);
+        s
+    });
+
+    let token_endpoint_auth_methods_supported = Some({
+        let mut s = HashSet::new();
+        s.insert(ClientAuthenticationMethod::ClientSecretBasic);
+        s.insert(ClientAuthenticationMethod::ClientSecretPost);
+        s.insert(ClientAuthenticationMethod::None);
+        s
+    });
+
+    let metadata = Metadata {
+        authorization_endpoint: base.join("oauth2/authorize").ok(),
+        token_endpoint: base.join("oauth2/token").ok(),
+        jwks_uri: base.join("oauth2/keys.json").ok(),
+        introspection_endpoint: base.join("oauth2/introspect").ok(),
+        userinfo_endpoint: base.join("oauth2/userinfo").ok(),
+        issuer: base,
+        registration_endpoint: None,
+        scopes_supported: None,
+        response_types_supported,
+        response_modes_supported,
+        grant_types_supported,
+        token_endpoint_auth_methods_supported,
+        code_challenge_methods_supported: None,
+    };
+
+    let cors = warp::cors().allow_any_origin();
+
+    warp::path!(".well-known" / "openid-configuration")
+        .and(warp::get())
+        .map(move || warp::reply::json(&metadata))
+        .with(cors)
+}

crates/core/src/handlers/oauth2/introspection.rs

@@ -0,0 +1,136 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use chrono::Utc;
+use oauth2_types::requests::{IntrospectionRequest, IntrospectionResponse, TokenTypeHint};
+use sqlx::{pool::PoolConnection, PgPool, Postgres};
+use tracing::{info, warn};
+use warp::{Filter, Rejection, Reply};
+
+use crate::{
+    config::{OAuth2ClientConfig, OAuth2Config},
+    errors::WrapError,
+    filters::{
+        client::{with_client_auth, ClientAuthentication},
+        database::with_connection,
+    },
+    storage::oauth2::{access_token::lookup_access_token, refresh_token::lookup_refresh_token},
+    tokens,
+};
+
+pub fn filter(
+    pool: &PgPool,
+    oauth2_config: &OAuth2Config,
+) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone + Send + Sync + 'static {
+    warp::path!("oauth2" / "introspect")
+        .and(warp::post())
+        .and(with_connection(pool))
+        .and(with_client_auth(oauth2_config))
+        .and_then(introspect)
+        .recover(recover)
+}
+
+const INACTIVE: IntrospectionResponse = IntrospectionResponse {
+    active: false,
+    scope: None,
+    client_id: None,
+    username: None,
+    token_type: None,
+    exp: None,
+    iat: None,
+    nbf: None,
+    sub: None,
+    aud: None,
+    iss: None,
+    jti: None,
+};
+
+async fn introspect(
+    mut conn: PoolConnection<Postgres>,
+    auth: ClientAuthentication,
+    client: OAuth2ClientConfig,
+    params: IntrospectionRequest,
+) -> Result<impl Reply, Rejection> {
+    // Token introspection is only allowed by confidential clients
+    if auth.public() {
+        warn!(?client, "Client tried to introspect");
+        // TODO: have a nice error here
+        return Ok(warp::reply::json(&INACTIVE));
+    }
+
+    let token = &params.token;
+    let token_type = tokens::check(token).wrap_error()?;
+    if let Some(hint) = params.token_type_hint {
+        if token_type != hint {
+            info!("Token type hint did not match");
+            return Ok(warp::reply::json(&INACTIVE));
+        }
+    }
+
+    let reply = match token_type {
+        tokens::TokenType::AccessToken => {
+            let token = lookup_access_token(&mut conn, token).await.wrap_error()?;
+            let exp = token.exp();
+
+            // Check it is active and did not expire
+            if !token.active || exp < Utc::now() {
+                info!(?token, "Access token expired");
+                return Ok(warp::reply::json(&INACTIVE));
+            }
+
+            IntrospectionResponse {
+                active: true,
+                scope: None, // TODO: parse back scopes
+                client_id: Some(token.client_id.clone()),
+                username: Some(token.username.clone()),
+                token_type: Some(TokenTypeHint::AccessToken),
+                exp: Some(exp),
+                iat: Some(token.created_at),
+                nbf: Some(token.created_at),
+                sub: None,
+                aud: None,
+                iss: None,
+                jti: None,
+            }
+        }
+        tokens::TokenType::RefreshToken => {
+            let token = lookup_refresh_token(&mut conn, token).await.wrap_error()?;
+
+            IntrospectionResponse {
+                active: true,
+                scope: None, // TODO: parse back scopes
+                client_id: Some(token.client_id),
+                username: None,
+                token_type: Some(TokenTypeHint::RefreshToken),
+                exp: None,
+                iat: None,
+                nbf: None,
+                sub: None,
+                aud: None,
+                iss: None,
+                jti: None,
+            }
+        }
+    };
+
+    Ok(warp::reply::json(&reply))
+}
+
+async fn recover(rejection: Rejection) -> Result<impl Reply, Rejection> {
+    if rejection.is_not_found() {
+        Err(rejection)
+    } else {
+        Ok(warp::reply::json(&INACTIVE))
+    }
+}

crates/core/src/handlers/oauth2/keys.rs

@@ -0,0 +1,30 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use warp::{Filter, Rejection, Reply};
+
+use crate::config::OAuth2Config;
+
+pub(super) fn filter(
+    config: &OAuth2Config,
+) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone + Send + Sync + 'static {
+    let jwks = config.keys.to_public_jwks();
+
+    let cors = warp::cors().allow_any_origin();
+
+    warp::path!("oauth2" / "keys.json")
+        .and(warp::get())
+        .map(move || warp::reply::json(&jwks))
+        .with(cors)
+}

crates/core/src/handlers/oauth2/mod.rs

@@ -0,0 +1,53 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use sqlx::PgPool;
+use warp::{Filter, Rejection, Reply};
+
+use crate::{
+    config::{CookiesConfig, OAuth2Config},
+    templates::Templates,
+};
+
+mod authorization;
+mod discovery;
+mod introspection;
+mod keys;
+mod token;
+mod userinfo;
+
+use self::{
+    authorization::filter as authorization, discovery::filter as discovery,
+    introspection::filter as introspection, keys::filter as keys, token::filter as token,
+    userinfo::filter as userinfo,
+};
+
+pub fn filter(
+    pool: &PgPool,
+    templates: &Templates,
+    oauth2_config: &OAuth2Config,
+    cookies_config: &CookiesConfig,
+) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone + Send + Sync + 'static {
+    discovery(oauth2_config)
+        .or(keys(oauth2_config))
+        .or(authorization(
+            pool,
+            templates,
+            oauth2_config,
+            cookies_config,
+        ))
+        .or(userinfo(pool, oauth2_config))
+        .or(introspection(pool, oauth2_config))
+        .or(token(pool, oauth2_config))
+}

crates/core/src/handlers/oauth2/token.rs

@@ -0,0 +1,276 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use anyhow::Context;
+use chrono::Duration;
+use data_encoding::BASE64URL_NOPAD;
+use headers::{CacheControl, Pragma};
+use hyper::StatusCode;
+use jwt_compact::{Claims, Header, TimeOptions};
+use oauth2_types::{
+    errors::{InvalidGrant, OAuth2Error, OAuth2ErrorCode, UnauthorizedClient},
+    requests::{
+        AccessTokenRequest, AccessTokenResponse, AuthorizationCodeGrant, RefreshTokenGrant,
+    },
+};
+use rand::thread_rng;
+use serde::Serialize;
+use serde_with::skip_serializing_none;
+use sha2::{Digest, Sha256};
+use sqlx::{pool::PoolConnection, Acquire, PgPool, Postgres};
+use url::Url;
+use warp::{
+    reject::Reject,
+    reply::{json, with_status},
+    Filter, Rejection, Reply,
+};
+
+use crate::{
+    config::{KeySet, OAuth2ClientConfig, OAuth2Config},
+    errors::WrapError,
+    filters::{
+        client::{with_client_auth, ClientAuthentication},
+        database::with_connection,
+        headers::typed_header,
+        with_keys,
+    },
+    storage::oauth2::{
+        access_token::{add_access_token, revoke_access_token},
+        authorization_code::lookup_code,
+        refresh_token::{add_refresh_token, lookup_refresh_token, replace_refresh_token},
+    },
+    tokens,
+};
+
+#[skip_serializing_none]
+#[derive(Serialize, Debug)]
+struct CustomClaims {
+    #[serde(rename = "iss")]
+    issuer: Url,
+    #[serde(rename = "sub")]
+    subject: String,
+    #[serde(rename = "aud")]
+    audiences: Vec<String>,
+    nonce: Option<String>,
+    at_hash: String,
+    c_hash: String,
+}
+
+#[derive(Debug)]
+struct Error {
+    json: serde_json::Value,
+    status: StatusCode,
+}
+
+impl Reject for Error {}
+
+fn error<T, E>(e: E) -> Result<T, Rejection>
+where
+    E: OAuth2ErrorCode + 'static,
+{
+    let status = e.status();
+    let json = serde_json::to_value(e.into_response()).wrap_error()?;
+    Err(Error { json, status }.into())
+}
+
+pub fn filter(
+    pool: &PgPool,
+    oauth2_config: &OAuth2Config,
+) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone + Send + Sync + 'static {
+    let issuer = oauth2_config.issuer.clone();
+    warp::path!("oauth2" / "token")
+        .and(warp::post())
+        .and(with_client_auth(oauth2_config))
+        .and(with_keys(oauth2_config))
+        .and(warp::any().map(move || issuer.clone()))
+        .and(with_connection(pool))
+        .and_then(token)
+        .recover(recover)
+}
+
+async fn recover(rejection: Rejection) -> Result<impl Reply, Rejection> {
+    if let Some(Error { json, status }) = rejection.find::<Error>() {
+        Ok(with_status(warp::reply::json(json), *status))
+    } else {
+        Err(rejection)
+    }
+}
+
+async fn token(
+    _auth: ClientAuthentication,
+    client: OAuth2ClientConfig,
+    req: AccessTokenRequest,
+    keys: KeySet,
+    issuer: Url,
+    mut conn: PoolConnection<Postgres>,
+) -> Result<impl Reply, Rejection> {
+    let reply = match req {
+        AccessTokenRequest::AuthorizationCode(grant) => {
+            let reply = authorization_code_grant(&grant, &client, &keys, issuer, &mut conn).await?;
+            json(&reply)
+        }
+        AccessTokenRequest::RefreshToken(grant) => {
+            let reply = refresh_token_grant(&grant, &client, &mut conn).await?;
+            json(&reply)
+        }
+        _ => {
+            let reply = InvalidGrant.into_response();
+            json(&reply)
+        }
+    };
+
+    Ok(typed_header(
+        Pragma::no_cache(),
+        typed_header(CacheControl::new().with_no_store(), reply),
+    ))
+}
+
+fn hash<H: Digest>(mut hasher: H, token: &str) -> anyhow::Result<String> {
+    hasher.update(token);
+    let hash = hasher.finalize();
+    // Left-most 128bit
+    let bits = hash
+        .get(..16)
+        .context("failed to get first 128 bits of hash")?;
+    Ok(BASE64URL_NOPAD.encode(bits))
+}
+
+async fn authorization_code_grant(
+    grant: &AuthorizationCodeGrant,
+    client: &OAuth2ClientConfig,
+    keys: &KeySet,
+    issuer: Url,
+    conn: &mut PoolConnection<Postgres>,
+) -> Result<AccessTokenResponse, Rejection> {
+    let mut txn = conn.begin().await.wrap_error()?;
+    let code = lookup_code(&mut txn, &grant.code).await.wrap_error()?;
+    if client.client_id != code.client_id {
+        return error(UnauthorizedClient);
+    }
+
+    // TODO: verify PKCE
+    // TODO: make the code invalid
+    let ttl = Duration::minutes(5);
+    let (access_token, refresh_token) = {
+        let mut rng = thread_rng();
+        (
+            tokens::generate(&mut rng, tokens::TokenType::AccessToken),
+            tokens::generate(&mut rng, tokens::TokenType::RefreshToken),
+        )
+    };
+
+    let access_token = add_access_token(&mut txn, code.oauth2_session_id, &access_token, ttl)
+        .await
+        .wrap_error()?;
+
+    let refresh_token = add_refresh_token(
+        &mut txn,
+        code.oauth2_session_id,
+        access_token.id,
+        &refresh_token,
+    )
+    .await
+    .wrap_error()?;
+
+    // TODO: generate id_token only if the "openid" scope was asked
+    let header = Header::default();
+    let options = TimeOptions::default();
+    let claims = Claims::new(CustomClaims {
+        issuer,
+        // TODO: get that from the session
+        subject: "random-subject".to_string(),
+        audiences: vec![client.client_id.clone()],
+        nonce: code.nonce,
+        at_hash: hash(Sha256::new(), &access_token.token).wrap_error()?,
+        c_hash: hash(Sha256::new(), &grant.code).wrap_error()?,
+    })
+    .set_duration_and_issuance(&options, Duration::minutes(30));
+    let id_token = keys
+        .token(crate::config::Algorithm::Rs256, header, claims)
+        .await
+        .context("could not sign ID token")
+        .wrap_error()?;
+
+    // TODO: have the scopes back here
+    let params = AccessTokenResponse::new(access_token.token)
+        .with_expires_in(ttl)
+        .with_refresh_token(refresh_token.token)
+        .with_id_token(id_token);
+
+    txn.commit().await.wrap_error()?;
+
+    Ok(params)
+}
+
+async fn refresh_token_grant(
+    grant: &RefreshTokenGrant,
+    client: &OAuth2ClientConfig,
+    conn: &mut PoolConnection<Postgres>,
+) -> Result<AccessTokenResponse, Rejection> {
+    let mut txn = conn.begin().await.wrap_error()?;
+    // TODO: scope handling
+    let refresh_token_lookup = lookup_refresh_token(&mut txn, &grant.refresh_token)
+        .await
+        .wrap_error()?;
+
+    if client.client_id != refresh_token_lookup.client_id {
+        // As per https://datatracker.ietf.org/doc/html/rfc6749#section-5.2
+        return error(InvalidGrant);
+    }
+
+    let ttl = Duration::minutes(5);
+    let (access_token, refresh_token) = {
+        let mut rng = thread_rng();
+        (
+            tokens::generate(&mut rng, tokens::TokenType::AccessToken),
+            tokens::generate(&mut rng, tokens::TokenType::RefreshToken),
+        )
+    };
+
+    let access_token = add_access_token(
+        &mut txn,
+        refresh_token_lookup.oauth2_session_id,
+        &access_token,
+        ttl,
+    )
+    .await
+    .wrap_error()?;
+
+    let refresh_token = add_refresh_token(
+        &mut txn,
+        refresh_token_lookup.oauth2_session_id,
+        access_token.id,
+        &refresh_token,
+    )
+    .await
+    .wrap_error()?;
+
+    replace_refresh_token(&mut txn, refresh_token_lookup.id, refresh_token.id)
+        .await
+        .wrap_error()?;
+
+    if let Some(access_token_id) = refresh_token_lookup.oauth2_access_token_id {
+        revoke_access_token(&mut txn, access_token_id)
+            .await
+            .wrap_error()?;
+    }
+
+    let params = AccessTokenResponse::new(access_token.token)
+        .with_expires_in(ttl)
+        .with_refresh_token(refresh_token.token);
+
+    txn.commit().await.wrap_error()?;
+
+    Ok(params)
+}

crates/core/src/handlers/oauth2/userinfo.rs

@@ -0,0 +1,43 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use serde::Serialize;
+use sqlx::PgPool;
+use warp::{Filter, Rejection, Reply};
+
+use crate::{
+    config::OAuth2Config, filters::authenticate::with_authentication,
+    storage::oauth2::access_token::OAuth2AccessTokenLookup,
+};
+
+#[derive(Serialize)]
+struct UserInfo {
+    sub: String,
+}
+
+pub(super) fn filter(
+    pool: &PgPool,
+    _config: &OAuth2Config,
+) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone + Send + Sync + 'static {
+    warp::path!("oauth2" / "userinfo")
+        .and(warp::get().or(warp::post()).unify())
+        .and(with_authentication(pool))
+        .and_then(userinfo)
+}
+
+async fn userinfo(token: OAuth2AccessTokenLookup) -> Result<impl Reply, Rejection> {
+    Ok(warp::reply::json(&UserInfo {
+        sub: token.username,
+    }))
+}