Merge pull request #33 from cxz66666/add-custom-dns

feat: support custom dns config

README.md

@@ -226,6 +226,8 @@ $ docker compose up -d
 
 + `udp-port-forwarding`: UDP 端口转发，格式为 `本地地址-远程地址,本地地址-远程地址,...`，例如 `127.0.0.1:53-10.10.0.21:53`。多个转发用 `,` 分隔
 
++ `custom-dns`: 指定自定义DNS解析结果，格式为 `域名:IP,域名:IP,...`，例如 `www.cc98.org:10.10.98.98,appservice.zju.edu.cn:10.203.8.198`。多个解析用 `,` 分隔:
+
 + `twf-id`: twfID 登录，调试用途，一般不需要加此参数
 
 + `config`: 指定配置文件，内容参考 `config.toml.example`。启用配置文件时其他参数无效

config.toml.example

@@ -23,4 +23,9 @@ port_forwarding = [
     { network_type = "tcp", bind_address = "127.0.0.1:9898", remote_address = "10.10.98.98:80" },
     { network_type = "tcp", bind_address = "127.0.0.1:9899", remote_address = "10.10.98.98:80" },
     { network_type = "udp", bind_address = "127.0.0.1:1053", remote_address = "10.10.0.21:53" }
+]
+
+custom_dns = [
+    { host_name = "appservice.zju.edu.cn", ip = "10.203.8.198"},
+    { host_name = "www.cc98.org", ip = "10.10.98.98"}
 ]

core/EasyConnectClient.go

@@ -21,6 +21,11 @@ type Forwarding struct {
 	RemoteAddress string
 }
 
+type CustomDNS struct {
+	HostName string
+	IP       string
+}
+
 var SocksBind string
 var SocksUser string
 var SocksPasswd string
@@ -35,6 +40,7 @@ var ProxyAll bool
 var ForwardingList []Forwarding
 var EnableKeepAlive bool
 var ZjuDnsServer string
