Implement the device consent logic

crates/handlers/src/lib.rs

@@ -412,6 +412,10 @@ where
             mas_router::DeviceCodeLink::route(),
             get(self::oauth2::device::link::get).post(self::oauth2::device::link::post),
         )
+        .route(
+            mas_router::DeviceCodeConsent::route(),
+            get(self::oauth2::device::consent::get).post(self::oauth2::device::consent::post),
+        )
         .layer(AndThenLayer::new(
             move |response: axum::response::Response| async move {
                 if response.status().is_server_error() {

crates/handlers/src/oauth2/device/authorize.rs

@@ -155,8 +155,7 @@ pub(crate) async fn post(
         TypedHeader(CacheControl::new().with_no_store()),
         TypedHeader(Pragma::no_cache()),
         Json(response),
-    )
-        .into_response())
+    ))
 }
 
 #[cfg(test)]

crates/handlers/src/oauth2/device/consent.rs

@@ -0,0 +1,184 @@
+// Copyright 2023 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 axum::{
+    extract::{Path, State},
+    response::{IntoResponse, Response},
+    Form,
+};
+use axum_extra::response::Html;
+use mas_axum_utils::{
+    cookies::CookieJar,
+    csrf::{CsrfExt, ProtectedForm},
+    FancyError, SessionInfoExt,
+};
+use mas_router::UrlBuilder;
+use mas_storage::{BoxClock, BoxRepository, BoxRng};
+use mas_templates::{DeviceConsentContext, TemplateContext, Templates};
+use serde::Deserialize;
+use tracing::warn;
+use ulid::Ulid;
+
+use crate::{BoundActivityTracker, PreferredLanguage};
+
+#[derive(Deserialize, Debug)]
+#[serde(rename_all = "lowercase")]
+enum Action {
+    Consent,
+    Reject,
+}
+
+#[derive(Deserialize, Debug)]
+pub(crate) struct ConsentForm {
+    action: Action,
+}
+
+pub(crate) async fn get(
+    mut rng: BoxRng,
+    clock: BoxClock,
+    PreferredLanguage(locale): PreferredLanguage,
+    State(templates): State<Templates>,
+    State(url_builder): State<UrlBuilder>,
+    mut repo: BoxRepository,
+    activity_tracker: BoundActivityTracker,
+    cookie_jar: CookieJar,
+    Path(grant_id): Path<Ulid>,
+) -> Result<Response, FancyError> {
+    let (session_info, cookie_jar) = cookie_jar.session_info();
+    let (csrf_token, cookie_jar) = cookie_jar.csrf_token(&clock, &mut rng);
+
+    let maybe_session = session_info.load_session(&mut repo).await?;
+
+    let Some(session) = maybe_session else {
+        let login = mas_router::Login::and_continue_device_code_grant(grant_id);
+        return Ok((cookie_jar, url_builder.redirect(&login)).into_response());
+    };
+
+    activity_tracker
+        .record_browser_session(&clock, &session)
+        .await;
+
+    // TODO: better error handling
+    let grant = repo
+        .oauth2_device_code_grant()
+        .lookup(grant_id)
+        .await?
+        .context("Device grant not found")?;
+
+    if grant.expires_at < clock.now() {
+        return Err(FancyError::from(anyhow::anyhow!("Grant is expired")));
+    }
+
+    let client = repo
+        .oauth2_client()
+        .lookup(grant.client_id)
+        .await?
+        .context("Client not found")?;
+
+    let ctx = DeviceConsentContext::new(grant, client)
+        .with_session(session)
+        .with_csrf(csrf_token.form_value())
+        .with_language(locale);
+
+    let rendered = templates
+        .render_device_consent(&ctx)
+        .context("Failed to render template")?;
+
+    Ok((cookie_jar, Html(rendered)).into_response())
+}
+
+pub(crate) async fn post(
+    mut rng: BoxRng,
+    clock: BoxClock,
+    PreferredLanguage(locale): PreferredLanguage,
+    State(templates): State<Templates>,
+    State(url_builder): State<UrlBuilder>,
+    mut repo: BoxRepository,
+    activity_tracker: BoundActivityTracker,
+    cookie_jar: CookieJar,
+    Path(grant_id): Path<Ulid>,
+    Form(form): Form<ProtectedForm<ConsentForm>>,
+) -> Result<Response, FancyError> {
+    let (session_info, cookie_jar) = cookie_jar.session_info();
+    let form = cookie_jar.verify_form(&clock, form)?;
+    let (csrf_token, cookie_jar) = cookie_jar.csrf_token(&clock, &mut rng);
+
+    let maybe_session = session_info.load_session(&mut repo).await?;
+
+    let Some(session) = maybe_session else {
+        let login = mas_router::Login::and_continue_device_code_grant(grant_id);
+        return Ok((cookie_jar, url_builder.redirect(&login)).into_response());
+    };
+
+    activity_tracker
+        .record_browser_session(&clock, &session)
+        .await;
+
+    // TODO: better error handling
+    let grant = repo
+        .oauth2_device_code_grant()
+        .lookup(grant_id)
+        .await?
+        .context("Device grant not found")?;
+
+    if grant.expires_at < clock.now() {
+        return Err(FancyError::from(anyhow::anyhow!("Grant is expired")));
+    }
+
+    let client = repo
+        .oauth2_client()
+        .lookup(grant.client_id)
+        .await?
+        .context("Client not found")?;
+
+    // TODO: run through the policy
+    let grant = if grant.is_pending() {
+        match form.action {
+            Action::Consent => {
+                repo.oauth2_device_code_grant()
+                    .fulfill(&clock, grant, &session)
+                    .await?
+            }
+            Action::Reject => {
+                repo.oauth2_device_code_grant()
+                    .reject(&clock, grant, &session)
+                    .await?
+            }
+        }
+    } else {
+        // XXX: In case we're not pending, let's just return the grant as-is
+        // since it might just be a form resubmission, and feedback is nice enough
+        warn!(
+            oauth2_device_code.id = %grant.id,
+            browser_session.id = %session.id,
+            user.id = %session.user.id,
+            "Grant is not pending",
+        );
+        grant
+    };
+
+    repo.save().await?;
+
+    let ctx = DeviceConsentContext::new(grant, client)
+        .with_session(session)
+        .with_csrf(csrf_token.form_value())
+        .with_language(locale);
+
+    let rendered = templates
+        .render_device_consent(&ctx)
+        .context("Failed to render template")?;
+
+    Ok((cookie_jar, Html(rendered)).into_response())
+}

