feat(server): add aes_key to request

noctisynth · Feb 13, 2024 · f1a95e4 · f1a95e4
1 parent 21fbe82
commit f1a95e4
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 14 deletions.
diff --git a/src/models/router.rs b/src/models/router.rs
@@ -48,10 +48,10 @@ impl RoutePath {
         }
     }
 
-    pub fn check(&mut self, olps: String) -> bool {
+    pub fn check(&mut self, olps: &str) -> bool {
         if self.route_type == RouteType::RegexPath {
             let regex = Regex::new(&self.route).unwrap();
-            regex.is_match(&olps)
+            regex.is_match(olps)
         } else if self.route_type == RouteType::StartswithPath {
             olps.starts_with(&self.route)
         } else {
@@ -86,10 +86,10 @@ impl Router {
         self.routes.insert(path.clone(), route);
     }
 
-    pub fn get_handler(&self, path: String) -> Route {
+    pub fn get_handler(&self, path: &str) -> Route {
         for (route_path, route) in &self.routes {
             let mut route_path = route_path.clone();
-            if route_path.check(path.clone()) {
+            if route_path.check(path) {
                 return route.clone();
             };
         }

diff --git a/src/models/server.rs b/src/models/server.rs
@@ -50,6 +50,8 @@ impl ServerConnection {
         let mut oke = OKE::new(Some(&self.private_key), Some(self.public_key))?;
         oke.to_stream_with_salt(stream).await?;
         oke.from_stream(stream).await?;
+
+        request.aes_key = Some(oke.get_aes_key());
         self.aes_key = Some(oke.get_aes_key());
 
         if request.method == "POST" {
@@ -122,7 +124,7 @@ async fn _handle(
         }
     };
 
-    let mut route = router.get_handler(request.get_olps());
+    let mut route = router.get_handler(&request.olps);
     let status_code = match response(
         &mut route,
         stream,
@@ -142,7 +144,7 @@ async fn _handle(
         }
     };
 
-    Ok((request.clone(), status_code))
+    Ok((request, status_code))
 }
 
 pub async fn handle(router: Router, stream: TcpStream, peer: SocketAddr) {
@@ -160,9 +162,7 @@ pub async fn handle(router: Router, stream: TcpStream, peer: SocketAddr) {
                 status_code
             );
         }
-        Err(error) => {
-            println!("{}", error);
-        }
+        Err(error) => println!("{}", error),
     }
 }
 

diff --git a/src/utils/parser.rs b/src/utils/parser.rs
@@ -158,7 +158,7 @@ impl Oblivion {
 #[derive(Clone, Debug, PartialEq)]
 pub struct OblivionRequest {
     pub(crate) method: String,
-    olps: String,
+    pub(crate) olps: String,
     protocol: String,
     version: String,
     data: Option<String>,
@@ -167,6 +167,7 @@ pub struct OblivionRequest {
     put: Option<Vec<u8>>,
     remote_addr: Option<String>,
     remote_port: Option<i32>,
+    pub(crate) aes_key: Option<Vec<u8>>,
 }
 
 impl OblivionRequest {
@@ -207,16 +208,17 @@ impl OblivionRequest {
                 .unwrap_or_default()
                 .to_string();
             Ok(Self {
-                method: method,
-                olps: olps,
-                protocol: protocol,
-                version: version,
+                method,
+                olps,
+                protocol,
+                version,
                 data: None,
                 plain_text: plain_text.to_owned(),
                 post: None,
                 put: None,
                 remote_addr: None,
                 remote_port: None,
+                aes_key: None,
             })
         } else {
             Err(OblivionException::BadProtocol {