From dab90bc4a9d160a0c3f072cdae003b214dc7a174 Mon Sep 17 00:00:00 2001
From: Jef Roosens <roosensjef@gmail.com>
Date: Sat, 21 Aug 2021 22:21:42 +0200
Subject: [PATCH] Separated JWT header into own guard

---
 Rocket.toml       |  2 +-
 src/rb/auth.rs    | 32 +++++---------------------------
 src/rb/lib.rs     |  8 ++++++++
 src/rbs/guards.rs | 19 ++++++++++++++++---
 4 files changed, 30 insertions(+), 31 deletions(-)

diff --git a/Rocket.toml b/Rocket.toml
index 081ed68..e931e4d 100644
--- a/Rocket.toml
+++ b/Rocket.toml
@@ -10,4 +10,4 @@ limits = { forms = 32768 }
 postgres_rb = { url = "postgres://rb:rb@localhost:5432/rb" }
 
 [release.databases]
-postgres_rb = { url = "postgres://rb:rb@db:5432/rb" }
+postgres_rb = { url = "postgres://rb:rb@localhost:5432/rb" }
diff --git a/src/rb/auth.rs b/src/rb/auth.rs
index da60c78..4951e1c 100644
--- a/src/rb/auth.rs
+++ b/src/rb/auth.rs
@@ -11,20 +11,6 @@ use jwt::SignWithKey;
 use rand::{thread_rng, Rng};
 use serde::{Deserialize, Serialize};
 use sha2::Sha256;
-use std::collections::HashMap;
-
-/// Expire time for the JWT tokens in seconds.
-const JWT_EXP_SECONDS: i64 = 900;
-/// Amount of bytes the refresh tokens should consist of
-const REFRESH_TOKEN_N_BYTES: usize = 64;
-/// Expire time for refresh tokens; here: one week
-const REFRESH_TOKEN_EXP_SECONDS: i64 = 604800;
-
-fn log<T>(message: &str, o: T) -> T {
-    println!("{}", message);
-
-    o
-}
 
 pub fn verify_user(conn: &PgConnection, username: &str, password: &str) -> crate::Result<User> {
     // TODO handle non-"NotFound" Diesel errors accordingely
@@ -62,7 +48,7 @@ pub struct Claims {
 pub fn generate_jwt_token(conn: &PgConnection, user: &User) -> crate::Result<JWTResponse> {
     let secret = std::env::var("JWT_KEY").map_err(|_| RBError::MissingJWTKey)?;
     let key: Hmac<Sha256> = Hmac::new_from_slice(secret.as_bytes())
-        .map_err(|_| log("Failed to create key", RBError::JWTCreationError))?;
+        .map_err(|_| RBError::JWTCreationError)?;
 
     let current_time = Utc::now();
 
@@ -71,28 +57,20 @@ pub fn generate_jwt_token(conn: &PgConnection, user: &User) -> crate::Result<JWT
         id: user.id,
         username: user.username.clone(),
         admin: user.admin,
-        exp: current_time.timestamp() + JWT_EXP_SECONDS,
+        exp: current_time.timestamp() + crate::JWT_EXP_SECONDS,
     };
-    // let mut claims = HashMap::new();
-    // claims.insert("id", user.id.to_string());
-    // claims.insert("username", user.username.clone());
-    // claims.insert("admin", user.admin.to_string());
-    // claims.insert(
-    //     "exp",
-    //     (current_time.timestamp() + JWT_EXP_SECONDS).to_string(),
-    // );
 
     // Sign the claims into a new token
     let token = claims
         .sign_with_key(&key)
-        .map_err(|_| log("Failed to sign token", RBError::JWTCreationError))?;
+        .map_err(|_| RBError::JWTCreationError)?;
 
     // Generate a random refresh token
-    let mut refresh_token = [0u8; REFRESH_TOKEN_N_BYTES];
+    let mut refresh_token = [0u8; crate::REFRESH_TOKEN_N_BYTES];
     thread_rng().fill(&mut refresh_token[..]);
 
     let refresh_expire =
-        (current_time + chrono::Duration::seconds(REFRESH_TOKEN_EXP_SECONDS)).naive_utc();
+        (current_time + chrono::Duration::seconds(crate::REFRESH_TOKEN_EXP_SECONDS)).naive_utc();
 
     // Store refresh token in database
     // TODO add expires_at here (it's what's causing the errors)
diff --git a/src/rb/lib.rs b/src/rb/lib.rs
index 13c8330..e3ca322 100644
--- a/src/rb/lib.rs
+++ b/src/rb/lib.rs
@@ -7,3 +7,11 @@ mod models;
 pub(crate) mod schema;
 
 pub use errors::Result;
+
+// Any import defaults are defined here
+/// Expire time for the JWT tokens in seconds.
+const JWT_EXP_SECONDS: i64 = 900;
+/// Amount of bytes the refresh tokens should consist of
+const REFRESH_TOKEN_N_BYTES: usize = 64;
+/// Expire time for refresh tokens; here: one week
+const REFRESH_TOKEN_EXP_SECONDS: i64 = 604800;
diff --git a/src/rbs/guards.rs b/src/rbs/guards.rs
index 95e3ef0..96891d9 100644
--- a/src/rbs/guards.rs
+++ b/src/rbs/guards.rs
@@ -8,10 +8,10 @@ use jwt::VerifyWithKey;
 use rb::auth::Claims;
 use sha2::Sha256;
 
-pub struct User(Claims);
+pub struct JWT(String);
 
 #[rocket::async_trait]
-impl<'r> FromRequest<'r> for User {
+impl<'r> FromRequest<'r> for JWT {
     type Error = rb::errors::RBError;
 
     async fn from_request(req: &'r Request<'_>) -> Outcome<Self, Self::Error> {
@@ -28,9 +28,22 @@ impl<'r> FromRequest<'r> for User {
         // Extract the jwt token from the header
         let jwt_token = match header.get(7..) {
             Some(token) => token,
-            None => return Outcome::Failure((Status::Unauthorized, Self::Error::JWTError)),
+            None => return Outcome::Forward(()),
         };
 
+        Outcome::Success(Self(jwt_token.to_string()))
+    }
+}
+
+pub struct User(Claims);
+
+#[rocket::async_trait]
+impl<'r> FromRequest<'r> for User {
+    type Error = rb::errors::RBError;
+
+    async fn from_request(req: &'r Request<'_>) -> Outcome<Self, Self::Error> {
+        let jwt_token = try_outcome!(req.guard::<JWT>().await).0;
+
         // Get secret & key
         let secret = match std::env::var("JWT_KEY") {
             Ok(key) => key,