feat(packet): allow hkdf and scrpyt algorithm

src/bin/main.rs

@@ -1,10 +1,11 @@
 use anyhow::Result;
-use oblivion::api::get;
+use oblivion::models::client::Client;
 use oblivion::models::render::{BaseResponse, Response};
 use oblivion::models::router::{RoutePath, RouteType, Router};
 use oblivion::models::server::Server;
 use oblivion::models::session::Session;
 use oblivion::path_route;
+use oblivion::utils::generator::{generate_key_pair, generate_random_salt, SharedKey};
 use oblivion_codegen::async_route;
 use serde_json::json;
 use std::env::args;
@@ -39,7 +40,8 @@ fn json(_sess: Session) -> Response {
 }
 
 #[async_route]
-async fn alive(mut _sess: Session) -> Response {
+async fn alive(mut sess: Session) -> Response {
+    sess.send("test".into(), 200).await?;
     Ok(BaseResponse::JsonResponse(
         json!({"status": true, "msg": "结束"}),
         200,
@@ -52,21 +54,50 @@ async fn main() -> Result<()> {
     if args.len() <= 1 {
         args.push("serve".to_string());
     }
+    if args.len() <= 2 {
+        args.push("/welcome".to_string());
+    }
     match args[1].as_str() {
+        "keygen" => {
+            let now = Instant::now();
+            generate_key_pair()?;
+            println!("执行时间: {}", now.elapsed().as_millis());
+        }
+        "dh" => {
+            let now = Instant::now();
+            let (pr, pu) = generate_key_pair()?;
+            let (alice_pr, alice_pu) = generate_key_pair()?;
+            let salt = generate_random_salt();
+            let mut shared_bob = SharedKey::new(&pr, &alice_pu);
+            let mut shared_alice = SharedKey::new(&alice_pr, &pu);
+            let bob_key = shared_bob.hkdf(&salt)?;
+            let alice_key = shared_alice.hkdf(&salt)?;
+            assert_eq!(bob_key, alice_key);
+            println!("执行时间: {}", now.elapsed().as_millis());
+        }
         "bench" => loop {
             let now = Instant::now();
-            let mut res = get("127.0.0.1:7076/welcome").await?;
-            println!("{}", res.text()?);
+            let mut client = Client::new("CONNECT", format!("127.0.0.1:7076{}", args[2]))?;
+            client.connect().await?;
+            client.recv().await?.text()?;
+            client.close().await?;
             println!("执行时间: {}", now.elapsed().as_millis());
         },
-        "socket" => todo!(),
+        "socket" => {
+            let mut client = Client::new("CONNECT", format!("127.0.0.1:7076{}", args[2]))?;
+            client.connect().await?;
+            client.recv().await?.text()?;
+            client.recv().await?.json()?;
+            client.close().await?;
+        }
         "serve" => {
             let mut router = Router::new();
 
             router.route(RoutePath::new("/handler", RouteType::Path), handler);
 
             path_route!(&mut router, "/welcome" => welcome);
             path_route!(&mut router, "/json" => json);
+            path_route!(&mut router, "/alive" => alive);
 
             let mut server = Server::new("0.0.0.0", 7076, router);
             server.run().await?;

src/lib.rs

@@ -8,9 +8,9 @@
 //! which makes it possible to apply it to message dispatching and just-in-time communication.
 pub extern crate oblivion_codegen;
 pub extern crate proc_macro;
-pub mod api;
+
+/// # Oblivion Exceptions
 pub mod exceptions;
-pub mod sessions;
 
 /// # Oblivion Utilities
 ///
@@ -34,15 +34,15 @@ pub mod models;
 ///
 /// ```rust
 /// use oblivion::path_route;
-/// use oblivion::utils::parser::OblivionRequest;
 /// use oblivion::models::render::{BaseResponse, Response};
 /// use oblivion_codegen::async_route;
 /// use oblivion::models::router::Router;
+/// use oblivion::models::session::Session;
 ///
 /// #[async_route]
-/// fn welcome(mut req: OblivionRequest) -> Response {
+/// fn welcome(mut sess: Session) -> Response {
 ///     Ok(BaseResponse::TextResponse(
-///        format!("欢迎进入信息绝对安全区, 来自[{}]的朋友", req.get_ip()),
+///        format!("欢迎进入信息绝对安全区, 来自[{}]的朋友", sess.get_ip()),
 ///        200,
 ///     ))
 /// }
@@ -69,15 +69,15 @@ macro_rules! path_route {
 ///
 /// ```rust
 /// use oblivion::startswith_route;
-/// use oblivion::utils::parser::OblivionRequest;
 /// use oblivion::models::render::{BaseResponse, Response};
 /// use oblivion_codegen::async_route;
 /// use oblivion::models::router::Router;
+/// use oblivion::models::session::Session;
 ///
 /// #[async_route]
-/// fn welcome(mut req: OblivionRequest) -> Response {
+/// fn welcome(mut sess: Session) -> Response {
 ///     Ok(BaseResponse::TextResponse(
-///        format!("欢迎进入信息绝对安全区, 来自[{}]的朋友", req.get_ip()),
+///        format!("欢迎进入信息绝对安全区, 来自[{}]的朋友", sess.get_ip()),
 ///        200,
 ///     ))
 /// }
@@ -106,17 +106,16 @@ macro_rules! startswith_route {
 /// Regular routing can be simply implemented using regular routing macros:
 ///
 /// ```rust
-/// use futures::future::{BoxFuture, FutureExt};
 /// use oblivion::regex_route;
-/// use oblivion::utils::parser::OblivionRequest;
 /// use oblivion::models::render::{BaseResponse, Response};
 /// use oblivion_codegen::async_route;
 /// use oblivion::models::router::Router;
+/// use oblivion::models::session::Session;
 ///
 /// #[async_route]
-/// fn welcome(mut req: OblivionRequest) -> Response {
+/// fn welcome(mut sess: Session) -> Response {
 ///     Ok(BaseResponse::TextResponse(
-///        format!("欢迎进入信息绝对安全区, 来自[{}]的朋友", req.get_ip()),
+///        format!("欢迎进入信息绝对安全区, 来自[{}]的朋友", sess.get_ip()),
 ///        200,
 ///     ))
 /// }

src/models/packet.rs

@@ -3,7 +3,7 @@ use crate::exceptions::OblivionException;
 use crate::utils::decryptor::decrypt_bytes;
 use crate::utils::encryptor::{encrypt_bytes, encrypt_plaintext};
 use crate::utils::gear::Socket;
-use crate::utils::generator::{generate_random_salt, generate_shared_key};
+use crate::utils::generator::{generate_random_salt, SharedKey};
 use crate::utils::parser::length;
 
 use anyhow::{Error, Result};
@@ -66,6 +66,12 @@ impl OSC {
     }
 }
 
+impl From<u32> for OSC {
+    fn from(value: u32) -> Self {
+        Self { status_code: value }
+    }
+}
+
 pub struct OKE<'a> {
     public_key: Option<PublicKey>,
     private_key: Option<&'a EphemeralSecret>,
@@ -97,11 +103,11 @@ impl<'a> OKE<'a> {
         let remote_public_key_length = stream.recv_usize().await?;
         let remote_public_key_bytes = stream.recv(remote_public_key_length).await?;
         self.remote_public_key = Some(PublicKey::from_sec1_bytes(&remote_public_key_bytes)?);
-        self.shared_aes_key = Some(generate_shared_key(
+        let mut shared_key = SharedKey::new(
             self.private_key.as_ref().unwrap(),
             self.remote_public_key.as_ref().unwrap(),
-            &self.salt.as_mut().unwrap(),
-        )?);
+        );
+        self.shared_aes_key = Some(shared_key.hkdf(&self.salt.as_mut().unwrap())?);
         Ok(self)
     }
 
@@ -111,11 +117,11 @@ impl<'a> OKE<'a> {
         self.remote_public_key = Some(PublicKey::from_sec1_bytes(&remote_public_key_bytes)?);
         let salt_length = stream.recv_usize().await?;
         self.salt = Some(stream.recv(salt_length).await?);
-        self.shared_aes_key = Some(generate_shared_key(
-            self.private_key.unwrap(),
-            &self.remote_public_key.unwrap(),
-            self.salt.as_mut().unwrap(),
-        )?);
+        let mut shared_key = SharedKey::new(
+            self.private_key.as_ref().unwrap(),
+            self.remote_public_key.as_ref().unwrap(),
+        );
+        self.shared_aes_key = Some(shared_key.hkdf(&self.salt.as_mut().unwrap())?);
         Ok(self)
     }
 

src/models/server.rs

@@ -29,18 +29,18 @@ async fn _handle(router: &mut Router, mut socket: Socket, peer: SocketAddr) -> R
         return Err(Error::from(error));
     }
 
-    let header = session.header.as_ref().unwrap().clone();
-    let ip_addr = session.request.as_mut().unwrap().get_ip();
+    let header = session.header();
+    let ip_addr = session.get_ip();
     let aes_key = session.aes_key.clone().unwrap();
 
     let arc_socket = Arc::clone(&session.socket);
-    let mut socket = arc_socket.lock().await;
 
     let mut route = router.get_handler(&session.request.as_ref().unwrap().olps)?;
     let mut callback = route.get_handler()(session).await?;
 
     let status_code = callback.get_status_code()?;
 
+    let mut socket = arc_socket.lock().await;
     OSC::from_u32(1).to_stream(&mut socket).await?;
     OED::new(Some(aes_key))
         .from_bytes(callback.as_bytes()?)?

src/models/session.rs

@@ -91,12 +91,9 @@ impl Session {
         Ok(())
     }
 
-    pub async fn send(
-        &mut self,
-        data: Vec<u8>,
-        status_code: u32,
-    ) -> Result<()> {
+    pub async fn send(&mut self, data: Vec<u8>, status_code: u32) -> Result<()> {
         let socket = &mut self.socket.lock().await;
+
         OSC::from_u32(0).to_stream(socket).await?;
         OED::new(Some(self.aes_key.clone().unwrap()))
             .from_bytes(data)?
@@ -105,4 +102,12 @@ impl Session {
         OSC::from_u32(status_code).to_stream(socket).await?;
         Ok(())
     }
+
+    pub fn header(&mut self) -> String {
+        self.header.clone().unwrap()
+    }
+
+    pub fn get_ip(&mut self) -> String {
+        self.request.as_mut().unwrap().get_ip()
+    }
 }