From 2cc4d53961455ccec5145b62e42c0b013f3aa94f Mon Sep 17 00:00:00 2001
From: Jef Roosens <roosensjef@gmail.com>
Date: Mon, 30 Aug 2021 15:28:01 +0200
Subject: [PATCH] Moved JWT config to config file

---
 Rb.yaml         |  6 +++++-
 src/auth/jwt.rs | 24 ++++++++++++++++--------
 src/auth/mod.rs | 20 +++++++++++++++-----
 src/main.rs     | 31 ++++++++++++++++++++++---------
 4 files changed, 58 insertions(+), 23 deletions(-)

diff --git a/Rb.yaml b/Rb.yaml
index 3081e7b..6b40537 100644
--- a/Rb.yaml
+++ b/Rb.yaml
@@ -12,7 +12,11 @@ debug:
 
   admin_user: "admin"
   admin_pass: "password"
-  jwt_key: "secret"
+  jwt:
+    key: "secret"
+    refresh_token_size: 64
+    # Just 5 seconds for debugging
+    refresh_token_expire: 5
 
   databases:
     postgres_rb:
diff --git a/src/auth/jwt.rs b/src/auth/jwt.rs
index 83e7a1e..89ce734 100644
--- a/src/auth/jwt.rs
+++ b/src/auth/jwt.rs
@@ -11,6 +11,7 @@ use crate::{
     db::{tokens::NewRefreshToken, users::User},
     errors::{RbError, RbResult},
     schema::{refresh_tokens::dsl as refresh_tokens, users::dsl as users},
+    RbJwtConf,
 };
 
 #[derive(Serialize)]
@@ -30,10 +31,13 @@ pub struct Claims
     pub exp: i64,
 }
 
-pub fn generate_jwt_token(conn: &PgConnection, user: &User) -> RbResult<JWTResponse>
+pub fn generate_jwt_token(
+    conn: &PgConnection,
+    jwt: &RbJwtConf,
+    user: &User,
+) -> RbResult<JWTResponse>
 {
-    let secret = std::env::var("JWT_KEY").map_err(|_| RbError::Custom("Missing JWT key."))?;
-    let key: Hmac<Sha256> = Hmac::new_from_slice(secret.as_bytes())
+    let key: Hmac<Sha256> = Hmac::new_from_slice(jwt.key.as_bytes())
         .map_err(|_| RbError::Custom("Couldn't create Hmac key."))?;
 
     let current_time = Utc::now();
@@ -43,7 +47,7 @@ pub fn generate_jwt_token(conn: &PgConnection, user: &User) -> RbResult<JWTRespo
         id: user.id,
         username: user.username.clone(),
         admin: user.admin,
-        exp: current_time.timestamp() + crate::JWT_EXP_SECONDS,
+        exp: current_time.timestamp() + jwt.refresh_token_expire,
     };
 
     // Sign the claims into a new token
@@ -52,11 +56,11 @@ pub fn generate_jwt_token(conn: &PgConnection, user: &User) -> RbResult<JWTRespo
         .map_err(|_| RbError::Custom("Couldn't sign JWT."))?;
 
     // Generate a random refresh token
-    let mut refresh_token = [0u8; crate::REFRESH_TOKEN_N_BYTES];
+    let mut refresh_token = vec![0u8; jwt.refresh_token_size];
     thread_rng().fill(&mut refresh_token[..]);
 
     let refresh_expire =
-        (current_time + chrono::Duration::seconds(crate::REFRESH_TOKEN_EXP_SECONDS)).naive_utc();
+        (current_time + chrono::Duration::seconds(jwt.refresh_token_expire)).naive_utc();
 
     // Store refresh token in database
     db::tokens::create(
@@ -74,7 +78,11 @@ pub fn generate_jwt_token(conn: &PgConnection, user: &User) -> RbResult<JWTRespo
     })
 }
 
-pub fn refresh_token(conn: &PgConnection, refresh_token: &str) -> RbResult<JWTResponse>
+pub fn refresh_token(
+    conn: &PgConnection,
+    jwt: &RbJwtConf,
+    refresh_token: &str,
+) -> RbResult<JWTResponse>
 {
     let token_bytes =
         base64::decode(refresh_token).map_err(|_| RbError::AuthInvalidRefreshToken)?;
@@ -108,5 +116,5 @@ pub fn refresh_token(conn: &PgConnection, refresh_token: &str) -> RbResult<JWTRe
         .execute(conn)
         .map_err(|_| RbError::Custom("Couldn't update last used time."))?;
 
-    generate_jwt_token(conn, &user)
+    generate_jwt_token(conn, jwt, &user)
 }
