From 4c59b345a5aaed520586c9cd4313c0703e807c9f Mon Sep 17 00:00:00 2001
From: trivernis <trivernis@protonmail.com>
Date: Sun, 13 Sep 2020 14:21:15 +0200
Subject: [PATCH] Fix panic on non-base64 session tokens

Signed-off-by: trivernis <trivernis@protonmail.com>
---
 src/database/tokens.rs | 11 ++++++-----
 src/server/user_rpc.rs |  3 ++-
 src/utils/mod.rs       | 10 +++++++---
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/database/tokens.rs b/src/database/tokens.rs
index 6a56a92..137ea0f 100644
--- a/src/database/tokens.rs
+++ b/src/database/tokens.rs
@@ -59,7 +59,7 @@ impl SessionTokens {
 
     /// Returns the user id that is stored in the first four bytes of the refresh token
     pub fn get_user_id(&self) -> i32 {
-        get_user_id_from_token(&self.refresh_token)
+        get_user_id_from_token(&self.refresh_token).unwrap()
     }
 
     /// Saves the tokens into the database
@@ -178,7 +178,7 @@ impl TokenStore {
 
     /// Returns the token store entry for a given request token
     pub fn get_by_request_token(&self, request_token: &String) -> Option<&TokenStoreEntry> {
-        let user_id = get_user_id_from_token(&request_token);
+        let user_id = get_user_id_from_token(&request_token)?;
         if let Some(user_tokens) = self.tokens.get(&user_id) {
             user_tokens.iter().find(|e| {
                 if let Some(token) = e.request_token() {
@@ -194,7 +194,7 @@ impl TokenStore {
 
     /// Returns the token store entry by the given refresh token
     pub fn get_by_refresh_token(&self, refresh_token: &String) -> Option<&TokenStoreEntry> {
-        let user_id = get_user_id_from_token(&refresh_token);
+        let user_id = get_user_id_from_token(&refresh_token)?;
         if let Some(user_tokens) = self.tokens.get(&user_id) {
             user_tokens.iter().find(|e| {
                 if let Some(token) = e.refresh_token() {
@@ -212,7 +212,7 @@ impl TokenStore {
     /// Also clears all expired token entries.
     pub fn set_request_token(&mut self, refresh_token: &String, request_token: &String) {
         self.clear_expired();
-        let user_id = get_user_id_from_token(&request_token);
+        let user_id = get_user_id_from_token(&request_token).unwrap();
         if let Some(user_tokens) = self.tokens.get_mut(&user_id) {
             user_tokens.iter_mut().for_each(|e| {
                 if let Some(ref_token) = &e.refresh_token() {
@@ -226,7 +226,8 @@ impl TokenStore {
 
     /// Inserts a new pair of request and refresh token
     pub fn insert(&mut self, request_token: &String, refresh_token: &String) -> Result<(), String> {
-        let user_id = get_user_id_from_token(refresh_token);
+        let user_id =
+            get_user_id_from_token(refresh_token).ok_or("Invalid request token".to_string())?;
         let user_tokens = if let Some(user_tokens) = self.tokens.get_mut(&user_id) {
             user_tokens
         } else {
diff --git a/src/server/user_rpc.rs b/src/server/user_rpc.rs
index 3e9aa33..08c5d07 100644
--- a/src/server/user_rpc.rs
+++ b/src/server/user_rpc.rs
@@ -161,7 +161,8 @@ impl UserRpcServer {
         {
             return Err(ErrorMessage::new("Invalid request token".to_string()));
         }
-        let user_id = get_user_id_from_token(&message.token);
+        let user_id = get_user_id_from_token(&message.token)
+            .ok_or(ErrorMessage::new("Invalid request token".to_string()))?;
         let response_data = database.user_roles.by_user(user_id)?;
 
         Ok(Message::new_with_serialize(GET_ROLES, response_data))
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index 3354d07..c31a509 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -29,9 +29,13 @@ pub fn create_user_token(user_id: i32) -> [u8; TOKEN_LENGTH] {
 }
 
 /// Extracts the userId from a request token
-pub fn get_user_id_from_token(token: &String) -> i32 {
-    let token = base64::decode(&token).unwrap();
-    BigEndian::read_i32(token.as_slice())
+pub fn get_user_id_from_token(token: &String) -> Option<i32> {
+    let token = base64::decode(&token).ok()?;
+    if token.len() > 4 {
+        Some(BigEndian::read_i32(token.as_slice()))
+    } else {
+        None
+    }
 }
 
 /// Hashes a password with a salt by using BCrypt