+var CustomDNSList []CustomDNS
 
 type EasyConnectClient struct {
 	queryConn net.Conn
@@ -140,6 +146,15 @@ func StartClient(host string, port int, username string, password string, twfId
 		go client.ServeForwarding(strings.ToLower(singleForwarding.NetworkType), singleForwarding.BindAddress, singleForwarding.RemoteAddress)
 	}
 
+	for _, customDNS := range CustomDNSList {
+		ipAddr := net.ParseIP(customDNS.IP)
+		if ipAddr == nil {
+			log.Printf("Custom DNS for host_name %s is invalid, SKIP", customDNS.HostName)
+		}
+		SetPermantDns(customDNS.HostName, ipAddr)
+		log.Printf("Custom DNS %s -> %s\n", customDNS.HostName, customDNS.IP)
+	}
+
 	if EnableKeepAlive {
 		go client.ServeSocks5(SocksBind, ZjuDnsServer)
 		client.KeepAlive(ZjuDnsServer)

core/dns_cache.go

@@ -47,3 +47,15 @@ func SetDnsCache(host string, ip net.IP) {
 	}
 	dnsCaches.cache.Set(host, ip, cache.DefaultExpiration)
 }
+
+func SetPermantDns(host string, ip net.IP) {
+	once.Do(func() {
+		dnsCaches = &DnsCache{
+			cache: cache.New(time.Duration(DnsTTL)*time.Second, time.Duration(DnsTTL)*2*time.Second),
+		}
+	})
+	if DebugDump {
+		log.Printf("SetPermantDnsCache: %s -> %s", host, ip.String())
+	}
+	dnsCaches.cache.Set(host, ip, cache.NoExpiration)
+}

main.go

@@ -30,13 +30,19 @@ type (
 		ZjuDnsServer        *string                `toml:"zju_dns_server"`
 		DebugDump           *bool                  `toml:"debug_dump"`
 		PortForwarding      []SinglePortForwarding `toml:"port_forwarding"`
+		CustomDns           []SingleCustomDns      `toml:"custom_dns"`
 	}
 
 	SinglePortForwarding struct {
 		NetworkType   *string `toml:"network_type"`
 		BindAddress   *string `toml:"bind_address"`
 		RemoteAddress *string `toml:"remote_address"`
 	}
+
+	SingleCustomDns struct {
+		HostName *string `toml:"host_name"`
+		IP       *string `toml:"ip"`
+	}
 )
 
 func getTomlVal[T int | uint64 | string | bool](valPointer *T, defaultVal T) T {
@@ -51,7 +57,7 @@ func main() {
 	// CLI args
 	host, port, username, password := "", 0, "", ""
 	disableServerConfig, disableZjuConfig, disableZjuDns, disableMultiLine, disableKeepAlive := false, false, false, false, false
-	twfId, configFile, tcpPortForwarding, udpPortForwarding := "", "", "", ""
+	twfId, configFile, tcpPortForwarding, udpPortForwarding, customDns := "", "", "", "", ""
 
 	flag.StringVar(&host, "server", "rvpn.zju.edu.cn", "EasyConnect server address")
 	flag.IntVar(&port, "port", 443, "EasyConnect port address")
@@ -70,6 +76,7 @@ func main() {
 	flag.BoolVar(&core.DebugDump, "debug-dump", false, "Enable traffic debug dump (only for debug usage)")
 	flag.StringVar(&tcpPortForwarding, "tcp-port-forwarding", "", "TCP port forwarding (e.g. 0.0.0.0:9898-10.10.98.98:80,127.0.0.1:9899-10.10.98.98:80)")
 	flag.StringVar(&udpPortForwarding, "udp-port-forwarding", "", "UDP port forwarding (e.g. 127.0.0.1:53-10.10.0.21:53)")
+	flag.StringVar(&customDns, "custom-dns", "", "Custom set dns lookup (e.g. www.cc98.org:10.10.98.98,appservice.zju.edu.cn:10.203.8.198)")
 	flag.BoolVar(&disableKeepAlive, "disable-keep-alive", false, "Disable keep alive")
 	flag.StringVar(&core.ZjuDnsServer, "zju-dns-server", "10.10.0.21", "ZJU DNS server address")
 	flag.StringVar(&twfId, "twf-id", "", "Login using twfID captured (mostly for debug usage)")
@@ -140,6 +147,23 @@ func main() {
 			})
 		}
 
+		for _, singleCustomDns := range conf.CustomDns {
+			if singleCustomDns.HostName == nil {
+				fmt.Println("ZJU Connect: host_name is not set")
+				return
+			}
+
+			if singleCustomDns.IP == nil {
+				fmt.Println("ZJU Connect: IP is not set")
+				return
+			}
+
+			core.CustomDNSList = append(core.CustomDNSList, core.CustomDNS{
+				HostName: *singleCustomDns.HostName,
+				IP:       *singleCustomDns.IP,
+			})
+		}
+
 		if host == "" || (username == "" || password == "") {
 			fmt.Println("ZJU Connect: host, username and password can not be empty")
 			return
@@ -185,6 +209,22 @@ func main() {
 			}
 		}
 
+		if customDns != "" {
+			customDnsList := strings.Split(customDns, ",")
+			for _, singleCustomDns := range customDnsList {
+				singleCustomDnsSplit := strings.Split(singleCustomDns, ":")
+				if len(singleCustomDnsSplit) != 2 {
+					fmt.Println("ZJU Connect: wrong custom dns format")
+					return
+				}
+
+				core.CustomDNSList = append(core.CustomDNSList, core.CustomDNS{
+					HostName: singleCustomDnsSplit[0],
+					IP:       singleCustomDnsSplit[1],
+				})
+			}
+		}
+
 		if host == "" || ((username == "" || password == "") && twfId == "") {
 			fmt.Println("ZJU Connect")
 			fmt.Println("Please see: https://github.com/mythologyli/zju-connect")