diff --git a/src/auth/mod.rs b/src/auth/mod.rs
index b4cbf0f..78a4fa0 100644
--- a/src/auth/mod.rs
+++ b/src/auth/mod.rs
@@ -1,11 +1,11 @@
-use rocket::serde::json::Json;
+use rocket::{serde::json::Json, State};
 use serde::Deserialize;
 
 use self::{
     jwt::{generate_jwt_token, JWTResponse},
     pass::verify_user,
 };
-use crate::{errors::RbResult, guards::User, RbDbConn};
+use crate::{errors::RbResult, guards::User, RbConfig, RbDbConn, RbJwtConf};
 
 pub mod jwt;
 pub mod pass;
@@ -24,16 +24,24 @@ pub async fn already_logged_in(_user: User) -> String
 }
 
 #[post("/login", data = "<credentials>", rank = 2)]
-pub async fn login(conn: RbDbConn, credentials: Json<Credentials>) -> RbResult<Json<JWTResponse>>
+pub async fn login(
+    conn: RbDbConn,
+    conf: &State<RbConfig>,
+    credentials: Json<Credentials>,
+) -> RbResult<Json<JWTResponse>>
 {
     let credentials = credentials.into_inner();
+    let jwt = conf.jwt.clone();
 
     // Get the user, if credentials are valid
     let user = conn
         .run(move |c| verify_user(c, &credentials.username, &credentials.password))
         .await?;
 
-    Ok(Json(conn.run(move |c| generate_jwt_token(c, &user)).await?))
+    Ok(Json(
+        conn.run(move |c| generate_jwt_token(c, &jwt, &user))
+            .await?,
+    ))
 }
 
 #[derive(Deserialize)]
@@ -46,13 +54,15 @@ pub struct RefreshTokenRequest
 #[post("/refresh", data = "<refresh_token_request>")]
 pub async fn refresh_token(
     conn: RbDbConn,
+    conf: &State<RbConfig>,
     refresh_token_request: Json<RefreshTokenRequest>,
 ) -> RbResult<Json<JWTResponse>>
 {
     let refresh_token = refresh_token_request.into_inner().refresh_token;
+    let jwt = conf.jwt.clone();
 
     Ok(Json(
-        conn.run(move |c| crate::auth::jwt::refresh_token(c, &refresh_token))
+        conn.run(move |c| crate::auth::jwt::refresh_token(c, &jwt, &refresh_token))
             .await?,
     ))
 }
diff --git a/src/main.rs b/src/main.rs
index 92d46fd..fa147e9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -8,9 +8,13 @@ extern crate diesel_migrations;
 #[macro_use]
 extern crate diesel;
 
+use figment::{
+    providers::{Env, Format, Yaml},
+    Figment,
+};
 use rocket::{fairing::AdHoc, Build, Rocket};
 use rocket_sync_db_pools::database;
-use figment::{Figment, providers::Env, providers::Yaml, providers::Format};
+use serde::{Deserialize, Serialize};
 
 mod admin;
 pub mod auth;
@@ -19,14 +23,6 @@ pub mod errors;
 pub mod guards;
 pub(crate) mod schema;
 
-// Any import defaults are defined here
-/// Expire time for the JWT tokens in seconds.
-const JWT_EXP_SECONDS: i64 = 600;
-/// 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;
-
 #[database("postgres_rb")]
 pub struct RbDbConn(diesel::PgConnection);
 
@@ -61,6 +57,22 @@ async fn create_admin_user(rocket: Rocket<Build>) -> Result<Rocket<Build>, Rocke
     Ok(rocket)
 }
 
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct RbJwtConf
+{
+    key: String,
+    refresh_token_size: usize,
+    refresh_token_expire: i64,
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+pub struct RbConfig
+{
+    admin_user: String,
+    admin_pass: String,
+    jwt: RbJwtConf,
+}
+
 #[launch]
 fn rocket() -> _
 {
@@ -75,6 +87,7 @@ fn rocket() -> _
             run_db_migrations,
         ))
         .attach(AdHoc::try_on_ignite("Create admin user", create_admin_user))
+        .attach(AdHoc::config::<RbConfig>())
         .mount(
             "/api/auth",
             routes![auth::already_logged_in, auth::login, auth::refresh_token,],