implement relay v2 protocol

Signed-off-by: Uncle Jack <[email protected]>
bepass-org · Jul 2, 2024 · 57a1816 · 57a1816
1 parent ad9fdbe
commit 57a1816
Show file tree

Hide file tree

Showing 8 changed files with 138 additions and 21 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -31,6 +31,7 @@ toml = "0.8"
 fastrand = "2.1"
 async-trait = "0.1"
 schemars = { version = "0.8", features = ["uuid1"] }
+bincode = "2.0.0-rc.3"
 
 [profile.release]
 opt-level = "s"

diff --git a/src/config.rs b/src/config.rs
@@ -1,6 +1,6 @@
 use std::net::IpAddr;
 
-use crate::proxy::RequestContext;
+use crate::proxy::{Network, RequestContext};
 
 use cidr::IpCidr;
 use schemars::JsonSchema;
@@ -14,7 +14,8 @@ pub enum Protocol {
     Vmess,
     Vless,
     Bepass,
-    Relay,
+    RelayV1,
+    RelayV2,
     Blackhole,
     Freedom,
 }
@@ -64,8 +65,15 @@ impl Config {
         None
     }
 
-    pub fn dispatch_outbound(&self, addr: &str, _port: u16) -> Outbound {
-        if let Ok(ip) = addr.parse::<IpAddr>() {
+    pub fn dispatch_outbound(&self, context: &RequestContext) -> Outbound {
+        match &context.network {
+            Network::Udp => {
+                return self.outbound.clone();
+            }
+            _ => {}
+        }
+
+        if let Ok(ip) = context.address.clone().parse::<IpAddr>() {
             if self.outbound.r#match.iter().any(|cidr| cidr.contains(&ip)) {
                 return self.outbound.clone();
             }

diff --git a/src/proxy/bepass/inbound.rs b/src/proxy/bepass/inbound.rs
@@ -56,16 +56,16 @@ impl<'a> Proxy for BepassStream<'a> {
             ))?;
         let header = encoding::decode_request_header(request)?;
 
-        let outbound = self.config.dispatch_outbound(&header.address, header.port);
-
         let mut context = self.inbound.context.clone();
         {
             context.address = header.address.clone();
             context.port = header.port;
             context.network = header.network;
         }
 
+        let outbound = self.config.dispatch_outbound(&context);
         let mut upstream = crate::proxy::connect_outbound(context, outbound).await?;
+
         tokio::io::copy_bidirectional(self, &mut upstream).await?;
 
         Ok(())

diff --git a/src/proxy/mod.rs b/src/proxy/mod.rs
@@ -56,9 +56,9 @@ impl Network {
 
 #[derive(Default)]
 pub struct RequestContext {
-    address: String,
-    port: u16,
-    network: Network,
+    pub address: String,
+    pub port: u16,
+    pub network: Network,
     pub request: Option<Request>,
 }
 
@@ -105,7 +105,16 @@ async fn connect_outbound(ctx: RequestContext, outbound: Outbound) -> Result<Box
 
     let mut stream: Box<dyn Proxy> = match outbound.protocol {
         Protocol::Vless => Box::new(vless::outbound::VlessStream::new(ctx, outbound, socket)),
-        Protocol::Relay => Box::new(relay::outbound::RelayStream::new(ctx, socket)),
+        Protocol::RelayV1 => Box::new(relay::outbound::RelayStream::new(
+            ctx,
+            socket,
+            relay::outbound::RelayVersion::V1,
+        )),
+        Protocol::RelayV2 => Box::new(relay::outbound::RelayStream::new(
+            ctx,
+            socket,
+            relay::outbound::RelayVersion::V2,
+        )),
         Protocol::Blackhole => Box::new(blackhole::outbound::BlackholeStream),
         _ => Box::new(socket),
     };

diff --git a/src/proxy/relay/outbound.rs b/src/proxy/relay/outbound.rs
@@ -1,26 +1,59 @@
-use crate::proxy::{Proxy, RequestContext};
+use crate::proxy::{self, Proxy, RequestContext};
 
+use std::net::IpAddr;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
 use async_trait::async_trait;
+use bincode::{Decode, Encode};
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf};
 use worker::*;
 
+#[derive(Decode, Encode)]
+pub enum RelayVersion {
+    V1,
+    V2,
+}
+
+#[derive(Decode, Encode)]
+enum Network {
+    Tcp,
+    Udp,
+}
+
+impl Network {
+    fn from(net: &proxy::Network) -> Self {
+        match net {
+            proxy::Network::Tcp => Self::Tcp,
+            proxy::Network::Udp => Self::Udp,
+        }
+    }
+}
+
+#[derive(Decode, Encode)]
+struct Header {
+    pub ver: RelayVersion,
+    pub net: Network,
+    pub addr: IpAddr,
+    pub port: u16,
+}
+
 pub struct RelayStream {
     pub stream: Socket,
     context: RequestContext,
+    version: RelayVersion,
 }
 
 impl RelayStream {
-    pub fn new(context: RequestContext, stream: Socket) -> Self {
-        Self { context, stream }
+    pub fn new(context: RequestContext, stream: Socket, version: RelayVersion) -> Self {
+        Self {
+            context,
+            stream,
+            version,
+        }
     }
-}
 
-#[async_trait]
-impl Proxy for RelayStream {
-    async fn process(&mut self) -> Result<()> {
+    async fn process_v1(&mut self) -> Result<()> {
         let header = {
             let addr = &self.context.address;
             let port = self.context.port;
@@ -32,6 +65,48 @@ impl Proxy for RelayStream {
         self.stream.write_all(&header).await?;
         Ok(())
     }
+
+    async fn process_v2(&mut self) -> Result<()> {
+        // +---------+---------+---------+---------+---------+
+        // | 2 Bytes | 1 Byte  | 1 Byte  | n Bytes | 2 Bytes |
+        // +---------+---------+---------+---------+---------+
+        // | length  | version | network | address | port    |
+        // +---------+---------+---------+---------+---------+
+
+        let address = self
+            .context
+            .address
+            .parse::<IpAddr>()
+            .map_err(|_| Error::RustError("invalid ip address".to_string()))?;
+
+        let header = Header {
+            ver: RelayVersion::V2,
+            net: Network::from(&self.context.network),
+            addr: address,
+            port: self.context.port,
+        };
+
+        let mut slice = [0u8; 128];
+        let len = bincode::encode_into_slice(header, &mut slice, bincode::config::standard())
+            .map_err(|e| Error::RustError(format!("bincode {e}")))?;
+
+        let len_bytes = (len as u16).to_be_bytes();
+        self.stream
+            .write_all(&[&len_bytes, &slice[..len]].concat())
+            .await?;
+
+        Ok(())
+    }
+}
+
+#[async_trait]
+impl Proxy for RelayStream {
+    async fn process(&mut self) -> Result<()> {
+        match &self.version {
+            RelayVersion::V1 => self.process_v1().await,
+            RelayVersion::V2 => self.process_v2().await,
+        }
+    }
 }
 
 impl AsyncRead for RelayStream {

diff --git a/src/proxy/vless/inbound.rs b/src/proxy/vless/inbound.rs
@@ -49,15 +49,14 @@ impl<'a> Proxy for VlessStream<'a> {
         let uuid = self.inbound.uuid;
         let header = encoding::decode_request_header(&mut self, &uuid.into_bytes()).await?;
 
-        let outbound = self.config.dispatch_outbound(&header.address, header.port);
-
         let mut context = self.inbound.context.clone();
         {
             context.address = header.address;
             context.port = header.port;
             context.network = header.network;
         }
 
+        let outbound = self.config.dispatch_outbound(&context);
         let mut upstream = crate::proxy::connect_outbound(context, outbound).await?;
 
         // +-----------------------------------------------+------------------------------------+------------------------------------+---------------+

diff --git a/src/proxy/vmess/inbound.rs b/src/proxy/vmess/inbound.rs
@@ -49,15 +49,14 @@ impl<'a> Proxy for VmessStream<'a> {
         let uuid = self.inbound.uuid;
         let header = encoding::decode_request_header(&mut self, &uuid.into_bytes()).await?;
 
-        let outbound = self.config.dispatch_outbound(&header.address, header.port);
-
         let mut context = self.inbound.context.clone();
         {
             context.address = header.address;
             context.port = header.port;
             context.network = header.network;
         }
 
+        let outbound = self.config.dispatch_outbound(&context);
         let mut upstream = crate::proxy::connect_outbound(context, outbound).await?;
 
         let header =