add extractor for jwt token

3 files changed, 89 insertions, 5 deletions

diff --git a/src/extractors.rs b/src/extractors.rs
index 6aa372b..7954cb2 100644
--- a/src/extractors.rs
+++ b/src/extractors.rs
@@ -1,17 +1,97 @@
 use axum::async_trait;
-use axum::extract::{rejection::JsonRejection, FromRequest, RequestParts};
+use axum::extract::FromRequest;
+use axum::extract::{rejection::JsonRejection, RequestParts};
 use axum::http::header::{self, HeaderValue};
 use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
 use axum::BoxError;
+use hmac::{Hmac, Mac};
+use jwt::VerifyWithKey;
+use lazy_static::lazy_static;
 use serde::de::DeserializeOwned;
-use serde::Serialize;
+use serde::{Deserialize, Serialize};
 use serde_json::json;
 use serde_json::Value;
+use sha2::Sha256;
+use std::collections::BTreeMap;
+use std::env;
 use std::error::Error;
 
+lazy_static! {
+    static ref JWT_SECRET: String = env::var("JWT_SECRET").unwrap();
+}
+
 pub struct Json<T>(pub T);
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Token {
+    pub iss: String,
+    pub sub: String,
+    pub iat: i64,
+    pub exp: i64,
+    pub dn: String,
+    pub email: String,
+    pub admin: bool,
+}
+pub struct Jwt(pub Token);
+
+#[async_trait]
+impl<B> FromRequest<B> for Jwt
+where
+    B: axum::body::HttpBody + Send,
+    B::Data: Send,
+    B::Error: Into<BoxError>,
+{
+    type Rejection = (axum::http::StatusCode, axum::Json<serde_json::Value>);
+
+    async fn from_request(
+        req: &mut axum::extract::RequestParts<B>,
+    ) -> Result<Self, Self::Rejection> {
+        // Grab the "Authorization" header from the request
+        let auth_header = req
+            .headers()
+            .get(axum::http::header::AUTHORIZATION)
+            .and_then(|value| value.to_str().ok());
+
+        match auth_header {
+            Some(header) => {
+                let key: Hmac<Sha256> = Hmac::new_from_slice((*JWT_SECRET).as_bytes()).unwrap();
+                let token = header.replace("Bearer ", "");
+                let claims: BTreeMap<String, String> = match token.verify_with_key(&key) {
+                    Ok(claims) => claims,
+                    Err(_) => return Err((StatusCode::UNAUTHORIZED, axum::Json(json!({ "error": "Invalid token" })))),
+                };
+
+                let token = Token {
+                    iss: claims.get("iss").unwrap().to_string(),
+                    sub: claims.get("sub").unwrap().to_string(),
+                    iat: claims.get("iat").unwrap().parse().unwrap(),
+                    exp: claims.get("exp").unwrap().parse().unwrap(),
+                    dn: claims.get("dn").unwrap().to_string(),
+                    email: claims.get("email").unwrap().to_string(),
+                    admin: claims.get("admin").unwrap().parse().unwrap(),
+                };
+
+                let now = chrono::Utc::now().timestamp();
+                if token.iat > now || token.exp < now {
+                    return Err((
+                        StatusCode::UNAUTHORIZED,
+                        axum::Json(json!({"error": "Invalid token"})),
+                    ));
+                }
+
+                return Ok(Self(token));
+            }
+            None => {
+                return Err((
+                    StatusCode::UNAUTHORIZED,
+                    axum::Json(json!({"error": "missing auth header"})),
+                ))
+            }
+        }
+    }
+}
+
 #[async_trait]
 impl<B, T> FromRequest<B> for Json<T>
 where
diff --git a/src/main.rs b/src/main.rs
index 0c2081e..1816574 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -51,7 +51,8 @@ async fn main() {
                         .route("/", post(routes::v1::users::create_user))
                         .route("/all", get(routes::v1::users::get_all_users))
                         .route("/totp", get(routes::v1::users::needs_totp))
-                        .route("/login", post(routes::v1::users::login)),
+                        .route("/login", post(routes::v1::users::login))
+                        .route("/whoami", get(routes::v1::users::whoami)),
                 ),
             ),
         )
diff --git a/src/routes/v1/users.rs b/src/routes/v1/users.rs
index 3a84c6b..5fb1595 100644
--- a/src/routes/v1/users.rs
+++ b/src/routes/v1/users.rs
@@ -1,11 +1,10 @@
 use crate::db;
-use crate::extractors;
+use crate::extractors::{Json, Jwt};
 use crate::routes::v1::requests;
 use axum::extract::Query;
 use axum::http::StatusCode;
 use axum::response::IntoResponse;
 use axum::Extension;
-use extractors::Json;
 use hmac::{Hmac, Mac};
 use jwt::SignWithKey;
 use lazy_static::lazy_static;
@@ -55,6 +54,10 @@ pub async fn get_all_users(Extension(pool): Extension<Arc<Pool<Postgres>>>) -> i
     }
 }
 
+pub async fn whoami(Jwt(user): Jwt) -> impl IntoResponse {
+    (StatusCode::OK, Json(json!(user)))
+}
+
 pub async fn needs_totp(
     Query(params): Query<HashMap<String, String>>,
     Extension(pool): Extension<Arc<Pool<Postgres>>>,