fix server

examples/proxy.rs

@@ -111,11 +111,9 @@ async fn main() {
         let password = "wmproxy".to_string();
 
         let proxy = ProxyConfig::builder()
-            .bind_addr(addr)
+            .bind(addr)
             .username(Some(username.clone()))
             .password(Some(password.clone()))
-            .center(true)
-            .mode("server".to_string())
             .into_value()
             .unwrap();
 
@@ -124,10 +122,9 @@ async fn main() {
             .unwrap();
 
         let proxy = ProxyConfig::builder()
-            .bind_addr(addr)
+            .bind(addr)
             .username(Some(username.clone()))
             .password(Some(password.clone()))
-            .center(true)
             .server(Some(format!("{}", server_addr)))
             .into_value()
             .unwrap();

src/arg.rs

@@ -14,13 +14,11 @@
 
 use std::process::id;
 use std::{
-    fmt::Display,
     fs::File,
     io::{self, Read, Write},
-    net::{AddrParseError, IpAddr, Ipv4Addr, SocketAddr},
+    net::{IpAddr, Ipv4Addr, SocketAddr},
     path::PathBuf,
     process::exit,
-    str::FromStr,
 };
 
 use bpaf::*;
@@ -29,38 +27,15 @@ use webparse::{Request, Url};
 use wenmeng::Client;
 
 use crate::reverse::StreamConfig;
-use crate::ConfigDuration;
 use crate::{
     option::proxy_config,
     reverse::{HttpConfig, LocationConfig, ServerConfig, UpstreamConfig},
     ConfigHeader, ConfigLog, ConfigOption, FileServer, ProxyConfig, ProxyResult,
 };
+use crate::{ConfigDuration, WrapAddr};
 
 const VERSION: &str = env!("CARGO_PKG_VERSION");
 
-#[derive(Debug, Clone, Copy)]
-pub struct WrapAddr(pub SocketAddr);
-
-impl FromStr for WrapAddr {
-    type Err = AddrParseError;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        if s.starts_with(":") {
-            let addr = format!("127.0.0.1{s}").parse::<SocketAddr>()?;
-            Ok(WrapAddr(addr))
-        } else {
-            let addr = s.parse::<SocketAddr>()?;
-            Ok(WrapAddr(addr))
-        }
-    }
-}
-
-impl Display for WrapAddr {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.write_fmt(format_args!("{}", self.0))
-    }
-}
-
 #[derive(Debug, Clone, Bpaf)]
 #[allow(dead_code)]
 struct Shared {

src/config/mod.rs

@@ -17,6 +17,7 @@ mod log;
 mod header;
 mod rate;
 mod ip_sets;
+mod wrap;
 
 use std::{str::FromStr, fmt::{Display, self}, marker::PhantomData};
 
@@ -26,6 +27,7 @@ pub use self::log::ConfigLog;
 pub use self::header::{ConfigHeader, HeaderOper};
 pub use self::rate::ConfigRate;
 pub use self::ip_sets::*;
+pub use self::wrap::*;
 
 use serde::{Serializer, Deserializer, de::{Visitor, Error, self}};
 use serde_with::{SerializeAs, DeserializeAs};

src/config/wrap.rs

@@ -0,0 +1,37 @@
+// Copyright 2022 - 2024 Wenmeng See the COPYRIGHT
+// file at the top-level directory of this distribution.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+//
+// Author: tickbh
+// -----
+// Created Date: 2024/01/25 02:13:35
+
+use std::{fmt::Display, net::{AddrParseError, SocketAddr}, str::FromStr};
+
+
+#[derive(Debug, Clone, Copy)]
+pub struct WrapAddr(pub SocketAddr);
+
+impl FromStr for WrapAddr {
+    type Err = AddrParseError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if s.starts_with(":") {
+            let addr = format!("127.0.0.1{s}").parse::<SocketAddr>()?;
+            Ok(WrapAddr(addr))
+        } else {
+            let addr = s.parse::<SocketAddr>()?;
+            Ok(WrapAddr(addr))
+        }
+    }
+}
+
+impl Display for WrapAddr {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_fmt(format_args!("{}", self.0))
+    }
+}

src/option.rs

@@ -14,7 +14,7 @@ use std::{
     collections::{HashMap, HashSet},
     fs::File,
     io::{self, BufReader},
-    net::{IpAddr, Ipv4Addr, SocketAddr},
+    net::{IpAddr, SocketAddr},
     process,
     sync::Arc,
     time::Duration,
@@ -30,8 +30,7 @@ use tokio::net::TcpListener;
 use tokio_rustls::{rustls, TlsAcceptor};
 
 use crate::{
-    reverse::{HttpConfig, StreamConfig, UpstreamConfig},
-    CenterClient, Flag, Helper, MappingConfig, OneHealth, ProxyError, ProxyResult,
+    reverse::{HttpConfig, StreamConfig, UpstreamConfig}, CenterClient, Flag, Helper, MappingConfig, OneHealth, ProxyError, ProxyResult, WrapAddr
 };
 
 pub struct Builder {
@@ -53,26 +52,34 @@ impl Builder {
         })
     }
 
