Added Admin & User guards

src/rb/auth.rs

@@ -1,5 +1,5 @@
 use crate::errors::RBError;
-use crate::models::{NewRefreshToken, User, NewUser};
+use crate::models::{NewRefreshToken, NewUser, User};
 use crate::schema::refresh_tokens::dsl as refresh_tokens;
 use crate::schema::users::dsl as users;
 use argon2::verify_encoded;
@@ -9,7 +9,7 @@ use diesel::{insert_into, PgConnection};
 use hmac::{Hmac, NewMac};
 use jwt::SignWithKey;
 use rand::{thread_rng, Rng};
-use serde::Serialize;
+use serde::{Deserialize, Serialize};
 use sha2::Sha256;
 use std::collections::HashMap;
 
@@ -51,22 +51,36 @@ pub struct JWTResponse {
     refresh_token: String,
 }
 
+#[derive(Serialize, Deserialize)]
+pub struct Claims {
+    pub id: uuid::Uuid,
+    pub username: String,
+    pub admin: bool,
+    pub exp: i64,
+}
+
 pub fn generate_jwt_token(conn: &PgConnection, user: &User) -> crate::Result<JWTResponse> {
-    // TODO actually use proper secret here
+    let secret = std::env::var("JWT_KEY").map_err(|_| RBError::MissingJWTKey)?;
-    let key: Hmac<Sha256> =
+    let key: Hmac<Sha256> = Hmac::new_from_slice(secret.as_bytes())
-        Hmac::new_from_slice(b"some-secret").map_err(|_| log("Failed to create key", RBError::JWTCreationError))?;
+        .map_err(|_| log("Failed to create key", RBError::JWTCreationError))?;
 
     let current_time = Utc::now();
 
     // Create the claims
-    let mut claims = HashMap::new();
+    let claims = Claims {
-    claims.insert("id", user.id.to_string());
+        id: user.id,
-    claims.insert("username", user.username.clone());
+        username: user.username.clone(),
-    claims.insert("admin", user.admin.to_string());
+        admin: user.admin,
-    claims.insert(
+        exp: current_time.timestamp() + JWT_EXP_SECONDS,
-        "exp",
+    };
-        (current_time.timestamp() + JWT_EXP_SECONDS).to_string(),
+    // 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
@@ -77,7 +91,8 @@ pub fn generate_jwt_token(conn: &PgConnection, user: &User) -> crate::Result<JWT
     let mut refresh_token = [0u8; REFRESH_TOKEN_N_BYTES];
     thread_rng().fill(&mut refresh_token[..]);
 
-    let refresh_expire = (current_time + chrono::Duration::seconds(REFRESH_TOKEN_EXP_SECONDS)).naive_utc();
+    let refresh_expire =
+        (current_time + chrono::Duration::seconds(REFRESH_TOKEN_EXP_SECONDS)).naive_utc();
 
     // Store refresh token in database
     // TODO add expires_at here (it's what's causing the errors)
@@ -85,7 +100,7 @@ pub fn generate_jwt_token(conn: &PgConnection, user: &User) -> crate::Result<JWT
         .values(NewRefreshToken {
             token: refresh_token.to_vec(),
             user_id: user.id,
-            expires_at: refresh_expire
+            expires_at: refresh_expire,
         })
         .execute(conn)
         .map_err(|_| RBError::JWTCreationError)?;
@@ -106,20 +121,25 @@ pub fn hash_password(password: &str) -> crate::Result<String> {
     argon2::hash_encoded(password.as_bytes(), &salt, &config).map_err(|_| RBError::PWSaltError)
 }
 
-pub fn create_admin_user(conn: &PgConnection, username: &str, password: &str) -> crate::Result<bool> {
+pub fn create_admin_user(
+    conn: &PgConnection,
+    username: &str,
+    password: &str,
+) -> crate::Result<bool> {
     let pass_hashed = hash_password(password)?;
     let new_user = NewUser {
-            username: username.to_string(),
+        username: username.to_string(),
-            password: pass_hashed,
+        password: pass_hashed,
-            admin: true,
+        admin: true,
-        };
+    };
 
     insert_into(users::users)
         .values(&new_user)
         .on_conflict(users::username)
         .do_update()
         .set(&new_user)
-        .execute(conn).map_err(|_| RBError::AdminCreationError)?;
+        .execute(conn)
+        .map_err(|_| RBError::AdminCreationError)?;
 
     Ok(true)
 }

src/rb/errors.rs

@@ -16,6 +16,8 @@ pub enum RBError {
     JWTTokenExpired,
     /// Umbrella error for when something goes wrong whilst creating a JWT token pair
     JWTCreationError,
+    JWTError,
+    MissingJWTKey,
     PWSaltError,
     AdminCreationError,
 }
@@ -28,8 +30,10 @@ impl<'r> Responder<'r, 'static> for RBError {
             RBError::InvalidPassword => (Status::Unauthorized, "Invalid password"),
             RBError::Unauthorized => (Status::Unauthorized, "Unauthorized"),
             RBError::JWTTokenExpired => (Status::Unauthorized, "Token expired"),
-            RBError::JWTCreationError => (Status::InternalServerError, "Failed to create tokens."),
+            RBError::JWTCreationError | RBError::MissingJWTKey => {
-            _ => (Status::InternalServerError, "Internal server error")
+                (Status::InternalServerError, "Failed to create tokens.")
+            }
+            _ => (Status::InternalServerError, "Internal server error"),
         };
 
         let mut res = Response::new();

src/rb/models.rs

@@ -1,5 +1,5 @@
 use crate::schema::{refresh_tokens, users};
-use diesel::{Insertable, Queryable, AsChangeset};
+use diesel::{AsChangeset, Insertable, Queryable};
 use serde::Serialize;
 use uuid::Uuid;
 
@@ -27,5 +27,5 @@ pub struct NewUser {
 pub struct NewRefreshToken {
     pub token: Vec<u8>,
     pub user_id: Uuid,
-    pub expires_at: chrono::NaiveDateTime
+    pub expires_at: chrono::NaiveDateTime,
 }

src/rbs/auth.rs

@@ -2,11 +2,13 @@ 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]
+    routes![login, me]
 }
 