crates/handlers/src/oauth2/device/link.rs

@@ -14,7 +14,7 @@
 
 use axum::{
     extract::{Query, State},
-    response::IntoResponse,
+    response::{IntoResponse, Response},
     Form,
 };
 use axum_extra::response::Html;
@@ -23,7 +23,8 @@ use mas_axum_utils::{
     csrf::{CsrfExt, ProtectedForm},
     FancyError,
 };
-use mas_storage::{BoxClock, BoxRng};
+use mas_router::UrlBuilder;
+use mas_storage::{BoxClock, BoxRepository, BoxRng};
 use mas_templates::{
     DeviceLinkContext, DeviceLinkFormField, FieldError, FormState, TemplateContext, Templates,
 };
@@ -76,27 +77,42 @@ pub(crate) async fn get(
 pub(crate) async fn post(
     mut rng: BoxRng,
     clock: BoxClock,
+    mut repo: BoxRepository,
     PreferredLanguage(locale): PreferredLanguage,
     State(templates): State<Templates>,
+    State(url_builder): State<UrlBuilder>,
     cookie_jar: CookieJar,
     Form(form): Form<ProtectedForm<Params>>,
-) -> Result<impl IntoResponse, FancyError> {
+) -> Result<Response, FancyError> {
     let form = cookie_jar.verify_form(&clock, form)?;
     let (csrf_token, cookie_jar) = cookie_jar.csrf_token(&clock, &mut rng);
 
-    let form_state = FormState::from_form(&form)
-        .with_error_on_field(DeviceLinkFormField::Code, FieldError::Required);
+    let code = form.code.to_uppercase();
+    let grant = repo
+        .oauth2_device_code_grant()
+        .find_by_user_code(&code)
+        .await?
+        // XXX: We should have different error messages for already exchanged and expired
+        .filter(|grant| grant.is_pending())
+        .filter(|grant| grant.expires_at > clock.now());
 
-    // TODO: find the device code grant in the database
-    // and redirect to /oauth2/device/link/:id
-    // That then will trigger a login if we don't have a session
+    let Some(grant) = grant else {
+        let form_state = FormState::from_form(&form)
+            .with_error_on_field(DeviceLinkFormField::Code, FieldError::Invalid);
 
-    let ctx = DeviceLinkContext::new()
-        .with_form_state(form_state)
-        .with_csrf(csrf_token.form_value())
-        .with_language(locale);
+        let ctx = DeviceLinkContext::new()
+            .with_form_state(form_state)
+            .with_csrf(csrf_token.form_value())
+            .with_language(locale);
 
-    let content = templates.render_device_link(&ctx)?;
+        let content = templates.render_device_link(&ctx)?;
 
-    Ok((cookie_jar, Html(content)))
+        return Ok((cookie_jar, Html(content)).into_response());
+    };
+
+    // Redirect to the consent page
+    // This will in turn redirect to the login page if the user is not logged in
+    let destination = url_builder.redirect(&mas_router::DeviceCodeConsent::new(grant.id));
+
+    Ok((cookie_jar, destination).into_response())
 }

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

@@ -13,4 +13,5 @@
 // limitations under the License.
 
 pub mod authorize;
+pub mod consent;
 pub mod link;

crates/handlers/src/views/shared.rs

@@ -63,6 +63,16 @@ impl OptionalPostAuthAction {
                 PostAuthContextInner::ContinueAuthorizationGrant { grant }
             }
 
+            PostAuthAction::ContinueDeviceCodeGrant { id } => {
+                let grant = repo
+                    .oauth2_device_code_grant()
+                    .lookup(id)
+                    .await?
+                    .context("Failed to load device code grant")?;
+                let grant = Box::new(grant);
+                PostAuthContextInner::ContinueDeviceCodeGrant { grant }
+            }
+
             PostAuthAction::ContinueCompatSsoLogin { id } => {
                 let login = repo
                     .compat_sso_login()