From 66249a5e82e94dccb9cb685fce9058aa81f634df Mon Sep 17 00:00:00 2001
From: traxys <quentin+dev@familleboyer.net>
Date: Sun, 28 May 2023 20:58:39 +0200
Subject: [PATCH] server: Add a extractor for the authenticated user

---
 Cargo.lock        | 26 +++++++++++++++++++++++
 Cargo.toml        |  2 +-
 src/routes/mod.rs | 54 ++++++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 49b5506..8fc9c5a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -267,6 +267,7 @@ dependencies = [
  "bitflags",
  "bytes",
  "futures-util",
+ "headers",
  "http",
  "http-body",
  "hyper",
@@ -1239,6 +1240,31 @@ dependencies = [
  "hashbrown 0.13.2",
 ]
 
+[[package]]
+name = "headers"
+version = "0.3.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f3e372db8e5c0d213e0cd0b9be18be2aca3d44cf2fe30a9d46a65581cd454584"
+dependencies = [
+ "base64 0.13.1",
+ "bitflags",
+ "bytes",
+ "headers-core",
+ "http",
+ "httpdate",
+ "mime",
+ "sha1",
+]
+
+[[package]]
+name = "headers-core"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e7f66481bfee273957b1f20485a4ff3362987f85b2c236580d81b4eb7a326429"
+dependencies = [
+ "http",
+]
+
 [[package]]
 name = "heck"
 version = "0.3.3"
diff --git a/Cargo.toml b/Cargo.toml
index 868bc97..814c01e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -9,7 +9,7 @@ members = [".", "api", "app", "migration"]
 
 [dependencies]
 anyhow = "1.0.71"
-axum = "0.6.18"
+axum = { version = "0.6.18", features = ["headers"] }
 base64 = "0.21.0"
 config = "0.13.3"
 jwt-simple = "0.11.5"
diff --git a/src/routes/mod.rs b/src/routes/mod.rs
index 815a1d6..c275062 100644
--- a/src/routes/mod.rs
+++ b/src/routes/mod.rs
@@ -2,11 +2,13 @@ use std::sync::Arc;
 
 use api::{LoginRequest, LoginResponse};
 use axum::{
-    extract::State,
-    http::{header::CONTENT_TYPE, HeaderValue, Method, StatusCode},
+    async_trait,
+    extract::{FromRef, FromRequestParts, State},
+    headers::{authorization::Bearer, Authorization},
+    http::{header::CONTENT_TYPE, request::Parts, HeaderValue, Method, StatusCode},
     response::IntoResponse,
     routing::post,
-    Json, Router,
+    Json, Router, TypedHeader,
 };
 use jwt_simple::prelude::*;
 use sea_orm::prelude::*;
@@ -23,6 +25,10 @@ enum RouteError {
     Db(#[from] DbErr),
     #[error("JWT error encountered")]
     Jwt(#[from] jwt_simple::Error),
+    #[error("User provided JWT token is invalid")]
+    UserJwt(jwt_simple::Error),
+    #[error("Request is missing the bearer token")]
+    MissingAuthorization,
 }
 
 impl IntoResponse for RouteError {
@@ -31,6 +37,13 @@ impl IntoResponse for RouteError {
             RouteError::UnknownAccount => {
                 (StatusCode::NOT_FOUND, "Account not found").into_response()
             }
+            RouteError::MissingAuthorization => {
+                (StatusCode::BAD_REQUEST, "Missing authorization header").into_response()
+            }
+            RouteError::UserJwt(e) => {
+                tracing::debug!("Invalid user JWT: {e:?}");
+                (StatusCode::BAD_REQUEST, "Invalid authorization header").into_response()
+            }
             e => {
                 tracing::error!("Internal error: {e:?}");
                 StatusCode::INTERNAL_SERVER_ERROR.into_response()
@@ -43,6 +56,41 @@ type JsonResult<T, E = RouteError> = Result<Json<T>, E>;
 
 type AppState = Arc<crate::AppState>;
 
+#[derive(Debug)]
+struct AuthenticatedUser {
+    pub id: Uuid,
+}
+
+#[async_trait]
+impl<S> FromRequestParts<S> for AuthenticatedUser
+where
+    S: Send + Sync,
+    AppState: FromRef<S>,
+{
+    type Rejection = RouteError;
+
+    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
+        let State(app_state): State<AppState> = State::from_request_parts(parts, state)
+            .await
+            .expect("Could not get state");
+
+        let TypedHeader(Authorization(bearer)) =
+            TypedHeader::<Authorization<Bearer>>::from_request_parts(parts, state)
+                .await
+                .map_err(|_| RouteError::MissingAuthorization)?;
+
+        let claims = app_state
+            .jwt_secret
+            .0
+            .verify_token::<NoCustomClaims>(bearer.token(), None)
+            .map_err(RouteError::UserJwt)?;
+
+        Ok(AuthenticatedUser {
+            id: claims.subject.unwrap().parse().unwrap(),
+        })
+    }
+}
+
 async fn login(
     State(state): State<AppState>,
     Json(req): Json<LoginRequest>,