From cc7a668ab0d00d84148bf31a8822eaac704ec37e Mon Sep 17 00:00:00 2001
From: Jef Roosens <roosensjef@gmail.com>
Date: Sat, 21 Aug 2021 22:41:38 +0200
Subject: [PATCH] Further split guards

---
 src/rb/auth.rs    |  4 ++--
 src/rb/errors.rs  |  1 +
 src/rbs/auth.rs   | 20 ++++++++------------
 src/rbs/guards.rs | 48 +++++++++++++++++++++++++++++++----------------
 4 files changed, 43 insertions(+), 30 deletions(-)

diff --git a/src/rb/auth.rs b/src/rb/auth.rs
index 4951e1c..18d3753 100644
--- a/src/rb/auth.rs
+++ b/src/rb/auth.rs
@@ -47,8 +47,8 @@ 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(|_| RBError::JWTCreationError)?;
+    let key: Hmac<Sha256> =
+        Hmac::new_from_slice(secret.as_bytes()).map_err(|_| RBError::JWTCreationError)?;
 
     let current_time = Utc::now();
 
diff --git a/src/rb/errors.rs b/src/rb/errors.rs
index 34bfe8f..9d048f3 100644
--- a/src/rb/errors.rs
+++ b/src/rb/errors.rs
@@ -20,6 +20,7 @@ pub enum RBError {
     MissingJWTKey,
     PWSaltError,
     AdminCreationError,
+    TokenExpired,
 }
 
 impl<'r> Responder<'r, 'static> for RBError {
diff --git a/src/rbs/auth.rs b/src/rbs/auth.rs
index 65d5255..2ff7561 100644
--- a/src/rbs/auth.rs
+++ b/src/rbs/auth.rs
@@ -1,23 +1,25 @@
+use crate::guards::User;
 use crate::RbDbConn;
 use rb::auth::{generate_jwt_token, verify_user, JWTResponse};
 use rocket::serde::json::Json;
 use serde::Deserialize;
-use crate::guards::User;
 
 pub(crate) fn routes() -> Vec<rocket::Route> {
-    routes![login, me]
+    routes![login, already_logged_in, me]
 }
 
-
 #[derive(Deserialize)]
 struct Credentials {
     username: String,
     password: String,
 }
 
-// TODO add catch for when user immediately requests new JWT token (they could totally spam this)
+#[post("/login")]
+async fn already_logged_in(_user: User) -> String {
+    String::from("You're already logged in!")
+}
 
-#[post("/login", data = "<credentials>")]
+#[post("/login", data = "<credentials>", rank = 2)]
 async fn login(conn: RbDbConn, credentials: Json<Credentials>) -> rb::Result<Json<JWTResponse>> {
     let credentials = credentials.into_inner();
 
@@ -29,10 +31,4 @@ async fn login(conn: RbDbConn, credentials: Json<Credentials>) -> rb::Result<Jso
     Ok(Json(conn.run(move |c| generate_jwt_token(c, &user)).await?))
 }
 
-#[get("/me")]
-async fn me(claims: User) -> String {
-    String::from("You are logged in!")
-}
-
-// /refresh
-// /logout
+// #[post("/refresh", data=)]
diff --git a/src/rbs/guards.rs b/src/rbs/guards.rs
index 96891d9..25d74a5 100644
--- a/src/rbs/guards.rs
+++ b/src/rbs/guards.rs
@@ -1,17 +1,18 @@
-use rocket::{
-    http::Status,
-    outcome::try_outcome,
-    request::{FromRequest, Outcome, Request}
-};
 use hmac::{Hmac, NewMac};
 use jwt::VerifyWithKey;
 use rb::auth::Claims;
+use rocket::{
+    http::Status,
+    outcome::try_outcome,
+    request::{FromRequest, Outcome, Request},
+};
 use sha2::Sha256;
 
-pub struct JWT(String);
+/// Extracts a "Authorization: Bearer" string from the headers.
+pub struct Bearer(String);
 
 #[rocket::async_trait]
-impl<'r> FromRequest<'r> for JWT {
+impl<'r> FromRequest<'r> for Bearer {
     type Error = rb::errors::RBError;
 
     async fn from_request(req: &'r Request<'_>) -> Outcome<Self, Self::Error> {
@@ -26,23 +27,24 @@ impl<'r> FromRequest<'r> for JWT {
         }
 
         // Extract the jwt token from the header
-        let jwt_token = match header.get(7..) {
-            Some(token) => token,
+        let auth_string = match header.get(7..) {
+            Some(s) => s,
             None => return Outcome::Forward(()),
         };
 
-        Outcome::Success(Self(jwt_token.to_string()))
+        Outcome::Success(Self(auth_string.to_string()))
     }
 }
 
-pub struct User(Claims);
+/// Verifies the provided JWT is valid.
+pub struct JWT(Claims);
 
 #[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> {
-        let jwt_token = try_outcome!(req.guard::<JWT>().await).0;
+        let bearer = try_outcome!(req.guard::<Bearer>().await).0;
 
         // Get secret & key
         let secret = match std::env::var("JWT_KEY") {
@@ -57,22 +59,36 @@ impl<'r> FromRequest<'r> for User {
                 return Outcome::Failure((Status::InternalServerError, Self::Error::JWTError))
             }
         };
-
         // Verify token using key
-        let claims: Claims = match jwt_token.verify_with_key(&key) {
+        let claims: Claims = match bearer.verify_with_key(&key) {
             Ok(claims) => claims,
             Err(_) => return Outcome::Failure((Status::Unauthorized, Self::Error::Unauthorized)),
         };
 
+        Outcome::Success(Self(claims))
+    }
+}
+
+/// Verifies the JWT has not expired.
+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 claims = try_outcome!(req.guard::<JWT>().await).0;
+
         // Verify key hasn't yet expired
         if chrono::Utc::now().timestamp() > claims.exp {
-            return Outcome::Failure((Status::Unauthorized, Self::Error::Unauthorized));
+            return Outcome::Failure((Status::Forbidden, Self::Error::TokenExpired));
         }
 
         Outcome::Success(Self(claims))
     }
 }
 
+/// Verifies the JWT belongs to an admin.
 pub struct Admin(Claims);
 
 #[rocket::async_trait]