From fea3fb235dbedc3895ce6a36e8f821535a8db63b Mon Sep 17 00:00:00 2001
From: Peter Cai <peter@typeblog.net>
Date: Fri, 21 Feb 2020 09:57:03 +0800
Subject: [PATCH] user: implement authorization guard

---
 src/api.rs   |  9 ++++++++-
 src/tests.rs | 39 ++++++++++++++++++++++++++++++++++++++-
 src/user.rs  | 44 +++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 87 insertions(+), 5 deletions(-)

diff --git a/src/api.rs b/src/api.rs
index 6e96df1..20d1b32 100644
--- a/src/api.rs
+++ b/src/api.rs
@@ -11,7 +11,8 @@ pub fn routes() -> impl Into<Vec<rocket::Route>> {
         auth,
         auth_change_pw,
         auth_sign_in,
-        auth_params
+        auth_params,
+        auth_ping
     ]
 }
 
@@ -119,4 +120,10 @@ fn auth_change_pw(db: DbConn, params: Json<ChangePwParams>) -> Custom<JsonResp<(
         Err(user::UserOpError(e)) =>
             error_resp(Status::InternalServerError, vec![e])
     }
+}
+
+// For testing the User request guard
+#[get("/auth/ping")]
+fn auth_ping(_db: DbConn, u: user::User) -> Custom<JsonResp<String>> {
+    Custom(Status::Ok, Json(Response::Success(u.email)))
 }
\ No newline at end of file
diff --git a/src/tests.rs b/src/tests.rs
index 61633fb..849e6be 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -1,6 +1,6 @@
 use crate::build_rocket;
 use rocket::local::Client;
-use rocket::http::{ContentType, Status};
+use rocket::http::{Header, ContentType, Status};
 use lazy_static::*;
 
 fn get_test_client() -> Client {
@@ -196,4 +196,41 @@ fn should_change_pw_successfully_and_log_in_successfully() {
         }"#)
         .dispatch();
     assert_eq!(resp.status(), Status::Ok);
+}
+
+#[test]
+fn should_fail_authorize() {
+    let resp = CLIENT.get("/auth/ping").dispatch();
+    assert_eq!(resp.status(), Status::Unauthorized);
+}
+
+#[test]
+fn should_fail_authorize_2() {
+    let resp = CLIENT.get("/auth/ping")
+        .header(Header::new("Authorization", "iwoe0nvie0bv024ibv043bv"))
+        .dispatch();
+    assert_eq!(resp.status(), Status::Unauthorized);
+}
+
+#[test]
+fn should_success_authorize() {
+    let token = CLIENT.post("/auth")
+        .header(ContentType::JSON)
+        .body(r#"{
+            "email": "test7@example.com",
+            "password": "testpw",
+            "pw_cost": "100",
+            "pw_nonce": "whatever",
+            "version": "001"
+        }"#)
+        .dispatch()
+        .body_string()
+        .unwrap()
+        .replace("{\"token\":\"", "")
+        .replace("\"}", "");
+    let mut resp = CLIENT.get("/auth/ping")
+        .header(Header::new("Authorization", token))
+        .dispatch();
+    assert_eq!(resp.status(), Status::Ok);
+    assert_eq!(resp.body_string().unwrap(), "\"test7@example.com\"");
 }
\ No newline at end of file
diff --git a/src/user.rs b/src/user.rs
index eb21354..d1df3e5 100644
--- a/src/user.rs
+++ b/src/user.rs
@@ -3,10 +3,17 @@ use crate::schema::users::dsl::*;
 use crate::{lock_db_write, lock_db_read};
 use diesel::prelude::*;
 use diesel::sqlite::SqliteConnection;
+use rocket::request;
+use rocket::http::Status;
 use serde::Deserialize;
 use std::env;
 use std::time::{SystemTime, UNIX_EPOCH};
 
+lazy_static! {
+    static ref JWT_SECRET: String = env::var("SFRS_JWT_SECRET")
+        .expect("Please have SFRS_JWT_SECRET set");
+}
+
 #[derive(Debug)]
 pub struct UserOpError(pub String);
 
@@ -133,6 +140,18 @@ impl User {
         }
     }
 
+    pub fn find_user_by_token(db: &SqliteConnection, token: &str) -> Result<User, UserOpError> {
+        let parsed: jwt::Token<jwt::Claims, jwt::Registered> =
+            jwt::Token::parse(token).map_err(|_| "Invalid JWT Token".into())?;
+        if !parsed.verify(JWT_SECRET.as_bytes(), crypto::sha2::Sha256::new()) {
+            Err("JWT Signature Verification Error".into())
+        } else {
+            parsed.claims.sub
+                .ok_or("Malformed Token".into())
+                .and_then(|mail| Self::find_user_by_email(db, &mail))
+        }
+    }
+
     // Create a JWT token for the current user if password matches
     pub fn create_token(&self, passwd: &str) -> Result<String, UserOpError> {
         if self.password != passwd {
@@ -150,9 +169,7 @@ impl User {
                     nbf: None,
                     jti: None
                 })
-            ).signed(env::var("SFRS_JWT_SECRET")
-                .expect("Please have SFRS_JWT_SECRET set")
-                .as_bytes(), crypto::sha2::Sha256::new())
+            ).signed(JWT_SECRET.as_bytes(), crypto::sha2::Sha256::new())
              .map_err(|_| UserOpError::new("Failed to generate token"))
         }
     }
@@ -174,4 +191,25 @@ impl User {
                     .map_err(|_| UserOpError::new("Database error")))
         }
     }
+}
+
+// Implement request guard for User type
+// This is intended for protecting authorized endpoints
+impl<'a, 'r> request::FromRequest<'a, 'r> for User {
+    type Error = UserOpError;
+
+    fn from_request(request: &'a request::Request<'r>) -> request::Outcome<Self, Self::Error> {
+        let token = request.headers().get_one("authorization");
+        match token {
+            None => request::Outcome::Failure((Status::Unauthorized, "Token missing".into())),
+            Some(token) => {
+                let result = Self::find_user_by_token(
+                    &request.guard::<crate::DbConn>().unwrap(), token);
+                match result {
+                    Ok(u) => request::Outcome::Success(u),
+                    Err(err) => request::Outcome::Failure((Status::Unauthorized, err))
+                }
+            }
+        }
+    }
 }
\ No newline at end of file