refactor file tree & support config file

6nu · Jul 8, 2021 · 352c841 · 352c841
1 parent d14d72b
commit 352c841
Show file tree

Hide file tree

Showing 12 changed files with 164 additions and 24 deletions.
diff --git a/src/cmd/mod.rs b/src/cmd/mod.rs
@@ -0,0 +1,25 @@
+use clap::{Arg, App, SubCommand};
+
+pub enum CmdInput {
+    Config(String),
+    None,
+}
+
+pub fn scan() -> CmdInput {
+    let matches = App::new("Realm")
+        .version("1.3-custom")
+        .about("A high efficiency proxy tool")
+        .arg(
+            Arg::with_name("config")
+                .short("c")
+                .long("config")
+                .value_name("json config file")
+                .help("specify a config file in json format")
+                .takes_value(true),
+        )
+        .get_matches();
+    if let Some(config) = matches.value_of("config") {
+        return CmdInput::Config(config.to_string());
+    }
+    CmdInput::None
+}
diff --git a/src/config/dns_mode.rs b/src/config/dns_mode.rs
@@ -0,0 +1,34 @@
+use serde::{Serialize, Deserialize};
+use trust_dns_resolver::config::LookupIpStrategy;
+
+#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
+pub enum DnsMode {
+    /// Only query for A (Ipv4) records
+    Ipv4Only,
+    /// Only query for AAAA (Ipv6) records
+    Ipv6Only,
+    /// Query for A and AAAA in parallel
+    Ipv4AndIpv6,
+    /// Query for Ipv4 if that fails, query for Ipv6 (default)
+    Ipv4thenIpv6,
+    /// Query for Ipv6 if that fails, query for Ipv4
+    Ipv6thenIpv4,
+}
+
+impl Default for DnsMode {
+    fn default() -> Self {
+        Self::Ipv4thenIpv6
+    }
+}
+
+impl DnsMode {
+    pub fn to_strategy(self) -> LookupIpStrategy {
+        match self {
+            Self::Ipv4Only => LookupIpStrategy::Ipv4Only,
+            Self::Ipv6Only => LookupIpStrategy::Ipv6Only,
+            Self::Ipv4AndIpv6 => LookupIpStrategy::Ipv4AndIpv6,
+            Self::Ipv4thenIpv6 => LookupIpStrategy::Ipv4thenIpv6,
+            Self::Ipv6thenIpv4 => LookupIpStrategy::Ipv6thenIpv4,
+        }
+    }
+}
diff --git a/src/config/endpoint_config.rs b/src/config/endpoint_config.rs
@@ -0,0 +1,16 @@
+use serde::{Serialize, Deserialize};
+
+use crate::relay::Endpoint;
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct EndpointConfig {
+    udp: bool,
+    local: String,
+    remote: String,
+}
+
+impl EndpointConfig {
+    pub fn to_endpoint(self) -> Endpoint {
+        Endpoint::new(&self.local, &self.remote, self.udp)
+    }
+}
diff --git a/src/config/mod.rs b/src/config/mod.rs
@@ -0,0 +1,23 @@
+use std::fs;
+
+use serde::{Serialize, Deserialize};
+
+mod dns_mode;
+mod endpoint_config;
+
+pub use dns_mode::DnsMode;
+pub use endpoint_config::EndpointConfig;
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct GlobalConfig {
+    #[serde(default)]
+    pub dns_mode: DnsMode,
+    pub endpoints: Vec<EndpointConfig>,
+}
+
+impl GlobalConfig {
+    pub fn from_config_file(file: &str) -> Self {
+        let config = fs::read_to_string(file).expect("invalid file path");
+        serde_json::from_str(&config).expect("failed to parse config file")
+    }
+}
diff --git a/src/main.rs b/src/main.rs
@@ -1,7 +1,34 @@
+mod cmd;
 mod relay;
+mod utils;
+mod config;
 
-#[tokio::main]
-async fn main() {
-    let eps = vec![relay::Endpoint::new("127.0.0.1:15000", "localhost:20000")];
-    relay::run(eps).await
+use cmd::CmdInput;
+use config::GlobalConfig;
+use relay::Endpoint;
+
+fn main() {
+    match cmd::scan() {
+        CmdInput::Config(c) => start_from_config(c),
+        CmdInput::None => {}
+    }
+}
+
+fn start_from_config(c: String) {
+    let config = GlobalConfig::from_config_file(&c);
+    relay::init_resolver(config.dns_mode.to_strategy());
+    let eps: Vec<Endpoint> = config
+        .endpoints
+        .into_iter()
+        .map(|epc| epc.to_endpoint())
+        .collect();
+    run_relay(eps);
+}
+
+fn run_relay(eps: Vec<Endpoint>) {
+    tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .build()
+        .unwrap()
+        .block_on(relay::run(eps))
 }
diff --git a/src/relay/dns.rs b/src/relay/dns.rs
@@ -3,17 +3,26 @@ use std::net::IpAddr;
 use tokio::io;
 use tokio::runtime::Runtime;
 use trust_dns_resolver::TokioAsyncResolver;
-use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};
+use trust_dns_resolver::config::{ResolverConfig, ResolverOpts, LookupIpStrategy};
 use lazy_static::lazy_static;
 
