From 6c94dc48743be29bfc991b1ba2518a5a6f8fcb59 Mon Sep 17 00:00:00 2001
From: Quentin Gliech <quenting@element.io>
Date: Wed, 26 Mar 2025 16:53:29 +0100
Subject: [PATCH] Record metrics for OAuth 2.0 token requests

---
 crates/handlers/src/oauth2/token.rs | 27 +++++++++++++++++++++++++--
 crates/oauth2-types/src/requests.rs | 14 ++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/crates/handlers/src/oauth2/token.rs b/crates/handlers/src/oauth2/token.rs
index d76c81f1b..4ed02ef8d 100644
--- a/crates/handlers/src/oauth2/token.rs
+++ b/crates/handlers/src/oauth2/token.rs
@@ -4,7 +4,7 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Please see LICENSE in the repository root for full details.
 
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
 
 use axum::{Json, extract::State, response::IntoResponse};
 use axum_extra::typed_header::TypedHeader;
@@ -40,12 +40,23 @@ use oauth2_types::{
     },
     scope,
 };
+use opentelemetry::{Key, KeyValue, metrics::Counter};
 use thiserror::Error;
 use tracing::{debug, info};
 use ulid::Ulid;
 
 use super::{generate_id_token, generate_token_pair};
-use crate::{BoundActivityTracker, impl_from_error_for_route};
+use crate::{BoundActivityTracker, METER, impl_from_error_for_route};
+
+static TOKEN_REQUEST_COUNTER: LazyLock<Counter<u64>> = LazyLock::new(|| {
+    METER
+        .u64_counter("mas.oauth2.token_request")
+        .with_description("How many OAuth 2.0 token requests have gone through")
+        .with_unit("{request}")
+        .build()
+});
+const GRANT_TYPE: Key = Key::from_static_str("grant_type");
+const RESULT: Key = Key::from_static_str("successful");
 
 #[derive(Debug, Error)]
 pub(crate) enum RouteError {
@@ -136,6 +147,8 @@ impl IntoResponse for RouteError {
     fn into_response(self) -> axum::response::Response {
         let event_id = sentry::capture_error(&self);
 
+        TOKEN_REQUEST_COUNTER.add(1, &[KeyValue::new(RESULT, "error")]);
+
         let response = match self {
             Self::Internal(_)
             | Self::NoSuchBrowserSession
@@ -254,6 +267,8 @@ pub(crate) async fn post(
 
     let form = client_authorization.form.ok_or(RouteError::BadRequest)?;
 
+    let grant_type = form.grant_type();
+
     let (reply, repo) = match form {
         AccessTokenRequest::AuthorizationCode(grant) => {
             authorization_code_grant(
@@ -321,6 +336,14 @@ pub(crate) async fn post(
 
     repo.save().await?;
 
+    TOKEN_REQUEST_COUNTER.add(
+        1,
+        &[
+            KeyValue::new(GRANT_TYPE, grant_type),
+            KeyValue::new(RESULT, "success"),
+        ],
+    );
+
     let mut headers = HeaderMap::new();
     headers.typed_insert(CacheControl::new().with_no_store());
     headers.typed_insert(Pragma::no_cache());
diff --git a/crates/oauth2-types/src/requests.rs b/crates/oauth2-types/src/requests.rs
index 631b33309..36ee36da6 100644
--- a/crates/oauth2-types/src/requests.rs
+++ b/crates/oauth2-types/src/requests.rs
@@ -638,6 +638,20 @@ pub enum AccessTokenRequest {
     Unsupported,
 }
 
+impl AccessTokenRequest {
+    /// Returns the string representation of the grant type of the request.
+    #[must_use]
+    pub fn grant_type(&self) -> &'static str {
+        match self {
+            Self::AuthorizationCode(_) => "authorization_code",
+            Self::RefreshToken(_) => "refresh_token",
+            Self::ClientCredentials(_) => "client_credentials",
+            Self::DeviceCode(_) => "urn:ietf:params:oauth:grant-type:device_code",
+            Self::Unsupported => "unsupported",
+        }
+    }
+}
+
 /// A successful response from the [Token Endpoint].
 ///
 /// [Token Endpoint]: https://www.rfc-editor.org/rfc/rfc6749#section-3.2