refactor: added tests to key_manager

mickvandijke · mickvandijke · commit 79bb1ed7bec0 · 2022-01-14T12:27:22.000+01:00
diff --git a/src/http_server.rs b/src/http_server.rs
@@ -173,7 +173,7 @@ impl HttpServer {
 
                         // query.info_hash somehow receives a corrupt string
                         // so we have to get the info_hash manually from the raw query
-                        let info_hashes = HttpServer::get_info_hashes_from_raw_query(&raw_query);
+                        let info_hashes = HttpServer::info_hashes_from_raw_query(&raw_query);
                         if info_hashes.len() < 1 { return HttpServer::send_error("info_hash not found") }
                         query.info_hash = info_hashes[0].to_string();
                         debug!("{:?}", query.info_hash);
@@ -198,7 +198,7 @@ impl HttpServer {
                 })
                 .and_then(move |(key, raw_query, http_server): (Option<String>, String, Arc<HttpServer>)| {
                     async move {
-                        let info_hashes = HttpServer::get_info_hashes_from_raw_query(&raw_query);
+                        let info_hashes = HttpServer::info_hashes_from_raw_query(&raw_query);
                         if info_hashes.len() < 1 { return HttpServer::send_error("info_hash not found") }
                         if info_hashes.len() > 50 { return HttpServer::send_error("exceeded the max of 50 info_hashes") }
                         debug!("{:?}", info_hashes);
@@ -214,7 +214,7 @@ impl HttpServer {
         warp::any().and(announce_route.or(scrape_route))
     }
 
-    fn get_info_hashes_from_raw_query(raw_query: &str) -> Vec<InfoHash> {
+    fn info_hashes_from_raw_query(raw_query: &str) -> Vec<InfoHash> {
         let split_raw_query: Vec<&str> = raw_query.split("&").collect();
         let mut info_hashes: Vec<InfoHash> = Vec::new();
 
diff --git a/src/key_manager.rs b/src/key_manager.rs
@@ -6,6 +6,29 @@ use serde::Serialize;
 use log::debug;
 use derive_more::{Display, Error};
 
+pub fn generate_auth_key(seconds_valid: u64) -> AuthKey {
+    let key: String = thread_rng()
+        .sample_iter(&Alphanumeric)
+        .take(AUTH_KEY_LENGTH)
+        .map(char::from)
+        .collect();
+
+    debug!("Generated key: {}, valid for: {} seconds", key, seconds_valid);
+
+    AuthKey {
+        key,
+        valid_until: Some(current_time() + seconds_valid),
+    }
+}
+
+pub fn verify_auth_key(auth_key: &AuthKey) -> Result<(), Error> {
+    let current_time = current_time();
+    if auth_key.valid_until.is_none() { return Err(Error::KeyInvalid) }
+    if auth_key.valid_until.unwrap() < current_time { return Err(Error::KeyExpired) }
+
+    Ok(())
+}
+
 #[derive(Serialize, Debug, Eq, PartialEq, Clone)]
 pub struct AuthKey {
     pub key: String,
@@ -14,27 +37,35 @@ pub struct AuthKey {
 
 impl AuthKey {
     pub fn from_buffer(key_buffer: [u8; AUTH_KEY_LENGTH]) -> Option<AuthKey> {
-        Some(AuthKey {
-            key: String::from_utf8(Vec::from(key_buffer)).unwrap(),
-            valid_until: None,
-        })
+        if let Ok(key) = String::from_utf8(Vec::from(key_buffer)) {
+            Some(AuthKey {
+                key,
+                valid_until: None,
+            })
+        } else {
+            None
+        }
     }
 
     pub fn from_string(key: &str) -> Option<AuthKey> {
-        if key.len() != AUTH_KEY_LENGTH { return None }
-
-        Some(AuthKey {
-            key: key.to_string(),
-            valid_until: None,
-        })
+        if key.len() != AUTH_KEY_LENGTH {
+            None
+        } else {
+            Some(AuthKey {
+                key: key.to_string(),
+                valid_until: None,
+            })
+        }
     }
 }
 
 #[derive(Debug, Display, PartialEq, Error)]
 #[allow(dead_code)]
 pub enum Error {
-    #[display(fmt = "Key is invalid.")]
+    #[display(fmt = "Key could not be verified.")]
     KeyVerificationError,
+    #[display(fmt = "Key is invalid.")]
+    KeyInvalid,
     #[display(fmt = "Key has expired.")]
     KeyExpired
 }
@@ -46,29 +77,49 @@ impl From<r2d2_sqlite::rusqlite::Error> for Error {
     }
 }
 
-pub struct KeyManager;
+#[cfg(test)]
+mod tests {
+    use crate::key_manager;
 
-impl KeyManager {
-    pub fn generate_auth_key(&self, seconds_valid: u64) -> AuthKey {
-        let key: String = thread_rng()
-            .sample_iter(&Alphanumeric)
-            .take(AUTH_KEY_LENGTH)
-            .map(char::from)
-            .collect();
+    #[test]
+    fn auth_key_from_buffer() {
+        let auth_key = key_manager::AuthKey::from_buffer(
+            [
+                89, 90, 83, 108,
+                52, 108, 77, 90,
+                117, 112, 82, 117,
+                79, 112, 83, 82,
+                67, 51, 107, 114,
+                73, 75, 82, 53,
+                66, 80, 66, 49,
+                52, 110, 114, 74]
+        );
 
-        debug!("Generated key: {}, valid for: {} seconds", key, seconds_valid);
+        assert!(auth_key.is_some());
+        assert_eq!(auth_key.unwrap().key, "YZSl4lMZupRuOpSRC3krIKR5BPB14nrJ");
+    }
 
-        AuthKey {
-            key,
-            valid_until: Some(current_time() + seconds_valid),
-        }
+    #[test]
+    fn auth_key_from_string() {
+        let key_string = "YZSl4lMZupRuOpSRC3krIKR5BPB14nrJ";
+        let auth_key = key_manager::AuthKey::from_string(key_string);
+
+        assert!(auth_key.is_some());
+        assert_eq!(auth_key.unwrap().key, key_string);
+    }
+
+    #[test]
+    fn generate_valid_auth_key() {
+        let auth_key = key_manager::generate_auth_key(9999);
+
+        assert!(key_manager::verify_auth_key(&auth_key).is_ok());
     }
 
-    pub async fn verify_auth_key(&self, auth_key: &AuthKey) -> Result<(), Error> {
-        let current_time = current_time();
-        if auth_key.valid_until.is_none() { return Err(Error::KeyVerificationError) }
-        if &auth_key.valid_until.unwrap() < &current_time { return Err(Error::KeyExpired) }
+    #[test]
+    fn generate_expired_auth_key() {
+        let mut auth_key = key_manager::generate_auth_key(0);
+        auth_key.valid_until = Some(0);
 
-        Ok(())
+        assert!(key_manager::verify_auth_key(&auth_key).is_err());
     }
 }
diff --git a/src/request.rs b/src/request.rs
@@ -4,6 +4,7 @@ use std::io;
 use std::io::{Cursor, Read};
 use byteorder::{NetworkEndian, ReadBytesExt};
 use std::convert::TryInto;
+use log::debug;
 use crate::key_manager::AuthKey;
 
 #[derive(PartialEq, Eq, Clone, Debug)]
@@ -169,6 +170,7 @@ impl Request {
                     // key should be the last bytes
                     cursor.set_position((bytes.len() - AUTH_KEY_LENGTH) as u64);
                     if cursor.read_exact(&mut key_buffer).is_ok() {
+                        debug!("AuthKey buffer: {:?}", key_buffer);
                         AuthKey::from_buffer(key_buffer)
                     } else {
                         None
diff --git a/src/tracker.rs b/src/tracker.rs
@@ -10,7 +10,7 @@ use std::collections::btree_map::Entry;
 use crate::database::SqliteDatabase;
 use std::sync::Arc;
 use log::debug;
-use crate::key_manager::{AuthKey, KeyManager};
+use crate::key_manager::{AuthKey};
 use r2d2_sqlite::rusqlite;
 use crate::http_server::HttpAnnounceRequest;
 
@@ -250,7 +250,6 @@ pub struct TorrentTracker {
     pub config: Arc<Configuration>,
     torrents: tokio::sync::RwLock<std::collections::BTreeMap<InfoHash, TorrentEntry>>,
     database: SqliteDatabase,
-    key_manager: KeyManager,
 }
 
 impl TorrentTracker {
@@ -263,12 +262,11 @@ impl TorrentTracker {
             config,
             torrents: RwLock::new(std::collections::BTreeMap::new()),
             database,
-            key_manager: KeyManager {}
         }
     }
 
     pub async fn generate_auth_key(&self, seconds_valid: u64) -> Result<AuthKey, rusqlite::Error> {
-        let auth_key = self.key_manager.generate_auth_key(seconds_valid);
+        let auth_key = key_manager::generate_auth_key(seconds_valid);
 
         // add key to database
         if let Err(error) = self.database.add_key_to_keys(&auth_key).await { return Err(error) }
@@ -282,7 +280,7 @@ impl TorrentTracker {
 
     pub async fn verify_auth_key(&self, auth_key: &AuthKey) -> Result<(), key_manager::Error> {
         let db_key = self.database.get_key_from_keys(&auth_key.key).await?;
-        self.key_manager.verify_auth_key(&db_key).await
+        key_manager::verify_auth_key(&db_key)
     }
 
     pub async fn authenticate_request(&self, info_hash: &InfoHash, key: &Option<AuthKey>) -> Result<(), TorrentError> {
@@ -298,7 +296,7 @@ impl TorrentTracker {
             TrackerMode::PrivateMode => {
                 match key {
                     Some(key) => {
-                        if self.key_manager.verify_auth_key(key).await.is_err() {
+                        if key_manager::verify_auth_key(key).is_err() {
                             return Err(TorrentError::PeerKeyNotValid)
                         }
 
@@ -312,7 +310,7 @@ impl TorrentTracker {
             TrackerMode::PrivateListedMode => {
                 match key {
                     Some(key) => {
-                        if self.key_manager.verify_auth_key(key).await.is_err() {
+                        if key_manager::verify_auth_key(key).is_err() {
                             return Err(TorrentError::PeerKeyNotValid)
                         }
 