-use super::utils;
+use crate::utils;
+
+static mut RESOLVE_STRATEGY: LookupIpStrategy = LookupIpStrategy::Ipv4thenIpv6;
 
 lazy_static! {
-    pub static ref DNS: TokioAsyncResolver = TokioAsyncResolver::tokio(
-        ResolverConfig::default(),
-        ResolverOpts::default()
-    )
-    .unwrap();
+    static ref DNS: TokioAsyncResolver =
+        TokioAsyncResolver::tokio(ResolverConfig::default(), {
+            let mut opts = ResolverOpts::default();
+            opts.ip_strategy = unsafe { RESOLVE_STRATEGY };
+            opts
+        })
+        .unwrap();
+}
+
+pub fn init_resolver(strategy: LookupIpStrategy) {
+    unsafe { RESOLVE_STRATEGY = strategy };
+    lazy_static::initialize(&DNS);
 }
 
 pub fn resolve_sync(addr: &str) -> io::Result<IpAddr> {

diff --git a/src/relay/mod.rs b/src/relay/mod.rs
@@ -7,17 +7,21 @@ use tokio::net::TcpListener;
 mod dns;
 mod tcp;
 mod udp;
-mod utils;
-pub use utils::{Endpoint, RemoteAddr};
+mod types;
+
+pub use types::{Endpoint, RemoteAddr};
+pub use dns::init_resolver;
 
 #[cfg(target_os = "linux")]
 mod zero_copy;
 
 pub async fn run(eps: Vec<Endpoint>) {
     let mut workers = vec![];
     for ep in eps.into_iter() {
-        workers.push(tokio::spawn(proxy_tcp(ep.local, ep.remote.clone())));
-        workers.push(tokio::spawn(proxy_udp(ep.local, ep.remote)))
+        if ep.udp {
+            workers.push(tokio::spawn(proxy_udp(ep.local, ep.remote.clone())))
+        }
+        workers.push(tokio::spawn(proxy_tcp(ep.local, ep.remote)));
     }
     join_all(workers).await;
 }

diff --git a/src/relay/tcp.rs b/src/relay/tcp.rs
@@ -4,7 +4,7 @@ use tokio::io::{self, AsyncWriteExt};
 use tokio::net::tcp::{ReadHalf, WriteHalf};
 use tokio::net::TcpStream;
 
-use super::utils::RemoteAddr;
+use super::types::RemoteAddr;
 
 pub async fn proxy(
     mut inbound: TcpStream,

diff --git a/src/relay/utils.rs → src/relay/types.rs b/src/relay/utils.rs → src/relay/types.rs
@@ -9,14 +9,11 @@ pub enum RemoteAddr {
 }
 
 pub struct Endpoint {
+    pub udp: bool,
     pub local: SocketAddr,
     pub remote: RemoteAddr,
 }
 
-pub fn new_io_err(e: &str) -> io::Error {
-    io::Error::new(io::ErrorKind::Other, e)
-}
-
 impl RemoteAddr {
     pub async fn to_sockaddr(self) -> io::Result<SocketAddr> {
         match self {
@@ -39,7 +36,7 @@ impl RemoteAddr {
 }
 
 impl Endpoint {
-    pub fn new(local: &str, remote: &str) -> Self {
+    pub fn new(local: &str, remote: &str, udp: bool) -> Self {
         let local = local
             .to_socket_addrs()
             .expect("invalid local address")
@@ -55,6 +52,6 @@ impl Endpoint {
             let _ = dns::resolve_sync(&addr).unwrap();
             RemoteAddr::DomainName(addr, port)
         };
-        Endpoint { local, remote }
+        Endpoint { local, remote, udp }
     }
 }
diff --git a/src/relay/udp.rs b/src/relay/udp.rs
@@ -8,7 +8,7 @@ use tokio::net::UdpSocket;
 use tokio::sync::oneshot;
 use tokio::time::sleep;
 
-use super::utils::RemoteAddr;
+use super::types::RemoteAddr;
 
 const BUFFERSIZE: usize = 2048;
 const TIMEOUT: Duration = Duration::from_secs(60 * 15);

diff --git a/src/relay/zero_copy.rs b/src/relay/zero_copy.rs
@@ -1,6 +1,6 @@
 use std::ops::Drop;
 use tokio::io;
-use super::utils;
+use crate::utils;
 
 pub struct Pipe(pub i32, pub i32);
 

diff --git a/src/utils/mod.rs b/src/utils/mod.rs
@@ -0,0 +1,5 @@
+use tokio::io;
+
+pub fn new_io_err(e: &str) -> io::Error {
+    io::Error::new(io::ErrorKind::Other, e)
+}