-    pub fn mode(self, mode: String) -> Builder {
+    // pub fn mode(self, mode: String) -> Builder {
+    //     self.and_then(|mut proxy| {
+    //         proxy.mode = mode;
+    //         Ok(proxy)
+    //     })
+    // }
+
+    pub fn add_flag(self, flag: Flag) -> Builder {
         self.and_then(|mut proxy| {
-            proxy.mode = mode;
+            proxy.flag.set(flag, true);
             Ok(proxy)
         })
     }
 
-    pub fn add_flag(self, flag: Flag) -> Builder {
+    pub fn bind(self, addr: SocketAddr) -> Builder {
         self.and_then(|mut proxy| {
-            proxy.flag.set(flag, true);
+            proxy.bind = Some(WrapAddr(addr));
             Ok(proxy)
         })
     }
 
-    pub fn bind_addr(self, addr: SocketAddr) -> Builder {
+    pub fn center_addr(self, addr: SocketAddr) -> Builder {
         self.and_then(|mut proxy| {
-            proxy.bind_addr = addr;
+            proxy.center_addr = Some(WrapAddr(addr));
             Ok(proxy)
         })
     }
+
 
     pub fn server(self, addr: Option<String>) -> Builder {
         self.and_then(|mut proxy| {
@@ -88,13 +95,6 @@ impl Builder {
         })
     }
 
-    pub fn center(self, center: bool) -> Builder {
-        self.and_then(|mut proxy| {
-            proxy.center = center;
-            Ok(proxy)
-        })
-    }
-
     pub fn tc(self, is_tls: bool) -> Builder {
         self.and_then(|mut proxy| {
             proxy.tc = is_tls;
@@ -197,33 +197,47 @@ fn default_bind_addr() -> SocketAddr {
     "127.0.0.1:8090".parse().unwrap()
 }
 
-fn default_bool_true() -> bool {
-    true
-}
-
 /// 代理类, 一个代理类启动一种类型的代理
 #[serde_as]
 #[derive(Debug, Clone, Serialize, Deserialize, Bpaf)]
 pub struct ProxyConfig {
+    /// 代理id
+    #[bpaf(
+        short('s'),
+        long
+    )]
+    #[serde(default)]
+    pub(crate) server_id: u32,
+
     /// 代理绑定端口地址
     #[bpaf(
-        fallback(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8090)),
-        display_fallback,
+        // fallback(Some(WrapAddr(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8090)))),
+        // display_fallback,
         short('b'),
         long
     )]
-    #[serde(default = "default_bind_addr")]
-    pub(crate) bind_addr: SocketAddr,
+    #[serde_as(as = "Option<DisplayFromStr>")]
+    pub(crate) bind: Option<WrapAddr>,
+
+    /// 代理绑定端口地址
+    #[bpaf(
+        // fallback(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8090)),
+        // display_fallback,
+        short('c'),
+        long
+    )]
+    #[serde_as(as = "Option<DisplayFromStr>")]
+    pub(crate) center_addr: Option<WrapAddr>,
 
     /// 代理种类, 如http https socks5
     #[bpaf(fallback(Flag::default()))]
     #[serde_as(as = "DisplayFromStr")]
     #[serde(default)]
     pub(crate) flag: Flag,
-    /// 启动程序类型
-    #[bpaf(fallback("client".to_string()))]
-    #[serde(default)]
-    pub(crate) mode: String,
+    // /// 启动程序类型
+    // #[bpaf(fallback("client".to_string()))]
+    // #[serde(default)]
+    // pub(crate) mode: String,
 
     /// 连接代理服务端地址
     #[bpaf(short('S'), long("server"))]
@@ -249,9 +263,6 @@ pub struct ProxyConfig {
     /// 内网映射的证书key
     pub(crate) map_key: Option<String>,
 
-    /// 是否启用协议转发
-    #[serde(default = "default_bool_true")]
-    pub(crate) center: bool,
     /// 连接服务端是否启用tls
     #[serde(default)]
     pub(crate) ts: bool,
@@ -314,9 +325,11 @@ impl Default for ConfigOption {
 impl Default for ProxyConfig {
     fn default() -> Self {
         Self {
+            server_id: 0,
             flag: Flag::HTTP | Flag::HTTPS | Flag::SOCKS5,
-            mode: "client".to_string(),
-            bind_addr: default_bind_addr(),
+            // mode: "client".to_string(),
+            bind: Some(WrapAddr(default_bind_addr())),
+            center_addr: None,
             server: None,
             username: None,
             password: None,
@@ -328,7 +341,6 @@ impl Default for ProxyConfig {
             map_cert: None,
             map_key: None,
 
-            center: false,
             ts: false,
             tc: false,
             two_way_tls: false,
@@ -561,27 +573,27 @@ cR+nZ6DRmzKISbcN9/m8I7xNWwU2cglrYa4NCHguQSrTefhRoZAfl8BEOW1rJVGC
         }
     }
 
-    pub fn is_client(&self) -> bool {
-        self.mode.eq_ignore_ascii_case("client")
-    }
+    // pub fn is_client(&self) -> bool {
+    //     self.mode.eq_ignore_ascii_case("client")
+    // }
 
-    pub fn is_server(&self) -> bool {
-        self.mode.eq_ignore_ascii_case("server")
-    }
+    // pub fn is_server(&self) -> bool {
+    //     self.mode.eq_ignore_ascii_case("server")
+    // }
 
     pub async fn bind(
         &self,
     ) -> ProxyResult<(
         Option<TlsAcceptor>,
         Option<Arc<ClientConfig>>,
         Option<TcpListener>,
+        Option<TcpListener>,
         Option<CenterClient>,
     )> {
-        let addr = self.bind_addr.clone();
         let proxy_accept = self.get_tls_accept().await.ok();
         let client = self.get_tls_request().await.ok();
         let mut center_client = None;
-        if self.center {
+        if self.bind.is_some() {
             if let Some(server) = self.server.clone() {
                 let mut center = CenterClient::new(
                     self.clone(),
@@ -609,9 +621,19 @@ cR+nZ6DRmzKISbcN9/m8I7xNWwU2cglrYa4NCHguQSrTefhRoZAfl8BEOW1rJVGC
                 center_client = Some(center);
             }
         }
-        log::info!("绑定代理：{:?}，提供代理功能。", addr);
-        let center_listener = Some(Helper::bind(addr).await?);
-        Ok((proxy_accept, client, center_listener, center_client))
+        let client_listener = if let Some(bind) = self.bind {
+            log::info!("绑定代理：{:?}，提供代理功能。", bind.0);
+            Some(Helper::bind(bind.0).await?)
+        } else {
+            None
+        };
+        let center_listener = if let Some(center) = self.center_addr {
+            log::info!("绑定代理：{:?}，提供中心代理功能。", center.0);
+            Some(Helper::bind(center.0).await?)
+        } else {
+            None
+        };
+        Ok((proxy_accept, client, client_listener, center_listener, center_client))
     }
 
     pub async fn bind_map(

src/prot/close.rs

@@ -22,29 +22,31 @@ use super::{ProtFrameHeader, read_short_string, write_short_string};
 /// 旧的Socket连接关闭, 接收到则关闭掉当前的连接
 #[derive(Debug)]
 pub struct ProtClose {
+    server_id: u32,
     sock_map: u32,
     reason: String,
 }
 
 impl ProtClose {
-    pub fn new(sock_map: u32) -> ProtClose {
-        ProtClose { sock_map, reason: String::new() }
+    pub fn new(server_id: u32, sock_map: u32) -> ProtClose {
+        ProtClose { server_id, sock_map, reason: String::new() }
     }
 
-    pub fn new_by_reason(sock_map: u32, reason: String) -> ProtClose {
-        ProtClose { sock_map, reason }
+    pub fn new_by_reason(server_id: u32, sock_map: u32, reason: String) -> ProtClose {
+        ProtClose { server_id, sock_map, reason }
     }
 
     pub fn parse<T: Buf>(header: ProtFrameHeader, mut buf: T) -> ProxyResult<ProtClose> {
         let reason = read_short_string(&mut buf)?;
         Ok(ProtClose {
+            server_id: header.server_id(),
             sock_map: header.sock_map(),
             reason,
         })
     }
 
     pub fn encode<B: Buf + BufMut>(self, buf: &mut B) -> ProxyResult<usize> {
-        let mut head = ProtFrameHeader::new(ProtKind::Close, ProtFlag::zero(), self.sock_map);
+        let mut head = ProtFrameHeader::new(ProtKind::Close, ProtFlag::zero(), self.sock_map, self.server_id);
         head.length = self.reason.as_bytes().len() as u32 + 1;
         let mut size = 0;
         size += head.encode(buf)?;