First JWT login implementation

Cargo.lock

@@ -1179,6 +1179,7 @@ checksum = "61b3909d758bb75c79f23d4736fac9433868679d3ad2ea7a61e3c25cfda9a088"
 name = "rusty-bever"
 version = "0.1.0"
 dependencies = [
+ "base64",
  "chrono",
  "diesel",
  "diesel_migrations",

Cargo.toml

@@ -28,6 +28,7 @@ jwt = "0.14.0"
 hmac = "*"
 sha2 = "*"
 chrono = "0.4.19"
+base64 = "0.13.0"
 
 # Backend web framework
 [dependencies.rocket]

src/rb/auth.rs

@@ -1,19 +1,22 @@
 use crate::errors::RBError;
-use crate::models::User;
+use crate::models::{User, NewRefreshToken};
 use crate::schema::users::dsl as users;
+use crate::schema::refresh_tokens::dsl as refresh_tokens;
 use argon2::verify_encoded;
 use diesel::prelude::*;
-use diesel::PgConnection;
+use diesel::{PgConnection, insert_into};
 use hmac::{Hmac, NewMac};
 use jwt::SignWithKey;
 use sha2::Sha256;
 use std::collections::HashMap;
 use chrono::Utc;
+use serde::Serialize;
+use rand::{thread_rng, Rng};
 
 /// 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: u32 = 64;
+const REFRESH_TOKEN_N_BYTES: usize = 64;
 
 pub fn verify_user(conn: &PgConnection, username: &str, password: &str) -> crate::Result<User> {
     // TODO handle non-"NotFound" Diesel errors accordingely
@@ -28,23 +31,39 @@ pub fn verify_user(conn: &PgConnection, username: &str, password: &str) -> crate
     }
 }
 
-struct JWTResponse {
+#[derive(Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct JWTResponse {
     token: String,
     refresh_token: String
 }
 
-pub fn generate_jwt_token(conn: &PgConnection, user: &User) -> JWTResponse {
+pub fn generate_jwt_token(conn: &PgConnection, user: &User) -> crate::Result<JWTResponse> {
     // TODO actually use proper secret here
-    // TODO don't just unwrap here
+    let key: Hmac<Sha256> = Hmac::new_from_slice(b"some-secret").map_err(|_| RBError::JWTCreationError)?;
-    let key: Hmac<Sha256> = Hmac::new_from_slice(b"some-secret").unwrap();
 
     // Create the claims
     let mut claims = HashMap::new();
     claims.insert("id", user.id.to_string());
-    claims.insert("username", user.username);
+    claims.insert("username", user.username.clone());
+    claims.insert("admin", user.admin.to_string());
     claims.insert("exp", (Utc::now().timestamp() + JWT_EXP_SECONDS).to_string());
 
     // Sign the claims into a new token
-    // TODO don't just unwrap here
+    let token = claims.sign_with_key(&key).map_err(|_| RBError::JWTCreationError)?;
-    let token = claims.sign_with_key(&key).unwrap();
+
+    // Generate a random refresh token
+    let mut refresh_token = [0u8; REFRESH_TOKEN_N_BYTES];
+    thread_rng().fill(&mut refresh_token[..]);
+
+    // Store refresh token in database
+    insert_into(refresh_tokens::refresh_tokens).values(NewRefreshToken {
+        token: refresh_token.to_vec(),
+        user_id: user.id
+    }).execute(conn).map_err(|_| RBError::JWTCreationError)?;
+
+    Ok(JWTResponse {
+        token: token,
+        refresh_token: base64::encode(refresh_token)
+    })
 }

src/rb/errors.rs

@@ -11,19 +11,22 @@ pub enum RBError {
     /// When a non-admin user tries to use an admin endpoint
     Unauthorized,
     /// When an expired JWT token is used for auth.
-    JWTTokenExpired
+    JWTTokenExpired,
+    /// Umbrella error for when something goes wrong whilst creating a JWT token pair
+    JWTCreationError
 }
 
 impl<'r> Responder<'r, 'static> for RBError {
     fn respond_to(self, _: &'r Request<'_>) -> response::Result<'static> {
         let (status, message): (Status, &str) = match self {
-            UnknownUser => (Status::NotFound, "Unknown user"),
+            RBError::UnknownUser => (Status::NotFound, "Unknown user"),
-            InvalidPassword => (Status::Unauthorized, "Invalid password"),
+            RBError::InvalidPassword => (Status::Unauthorized, "Invalid password"),
-            Unauthorized => (Status::Unauthorized, "Unauthorized"),
+            RBError::Unauthorized => (Status::Unauthorized, "Unauthorized"),
-            JWTTokenExpired => (Status::Unauthorized, "Token expired"),
+            RBError::JWTTokenExpired => (Status::Unauthorized, "Token expired"),
+            RBError::JWTCreationError => (Status::InternalServerError, "Failed to create tokens."),
         };
 
-        let res = Response::new();
+        let mut res = Response::new();
         res.set_status(status);
         res.set_sized_body(message.len(), io::Cursor::new(message));
 

src/rb/models.rs

@@ -1,6 +1,7 @@
-use diesel::Queryable;
+use diesel::{Queryable, Insertable};
 use uuid::Uuid;
 use serde::Serialize;
+use crate::schema::refresh_tokens;
 
 #[derive(Queryable, Serialize)]
 pub struct User {
@@ -10,5 +11,13 @@ pub struct User {
     pub password: String,
     #[serde(skip_serializing)]
     blocked: bool,
-    admin: bool,
+    pub admin: bool,
+}
+
+
+#[derive(Insertable)]
+#[table_name = "refresh_tokens"]
+pub struct NewRefreshToken {
+    pub token: Vec<u8>,
+    pub user_id: Uuid
 }

src/rbs/auth.rs

@@ -1,5 +1,5 @@
 use crate::RbDbConn;
-use rb::auth::verify_user;
+use rb::auth::{verify_user, JWTResponse, generate_jwt_token};
 use rocket::serde::json::Json;
 use serde::Deserialize;
 
@@ -9,14 +9,18 @@ struct Credentials {
     password: String,
 }
 
+// TODO add catch for when user immediately requests new JWT token (they could totally spam this)
+
 #[post("/login", data = "<credentials>")]
-async fn login(conn: RbDbConn, credentials: Json<Credentials>) {
+async fn login(conn: RbDbConn, credentials: Json<Credentials>) -> rb::Result<Json<JWTResponse>> {
     let credentials = credentials.into_inner();
 
+    // Get the user, if credentials are valid
     let user = conn
         .run(move |c| verify_user(c, &credentials.username, &credentials.password))
-        .await;
+        .await?;
-    user
+
+    Ok(Json(conn.run(move |c| generate_jwt_token(c, &user)).await?))
 }
 
 // /refresh