diff --git a/server/src/config.rs b/server/src/config.rs index 79447ff..19f3ff6 100644 --- a/server/src/config.rs +++ b/server/src/config.rs @@ -74,6 +74,7 @@ pub struct ServerConfig { pub tcp_nodelay: bool, /// Whether or not to set "raw mode" for the file. pub file_raw_mode: bool, + #[serde(skip_serializing_if = "Option::is_none")] /// Keypair (public, private) in PEM format for TLS. pub tls_keypair: Option<[PathBuf; 2]>, @@ -82,12 +83,15 @@ pub struct ServerConfig { pub verbose_stats: bool, /// Whether or not to respond to stats requests over HTTP. pub enable_stats_endpoint: bool, + #[serde(skip_serializing_if = "String::is_empty")] /// Path of stats HTTP endpoint. pub stats_endpoint: String, + #[serde(skip_serializing_if = "String::is_empty")] /// String sent to a request that is not a websocket upgrade request. pub non_ws_response: String, + #[serde(skip_serializing_if = "String::is_empty")] /// Prefix of Wisp server. Do NOT add a trailing slash here. pub prefix: String, @@ -126,13 +130,17 @@ pub struct WispConfig { /// Whether or not to use Wisp version 2. pub wisp_v2: bool, + #[serde(skip_serializing_if = "Vec::is_empty")] /// Wisp version 2 extensions advertised. pub extensions: Vec, + #[serde(skip_serializing_if = "Option::is_none")] /// Wisp version 2 authentication extension advertised. pub auth_extension: Option, + #[serde(skip_serializing_if = "HashMap::is_empty")] /// Wisp version 2 password authentication extension username/passwords. pub password_extension_users: HashMap, + #[serde(skip_serializing_if = "Vec::is_empty")] /// Wisp version 2 certificate authentication extension public ed25519 pem keys. pub certificate_extension_keys: Vec, @@ -154,6 +162,7 @@ pub struct StreamConfig { #[cfg(feature = "twisp")] pub allow_twisp: bool, + #[serde(skip_serializing_if = "Vec::is_empty")] /// DNS servers to resolve with. Will default to system configuration. pub dns_servers: Vec, @@ -169,23 +178,31 @@ pub struct StreamConfig { /// Whether or not to allow connections to non-globally-routable IP addresses. pub allow_non_global: bool, + #[serde(skip_serializing_if = "Vec::is_empty")] /// Regex whitelist of hosts for TCP connections. pub allow_tcp_hosts: Vec, + #[serde(skip_serializing_if = "Vec::is_empty")] /// Regex blacklist of hosts for TCP connections. pub block_tcp_hosts: Vec, + #[serde(skip_serializing_if = "Vec::is_empty")] /// Regex whitelist of hosts for UDP connections. pub allow_udp_hosts: Vec, + #[serde(skip_serializing_if = "Vec::is_empty")] /// Regex blacklist of hosts for UDP connections. pub block_udp_hosts: Vec, + #[serde(skip_serializing_if = "Vec::is_empty")] /// Regex whitelist of hosts. pub allow_hosts: Vec, + #[serde(skip_serializing_if = "Vec::is_empty")] /// Regex blacklist of hosts. pub block_hosts: Vec, + #[serde(skip_serializing_if = "Vec::is_empty")] /// Range whitelist of ports. Format is `[lower_bound, upper_bound]`. pub allow_ports: Vec>, + #[serde(skip_serializing_if = "Vec::is_empty")] /// Range blacklist of ports. Format is `[lower_bound, upper_bound]`. pub block_ports: Vec>, } diff --git a/server/src/main.rs b/server/src/main.rs index 2312df1..bf86b69 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -1,7 +1,7 @@ #![feature(ip)] #![deny(clippy::todo)] -use std::{fmt::Write, fs::read_to_string}; +use std::{fmt::Write, fs::read_to_string, net::IpAddr}; use clap::Parser; use config::{validate_config_cache, Cli, Config}; @@ -27,6 +27,29 @@ mod stream; type Client = (DashMap, bool); +pub enum Resolver { + Hickory(TokioAsyncResolver), + System, +} + +impl Resolver { + pub async fn resolve(&self, host: String) -> anyhow::Result>> { + match self { + Self::Hickory(resolver) => Ok(Box::new(resolver.lookup_ip(host).await?.into_iter())), + Self::System => Ok(Box::new( + tokio::net::lookup_host(host + ":0").await?.map(|x| x.ip()), + )), + } + } + + pub fn clear_cache(&self) { + match self { + Self::Hickory(resolver) => resolver.clear_cache(), + Self::System => {} + } + } +} + lazy_static! { pub static ref CLI: Cli = Cli::parse(); pub static ref CONFIG: Config = { @@ -37,21 +60,19 @@ lazy_static! { } }; pub static ref CLIENTS: DashMap = 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() + pub static ref RESOLVER: Resolver = { + if CONFIG.stream.dns_servers.is_empty() { + Resolver::System } else { - ( + Resolver::Hickory(TokioAsyncResolver::tokio( ResolverConfig::from_parts( None, Vec::new(), NameServerConfigGroup::from_ips_clear(&CONFIG.stream.dns_servers, 53, true), ), ResolverOpts::default(), - ) - }; - - TokioAsyncResolver::tokio(config, opts) + )) + } }; } diff --git a/server/src/stream.rs b/server/src/stream.rs index d6f4f46..a483b6b 100644 --- a/server/src/stream.rs +++ b/server/src/stream.rs @@ -126,10 +126,9 @@ impl ClientStream { } let packet = RESOLVER - .lookup_ip(packet.destination_hostname) + .resolve(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, diff --git a/wisp/src/lib.rs b/wisp/src/lib.rs index a9d846c..f30c8db 100644 --- a/wisp/src/lib.rs +++ b/wisp/src/lib.rs @@ -70,7 +70,7 @@ pub enum WispError { StreamAlreadyClosed, /// The websocket frame received had an invalid type. - WsFrameInvalidType, + WsFrameInvalidType(ws::OpCode), /// The websocket frame received was not finished. WsFrameNotFinished, /// Error specific to the websocket implementation. @@ -133,7 +133,7 @@ impl std::fmt::Display for WispError { Self::MaxStreamCountReached => write!(f, "Maximum stream count reached"), Self::IncompatibleProtocolVersion => write!(f, "Incompatible Wisp protocol version"), Self::StreamAlreadyClosed => write!(f, "Stream already closed"), - Self::WsFrameInvalidType => write!(f, "Invalid websocket frame type"), + Self::WsFrameInvalidType(ty) => write!(f, "Invalid websocket frame type: {:?}", ty), Self::WsFrameNotFinished => write!(f, "Unfinished websocket frame"), Self::WsImplError(err) => write!(f, "Websocket implementation error: {}", err), Self::WsImplSocketClosed => { diff --git a/wisp/src/packet.rs b/wisp/src/packet.rs index 3d143f4..82c6617 100644 --- a/wisp/src/packet.rs +++ b/wisp/src/packet.rs @@ -487,7 +487,7 @@ impl<'a> Packet<'a> { return Err(WispError::WsFrameNotFinished); } if frame.opcode != OpCode::Binary { - return Err(WispError::WsFrameInvalidType); + return Err(WispError::WsFrameInvalidType(frame.opcode)); } let mut bytes = frame.payload; if bytes.remaining() < 1 { @@ -511,7 +511,7 @@ impl<'a> Packet<'a> { return Err(WispError::WsFrameNotFinished); } if frame.opcode != OpCode::Binary { - return Err(WispError::WsFrameInvalidType); + return Err(WispError::WsFrameInvalidType(frame.opcode)); } let mut bytes = frame.payload; if bytes.remaining() < 5 { @@ -587,7 +587,7 @@ impl<'a> TryFrom> for Packet<'a> { return Err(Self::Error::WsFrameNotFinished); } if frame.opcode != ws::OpCode::Binary { - return Err(Self::Error::WsFrameInvalidType); + return Err(Self::Error::WsFrameInvalidType(frame.opcode)); } Packet::try_from(frame.payload) }