use hickory-dns resovler

server/src/config.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, ops::RangeInclusive, path::PathBuf};
+use std::{collections::HashMap, net::IpAddr, ops::RangeInclusive, path::PathBuf};
 
 use cfg_if::cfg_if;
 use clap::{Parser, ValueEnum};
@@ -11,7 +11,7 @@ use wisp_mux::extensions::{
 	ProtocolExtensionBuilder,
 };
 
-use crate::{CLI, CONFIG};
+use crate::{CLI, CONFIG, RESOLVER};
 
 #[derive(Serialize, Deserialize, Default, Debug)]
 #[serde(rename_all = "lowercase")]
@@ -114,6 +114,9 @@ pub struct StreamConfig {
 	#[cfg(feature = "twisp")]
 	pub allow_twisp: bool,
 
+	/// DNS servers to resolve with. Will default to system configuration.
+	pub dns_servers: Vec<IpAddr>,
+
 	/// Whether or not to allow connections to IP addresses.
 	pub allow_direct_ip: bool,
 	/// Whether or not to allow connections to loopback IP addresses.
@@ -200,8 +203,12 @@ lazy_static! {
 }
 
 pub fn validate_config_cache() {
+	// constructs regexes
 	let _ = CONFIG_CACHE.allowed_ports;
+	// constructs wisp config
 	CONFIG.wisp.to_opts().unwrap();
+	// constructs resolver
+	RESOLVER.clear_cache();
 }
 
 impl Default for ServerConfig {
@@ -274,6 +281,8 @@ impl Default for StreamConfig {
 			#[cfg(feature = "twisp")]
 			allow_twisp: false,
 
+			dns_servers: Vec::new(),
+
 			allow_direct_ip: true,
 			allow_loopback: true,
 			allow_multicast: true,

server/src/main.rs

@@ -6,6 +6,10 @@ use clap::Parser;
 use config::{validate_config_cache, Cli, Config};
 use dashmap::DashMap;
 use handle::{handle_wisp, handle_wsproxy};
+use hickory_resolver::{
+	config::{NameServerConfigGroup, ResolverConfig, ResolverOpts},
+	TokioAsyncResolver,
+};
 use lazy_static::lazy_static;
 use listener::ServerListener;
 use log::{error, info};
@@ -32,6 +36,22 @@ lazy_static! {
 		}
 	};
 	pub static ref CLIENTS: DashMap<String, Client> = DashMap::new();
+	pub static ref RESOLVER: TokioAsyncResolver = {
+		let (config, opts) = if CONFIG.stream.dns_servers.is_empty() {
+			hickory_resolver::system_conf::read_system_conf().unwrap()
+		} else {
+			(
+				ResolverConfig::from_parts(
+					None,
+					Vec::new(),
+					NameServerConfigGroup::from_ips_clear(&CONFIG.stream.dns_servers, 53, true),
+				),
+				ResolverOpts::default(),
+			)
+		};
+
+		TokioAsyncResolver::tokio(config, opts)
+	};
 }
 
 fn format_stream_type(stream_type: StreamType) -> &'static str {

server/src/stream.rs

@@ -10,10 +10,10 @@ use fastwebsockets::{FragmentCollector, Frame, OpCode, Payload, WebSocketError};
 use hyper::upgrade::Upgraded;
 use hyper_util::rt::TokioIo;
 use regex::RegexSet;
-use tokio::net::{lookup_host, TcpStream, UdpSocket};
+use tokio::net::{TcpStream, UdpSocket};
 use wisp_mux::{ConnectPacket, StreamType};
 
-use crate::CONFIG;
+use crate::{CONFIG, RESOLVER};
 
 fn match_addr(str: &str, allowed: &RegexSet, blocked: &RegexSet) -> bool {
 	blocked.is_match(str) && !allowed.is_match(str)
@@ -125,13 +125,15 @@ impl ClientStream {
 			return Ok(ResolvedPacket::Blocked);
 		}
 
-		let packet = lookup_host(packet.destination_hostname + ":0")
+		let packet = RESOLVER
+			.lookup_ip(packet.destination_hostname)
 			.await
 			.context("failed to resolve hostname")?
+			.iter()
 			.filter(|x| CONFIG.server.resolve_ipv6 || x.is_ipv4())
 			.map(|x| ConnectPacket {
 				stream_type: packet.stream_type,
-				destination_hostname: x.ip().to_string(),
+				destination_hostname: x.to_string(),
 				destination_port: packet.destination_port,
 			})
 			.next();