Add support for comments in nat.conf file

arloor · Apr 23, 2024 · 9c95cda · 9c95cda
1 parent e80cec6
commit 9c95cda
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 40 deletions.
diff --git a/README.md b/README.md
@@ -107,6 +107,7 @@ RANGE,50000,50010,baidu.com
 - SINGLE：单端口转发：本机49999端口转发到baidu.com:59999
 - RANGE：范围端口转发：本机50000-50010转发到baidu.com:50000-50010
 - 请确保配置文件符合格式要求，否则程序可能会出现不可预期的错误，包括但不限于你和你的服务器炸掉（认真
+- 以 `#` 开始的行会被当成注释
 
 高级用法：
 

diff --git a/nat.conf b/nat.conf
@@ -1,3 +1,6 @@
+# 测试1
 SINGLE,10000,443,baidu.com,tcp
+# 测试2
 RANGE,1000,2000,baidu.com
+    #测试3
 SINGLE,2222,22,localhost
diff --git a/src/config.rs b/src/config.rs
@@ -1,8 +1,9 @@
+#![deny(warnings)]
+use crate::ip;
+use log::info;
 use std::env;
 use std::fs;
-use crate::ip;
 use std::process::exit;
-use log::info;
 
 #[derive(Debug)]
 pub enum Protocol {
@@ -64,76 +65,76 @@ pub enum NatCell {
         dst_domain: String,
         protocol: Protocol,
     },
+    Comment {
+        content: String,
+    },
 }
 
 impl NatCell {
     pub fn build(&self) -> String {
         let dst_domain = match &self {
             NatCell::Single { dst_domain, .. } => dst_domain,
-            NatCell::Range { dst_domain, .. } => dst_domain
+            NatCell::Range { dst_domain, .. } => dst_domain,
+            NatCell::Comment { content } => return content.clone(),
         };
         let dst_ip = match ip::remote_ip(dst_domain) {
             Ok(s) => s,
             Err(_) => return "".to_string(),
         };
         // 从环境变量读取本机ip或自动探测
-        let local_ip = env::var("nat_local_ip").unwrap_or(
-            match ip::local_ip() {
-                Ok(s) => s,
-                Err(_) => return "".to_string(),
-            });
+        let local_ip = env::var("nat_local_ip").unwrap_or(match ip::local_ip() {
+            Ok(s) => s,
+            Err(_) => return "".to_string(),
+        });
 
         match &self {
-            NatCell::Range { port_start, port_end, dst_domain: _, protocol } =>
-                {
-                    format!("#{cell:?}\n\
+            NatCell::Range {
+                port_start,
+                port_end,
+                dst_domain: _,
+                protocol,
+            } => {
+                format!("# {cell:?}\n\
                     {tcpPrefix}add rule ip nat PREROUTING tcp dport {portStart}-{portEnd} counter dnat to {dstIp}:{portStart}-{portEnd}\n\
                     {udpPrefix}add rule ip nat PREROUTING udp dport {portStart}-{portEnd} counter dnat to {dstIp}:{portStart}-{portEnd}\n\
                     {tcpPrefix}add rule ip nat POSTROUTING ip daddr {dstIp} tcp dport {portStart}-{portEnd} counter snat to {localIP}\n\
                     {udpPrefix}add rule ip nat POSTROUTING ip daddr {dstIp} udp dport {portStart}-{portEnd} counter snat to {localIP}\n\n\
                     ", cell = self, portStart = port_start, portEnd = port_end, dstIp = dst_ip, localIP = local_ip, tcpPrefix = protocol.tcp_prefix(), udpPrefix = protocol.udp_prefix())
-                }
-            NatCell::Single { src_port, dst_port, dst_domain, protocol } =>
-                {
-                    if dst_domain == "localhost" || dst_domain == "127.0.0.1" { // 重定向到本机
-                        format!("#{cell:?}\n\
+            }
+            NatCell::Single {
+                src_port,
+                dst_port,
+                dst_domain,
+                protocol,
+            } => {
+                if dst_domain == "localhost" || dst_domain == "127.0.0.1" {
+                    // 重定向到本机
+                    format!("# {cell:?}\n\
                         {tcpPrefix}add rule ip nat PREROUTING tcp dport {localPort} redirect to :{remotePort}\n\
                         {udpPrefix}add rule ip nat PREROUTING udp dport {localPort} redirect to :{remotePort}\n\n\
                         ", cell = self, localPort = src_port, remotePort = dst_port, tcpPrefix = protocol.tcp_prefix(), udpPrefix = protocol.udp_prefix())
-                    } else { // 转发到其他机器
-                        format!("#{cell:?}\n\
+                } else {
+                    // 转发到其他机器
+                    format!("# {cell:?}\n\
                         {tcpPrefix}add rule ip nat PREROUTING tcp dport {localPort} counter dnat to {dstIp}:{dstPort}\n\
                         {udpPrefix}add rule ip nat PREROUTING udp dport {localPort} counter dnat to {dstIp}:{dstPort}\n\
                         {tcpPrefix}add rule ip nat POSTROUTING ip daddr {dstIp} tcp dport {dstPort} counter snat to {localIP}\n\
                         {udpPrefix}add rule ip nat POSTROUTING ip daddr {dstIp} udp dport {dstPort} counter snat to {localIP}\n\n\
                         ", cell = self, localPort = src_port, dstPort = dst_port, dstIp = dst_ip, localIP = local_ip, tcpPrefix = protocol.tcp_prefix(), udpPrefix = protocol.udp_prefix())
-                    }
                 }
-        }
-    }
-
-    pub fn get_target_ip(&self) -> (String, String) {
-        match &self {
-            NatCell::Range { port_start: _, port_end: _, dst_domain: remote_domain, protocol: _ } =>
-                (remote_domain.clone(), match ip::remote_ip(remote_domain) {
-                    Ok(s) => s,
-                    Err(_) => "".to_string()
-                })
-            ,
-            NatCell::Single { src_port: _local_port, dst_port: _remote_port, dst_domain: remote_domain, protocol: _ } =>
-                (remote_domain.clone(), match ip::remote_ip(remote_domain) {
-                    Ok(s) => s,
-                    Err(_) => "".to_string()
-                })
+            }
+            NatCell::Comment { .. } => "".to_string(),
         }
     }
 }
 
-
 pub fn example(conf: &String) {
     info!("请在 {} 编写转发规则，内容类似：", &conf);
-    info!("{}", "SINGLE,10000,443,baidu.com\n\
-                    RANGE,1000,2000,baidu.com")
+    info!(
+        "{}",
+        "SINGLE,10000,443,baidu.com\n\
+                    RANGE,1000,2000,baidu.com"
+    )
 }
 
 pub fn read_config(conf: String) -> Vec<NatCell> {
@@ -149,7 +150,10 @@ pub fn read_config(conf: String) -> Vec<NatCell> {
 
     let strs = contents.split('\n');
     for str in strs {
-        if str.starts_with('#'){
+        if str.trim().starts_with('#') {
+            nat_cells.push(NatCell::Comment {
+                content: str.trim().to_string()+"\n",
+            });
             continue;
         }
         let cells = str.trim().split(',').collect::<Vec<&str>>();

diff --git a/src/main.rs b/src/main.rs
@@ -58,7 +58,6 @@ fn main() {
         script += &script_prefix;
 
         for x in vec.iter() {
-            let (_domain, _ip) = x.get_target_ip();
             let string = x.build();
             script += &string;
         }