From f50008ff9994dfbe02216679fe2ae0e8cc63d4b8 Mon Sep 17 00:00:00 2001
From: Jef Roosens <roosensjef@gmail.com>
Date: Wed, 1 Sep 2021 16:18:48 +0200
Subject: [PATCH] Fixed guard still using env var for jwt key

---
 Rb.yaml        |  2 +-
 rustfmt.toml   |  2 +-
 src/guards.rs  | 31 ++++++++++++++----------
 tests/admin.py | 64 ++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 84 insertions(+), 15 deletions(-)
 create mode 100644 tests/admin.py

diff --git a/Rb.yaml b/Rb.yaml
index 6b40537..c944758 100644
--- a/Rb.yaml
+++ b/Rb.yaml
@@ -16,7 +16,7 @@ debug:
     key: "secret"
     refresh_token_size: 64
     # Just 5 seconds for debugging
-    refresh_token_expire: 5
+    refresh_token_expire: 60
 
   databases:
     postgres_rb:
diff --git a/rustfmt.toml b/rustfmt.toml
index 03acab0..5e52857 100644
--- a/rustfmt.toml
+++ b/rustfmt.toml
@@ -49,7 +49,7 @@ reorder_imports = true
 reorder_modules = true
 report_fixme = "Always"
 report_todo = "Always"
-required_version = "1.4.36"
+required_version = "1.4.37"
 skip_children = false
 space_after_colon = true
 space_before_colon = false
diff --git a/src/guards.rs b/src/guards.rs
index 532ba97..d515d77 100644
--- a/src/guards.rs
+++ b/src/guards.rs
@@ -1,13 +1,16 @@
+use std::convert::From;
+
 use hmac::{Hmac, NewMac};
 use jwt::VerifyWithKey;
 use rocket::{
     http::Status,
     outcome::try_outcome,
     request::{FromRequest, Outcome, Request},
+    State,
 };
 use sha2::Sha256;
 
-use crate::auth::jwt::Claims;
+use crate::{auth::jwt::Claims, errors::RbError, RbConfig};
 
 /// Extracts a "Authorization: Bearer" string from the headers.
 pub struct Bearer<'a>(&'a str);
@@ -42,26 +45,28 @@ impl<'r> FromRequest<'r> for Bearer<'r>
 /// Verifies the provided JWT is valid.
 pub struct Jwt(Claims);
 
+impl From<()> for RbError
+{
+    fn from(_: ()) -> Self
+    {
+        RbError::Custom("Couldn't get config guard.")
+    }
+}
+
 #[rocket::async_trait]
 impl<'r> FromRequest<'r> for Jwt
 {
-    type Error = crate::errors::RbError;
+    type Error = RbError;
 
     async fn from_request(req: &'r Request<'_>) -> Outcome<Self, Self::Error>
     {
         let bearer = try_outcome!(req.guard::<Bearer>().await).0;
+        let config = try_outcome!(req.guard::<&State<RbConfig>>().await.map_failure(|_| (
+            Status::InternalServerError,
+            RbError::Custom("Couldn't get config guard.")
+        )));
 
-        // Get secret & key
-        let secret = match std::env::var("JWT_KEY") {
-            Ok(key) => key,
-            Err(_) => {
-                return Outcome::Failure((
-                    Status::InternalServerError,
-                    Self::Error::AuthUnauthorized,
-                ))
-            }
-        };
-        let key: Hmac<Sha256> = match Hmac::new_from_slice(secret.as_bytes()) {
+        let key: Hmac<Sha256> = match Hmac::new_from_slice(&config.jwt.key.as_bytes()) {
             Ok(key) => key,
             Err(_) => {
                 return Outcome::Failure((
diff --git a/tests/admin.py b/tests/admin.py
new file mode 100644
index 0000000..19a1c5b
--- /dev/null
+++ b/tests/admin.py
@@ -0,0 +1,64 @@
+import requests
+
+
+class RbClient:
+    def __init__(self, username, password, base_url = "http://localhost:8000/api"):
+        self.username = username
+        self.password = password
+        self.base_url = base_url
+
+        self.jwt = None
+        self.refresh_token = None
+
+    def _login(self):
+        r = requests.post(f"{self.base_url}/auth/login", json={
+            "username": self.username,
+            "password": self.password,
+        })
+
+        if r.status_code != 200:
+            raise Exception("Couldn't login")
+
+        res = r.json()
+        self.jwt = res["token"]
+        self.refresh_token = res["refreshToken"]
+
+    def _refresh(self):
+        r = requests.post(f"{self.base_url}/auth/refresh", json={"refreshToken": self.refresh_token})
+
+        if r.status_code != 200:
+            raise Exception("Couldn't refresh")
+
+        res = r.json()
+        self.jwt = res["token"]
+        self.refresh_token = res["refreshToken"]
+
+    def _request(self, type_, url, retry=2, *args, **kwargs):
+        if self.jwt:
+            headers = kwargs.get("headers", {})
+            headers["Authorization"] = f"Bearer {self.jwt}"
+            kwargs["headers"] = headers
+            print(kwargs["headers"])
+
+        r = requests.request(type_, url, *args, **kwargs)
+
+        if r.status_code != 200 and retry > 0:
+            if self.refresh_token:
+                self._refresh()
+
+            else:
+                self._login()
+
+            r = self._request(type_, url, *args, **kwargs, retry=retry - 1)
+
+        return r
+
+    def get(self, url, *args, **kwargs):
+        return self._request("GET", f"{self.base_url}{url}", *args, **kwargs)
+
+
+
+if __name__ == "__main__":
+    client = RbClient("admin", "password")
+
+    print(client.get("/admin/users").json())