add udp(raknet) protocol.

b23r0 · Aug 26, 2022 · c998fd8 · c998fd8
1 parent 0f27ba7
commit c998fd8
Show file tree

Hide file tree

Showing 9 changed files with 530 additions and 6 deletions.
diff --git a/heroinn/src/main.rs b/heroinn/src/main.rs
@@ -167,6 +167,7 @@ impl HeroinnApp {
                                     .show_ui(ui, |ui| {
                                         ui.selectable_value(&mut self.combox_listen_protocol, HeroinnProtocol::TCP, "TCP");
                                         ui.selectable_value(&mut self.combox_listen_protocol, HeroinnProtocol::HTTP, "HTTP");
+                                        ui.selectable_value(&mut self.combox_listen_protocol, HeroinnProtocol::UDP, "UDP");
                                     });
                                 });
                                 strip.cell(|ui|{
@@ -235,6 +236,7 @@ impl HeroinnApp {
                                 .show_ui(ui, |ui| {
                                     ui.selectable_value(&mut self.combox_listen_protocol, HeroinnProtocol::TCP, format!("{:?}" , HeroinnProtocol::TCP));
                                     ui.selectable_value(&mut self.combox_listen_protocol, HeroinnProtocol::HTTP, format!("{:?}" , HeroinnProtocol::HTTP));
+                                    ui.selectable_value(&mut self.combox_listen_protocol, HeroinnProtocol::UDP, format!("{:?}" , HeroinnProtocol::UDP));
                                 });
                                 ui.end_row();
 

diff --git a/heroinn_client/src/config.rs b/heroinn_client/src/config.rs
@@ -10,7 +10,7 @@ pub fn master_configure() -> ConnectionInfo{
 
     if size == 0{
         return ConnectionInfo{
-            protocol : HeroinnProtocol::HTTP.to_u8(),
+            protocol : HeroinnProtocol::UDP.to_u8(),
             address : String::from("127.0.0.1:8000"),
             remark : String::from("Default"),
         };

diff --git a/heroinn_core/src/lib.rs b/heroinn_core/src/lib.rs
@@ -4,11 +4,13 @@ use std::net::SocketAddr;
 pub mod module;
 
 use heroinn_util::protocol::http::WSServer;
+use heroinn_util::protocol::udp::UDPServer;
 use heroinn_util::{*, protocol::{Server} , protocol::tcp::*, packet::Message};
 
 pub struct HeroinnServer{
     tcp_server : Option<TcpServer>,
     ws_server : Option<WSServer>,
+    udp_server : Option<UDPServer>,
     protocol : HeroinnProtocol
 }
 
@@ -34,6 +36,7 @@ impl HeroinnServer{
                     Ok(tcp_server) => Ok(Self{
                         tcp_server: Some(tcp_server) , 
                         ws_server : None,
+                        udp_server:  None, 
                         protocol
                     }),
                     Err(e) => Err(e),
@@ -44,6 +47,18 @@ impl HeroinnServer{
                     Ok(ws_server) => Ok(Self{
                         tcp_server:  None, 
                         ws_server : Some(ws_server),
+                        udp_server:  None, 
+                        protocol
+                    }),
+                    Err(e) => Err(e),
+                }
+            },
+            HeroinnProtocol::UDP => {
+                match UDPServer::new(format!("0.0.0.0:{}" , port).as_str() , HeroinnServer::cb_connection , cb_msg){
+                    Ok(udp_server) => Ok(Self{
+                        tcp_server:  None, 
+                        ws_server : None,
+                        udp_server : Some(udp_server),
                         protocol
                     }),
                     Err(e) => Err(e),
@@ -57,6 +72,7 @@ impl HeroinnServer{
         match self.protocol{
             HeroinnProtocol::TCP => self.tcp_server.as_mut().unwrap().sendto(peer_addr, buf),
             HeroinnProtocol::HTTP => self.ws_server.as_mut().unwrap().sendto(peer_addr, buf),
+            HeroinnProtocol::UDP => self.udp_server.as_mut().unwrap().sendto(peer_addr, buf),
             HeroinnProtocol::Unknow => panic!("unknow protocol"),
         }
     }
@@ -65,6 +81,7 @@ impl HeroinnServer{
         match self.protocol{
             HeroinnProtocol::TCP => self.tcp_server.as_ref().unwrap().local_addr(),
             HeroinnProtocol::HTTP => self.ws_server.as_ref().unwrap().local_addr(),
+            HeroinnProtocol::UDP => self.udp_server.as_ref().unwrap().local_addr(),
             HeroinnProtocol::Unknow => panic!("unknow protocol"),
         }
     }
@@ -77,6 +94,7 @@ impl HeroinnServer{
         match self.protocol{
             HeroinnProtocol::TCP => self.tcp_server.as_mut().unwrap().contains_addr(peer_addr),
             HeroinnProtocol::HTTP => self.ws_server.as_mut().unwrap().contains_addr(peer_addr),
+            HeroinnProtocol::UDP => self.udp_server.as_mut().unwrap().contains_addr(peer_addr),
             HeroinnProtocol::Unknow => panic!("unknow protocol"),
         }
     }
@@ -85,6 +103,7 @@ impl HeroinnServer{
         match self.protocol{
             HeroinnProtocol::TCP => self.tcp_server.as_mut().unwrap().close(),
             HeroinnProtocol::HTTP => self.ws_server.as_mut().unwrap().close(),
+            HeroinnProtocol::UDP => self.udp_server.as_mut().unwrap().close(),
             HeroinnProtocol::Unknow => panic!("unknow protocol"),
         }
     }

diff --git a/heroinn_util/Cargo.toml b/heroinn_util/Cargo.toml
@@ -17,4 +17,9 @@ path-absolutize = "3.0.13"
 simple_logger = { version = "2.2.0" , features = ["threads"]}
 md-5 = "0.10.1"
 net2 = "0.2.35"
-websocket = "0.26.5"
+websocket = "0.26.5"
+bytes = "1.1.0"
+rand = {version = "0.8.5", features = ["std"]}
+rust-raknet = {git = "https://github.com/b23r0/rust-raknet.git"}
+tokio = "1.20.1"
+lazy_static = "1.4.0"
diff --git a/heroinn_util/src/lib.rs b/heroinn_util/src/lib.rs
@@ -153,6 +153,7 @@ impl HeroinnServerCommandID{
 pub enum HeroinnProtocol{
     TCP,
     HTTP,
+    UDP,
     Unknow
 }
 
@@ -161,6 +162,7 @@ impl HeroinnProtocol{
         match self{
             HeroinnProtocol::TCP => 0x00,
             HeroinnProtocol::HTTP => 0x01,
+            HeroinnProtocol::UDP => 0x02,
             HeroinnProtocol::Unknow => 0xff,
         }
     }
@@ -169,6 +171,7 @@ impl HeroinnProtocol{
         match v{
             0x00 => HeroinnProtocol::TCP,
             0x01 => HeroinnProtocol::HTTP,
+            0x02 => HeroinnProtocol::UDP,
             _ => HeroinnProtocol::Unknow,
         }
     }

diff --git a/heroinn_util/src/protocol/http.rs b/heroinn_util/src/protocol/http.rs
@@ -87,8 +87,7 @@ impl Server for WSServer{
 
                                 if buf.len() == 6 {
                                     if buf[..4] == TUNNEL_FLAG{
-                                        let mut conns = connections_2.lock().unwrap();
-                                        let mut sender = conns.remove(&remote_addr).unwrap();
+                                        let mut sender = connections_2.lock().unwrap().remove(&remote_addr).unwrap();
 
                                         let port = [buf[4] , buf[5]];
                                         let port = u16::from_be_bytes(port);

diff --git a/heroinn_util/src/protocol/mod.rs b/heroinn_util/src/protocol/mod.rs
@@ -1,11 +1,12 @@
 
 pub mod tcp;
 pub mod http;
+pub mod udp;
 use std::{io::*, net::SocketAddr, ops::{Deref, DerefMut}};
 
 use crate::{HeroinnProtocol, packet::Message};
 
-use self::{tcp::TcpConnection, http::WSConnection};
+use self::{tcp::TcpConnection, http::WSConnection, udp::UDPConnection};
 
 static TUNNEL_FLAG : [u8;4] = [0x38, 0x38 , 0x38, 0x38];
 
@@ -38,6 +39,7 @@ pub struct ClientWrapper{
     typ : HeroinnProtocol,
     tcp_client : Option<TcpConnection>,
     http_client : Option<WSConnection>,
+    udp_client : Option<UDPConnection>,
 }
 
 impl Deref for ClientWrapper{
@@ -51,6 +53,9 @@ impl Deref for ClientWrapper{
             HeroinnProtocol::HTTP => {
                 self.http_client.as_ref().unwrap()
             },
+            HeroinnProtocol::UDP => {
+                self.udp_client.as_ref().unwrap()
+            },
             HeroinnProtocol::Unknow => panic!("unknow protocol"),
         }
     }
@@ -65,14 +70,17 @@ impl DerefMut for ClientWrapper{
             HeroinnProtocol::HTTP => {
                 self.http_client.as_mut().unwrap()
             },
+            HeroinnProtocol::UDP => {
+                self.udp_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(), http_client: self.http_client.clone() }
+        Self { typ: self.typ.clone(), tcp_client: self.tcp_client.clone(), http_client: self.http_client.clone() , udp_client : self.udp_client.clone() }
     }
 }
 
@@ -85,6 +93,7 @@ impl ClientWrapper{
                     typ : typ.clone(),
                     tcp_client : Some(client),
                     http_client: None,
+                    udp_client : None
                 })
             },
             HeroinnProtocol::HTTP => {
@@ -93,6 +102,16 @@ impl ClientWrapper{
                     typ : typ.clone(),
                     tcp_client : None,
                     http_client: Some(client),
+                    udp_client : None,
+                })
+            },
+            HeroinnProtocol::UDP => {
+                let client = UDPConnection::connect(address)?;
+                Ok(Self{
+                    typ : typ.clone(),
+                    tcp_client : None,
+                    http_client: None,
+                    udp_client : Some(client)
                 })
             },
             HeroinnProtocol::Unknow => {
@@ -110,6 +129,9 @@ pub fn create_tunnel(addr : &str , protocol : &HeroinnProtocol , server_local_po
         HeroinnProtocol::HTTP => {
             Box::new(WSConnection::tunnel(addr, server_local_port)?)
         },
+        HeroinnProtocol::UDP => {
+            Box::new(UDPConnection::tunnel(addr, server_local_port)?)
+        },
         HeroinnProtocol::Unknow => {
             return Err(std::io::Error::new(std::io::ErrorKind::NotFound, "not found"));
         },