+
 #[derive(Deserialize)]
 struct Credentials {
     username: String,
@@ -27,5 +29,10 @@ 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

src/rbs/guards.rs

@@ -0,0 +1,77 @@
+use rocket::{
+    http::Status,
+    outcome::try_outcome,
+    request::{FromRequest, Outcome, Request}
+};
+use hmac::{Hmac, NewMac};
+use jwt::VerifyWithKey;
+use rb::auth::Claims;
+use sha2::Sha256;
+
+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> {
+        // If the header isn't present, just forward to the next route
+        let header = match req.headers().get_one("Authorization") {
+            None => return Outcome::Forward(()),
+            Some(val) => val,
+        };
+
+        if !header.starts_with("Bearer ") {
+            return Outcome::Forward(());
+        }
+
+        // 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)),
+        };
+
+        // Get secret & key
+        let secret = match std::env::var("JWT_KEY") {
+            Ok(key) => key,
+            Err(_) => {
+                return Outcome::Failure((Status::InternalServerError, Self::Error::MissingJWTKey))
+            }
+        };
+        let key: Hmac<Sha256> = match Hmac::new_from_slice(secret.as_bytes()) {
+            Ok(key) => key,
+            Err(_) => {
+                return Outcome::Failure((Status::InternalServerError, Self::Error::JWTError))
+            }
+        };
+
+        // Verify token using key
+        let claims: Claims = match jwt_token.verify_with_key(&key) {
+            Ok(claims) => claims,
+            Err(_) => return Outcome::Failure((Status::Unauthorized, Self::Error::Unauthorized)),
+        };
+
+        // Verify key hasn't yet expired
+        if chrono::Utc::now().timestamp() > claims.exp {
+            return Outcome::Failure((Status::Unauthorized, Self::Error::Unauthorized));
+        }
+
+        Outcome::Success(Self(claims))
+    }
+}
+
+pub struct Admin(Claims);
+
+#[rocket::async_trait]
+impl<'r> FromRequest<'r> for Admin {
+    type Error = rb::errors::RBError;
+
+    async fn from_request(req: &'r Request<'_>) -> Outcome<Self, Self::Error> {
+        let user = try_outcome!(req.guard::<User>().await);
+        if user.0.admin {
+            Outcome::Success(Self(user.0))
+        } else {
+            Outcome::Forward(())
+        }
+    }
+}

src/rbs/main.rs

@@ -11,6 +11,7 @@ use rocket::{fairing::AdHoc, Build, Rocket};
 use rocket_sync_db_pools::{database, diesel};
 
 mod auth;
+pub(crate) mod guards;
 
 embed_migrations!();
 
@@ -35,7 +36,11 @@ async fn create_admin_user(rocket: Rocket<Build>) -> Result<Rocket<Build>, Rocke
     let conn = RbDbConn::get_one(&rocket)
         .await
         .expect("database connection");
-    conn.run(move |c| rb::auth::create_admin_user(c, &admin_user, &admin_password).expect("failed to create admin user")).await;
+    conn.run(move |c| {
+        rb::auth::create_admin_user(c, &admin_user, &admin_password)
+            .expect("failed to create admin user")
+    })
+    .await;
 
     Ok(rocket)
 }
@@ -48,9 +53,6 @@ fn rocket() -> _ {
             "Run database migrations",
             run_db_migrations,
         ))
-        .attach(AdHoc::try_on_ignite(
+        .attach(AdHoc::try_on_ignite("Create admin user", create_admin_user))
-                "Create admin user",
-                create_admin_user
-        ))
         .mount("/api/auth", auth::routes())
 }