diff --git a/.drone.jsonnet b/.drone.jsonnet index d6fcd1a..04716d2 100644 --- a/.drone.jsonnet +++ b/.drone.jsonnet @@ -1,5 +1,5 @@ local executableName = 'fourth'; -local build_image = 'img.kie.rs/jjkiers/rust-cross:rust1.70-zig'; +local build_image = 'img.kie.rs/jjkiers/rust-cross:rust1.71.1-zig'; local archs = [ { target: 'aarch64-unknown-linux-musl', short: 'arm64-musl' }, diff --git a/src/config.rs b/src/config.rs index 64c7699..154318a 100644 --- a/src/config.rs +++ b/src/config.rs @@ -173,6 +173,10 @@ fn load_config(path: &str) -> Result { name: name.to_string(), addr: format!("{}:{}", upstream_host, upsteam_port), protocol: upstream_url.scheme().to_string(), + addresses: Addr(Mutex::new(UpstreamAddress::new(format!( + "{}:{}", + upstream_host, upsteam_port + )))), ..Default::default() }), ); diff --git a/src/main.rs b/src/main.rs index 7d4d9b3..c223887 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,10 +7,10 @@ use crate::servers::Server; use log::{debug, error}; use std::env; +use std::path::Path; fn main() { - let config_path = - env::var("FOURTH_CONFIG").unwrap_or_else(|_| "/etc/fourth/config.yaml".to_string()); + let config_path = find_config(); let config = match Config::new(&config_path) { Ok(config) => config, @@ -27,3 +27,18 @@ fn main() { let _ = server.run(); error!("Server ended with errors"); } + +fn find_config() -> String { + let config_path = + env::var("FOURTH_CONFIG").unwrap_or_else(|_| "/etc/fourth/config.yaml".to_string()); + + if Path::new(&config_path).exists() { + return config_path; + } + + if Path::new("config.yaml").exists() { + return String::from("config.yaml"); + } + + String::from("") +} diff --git a/src/servers/mod.rs b/src/servers/mod.rs index bda4e79..0161e68 100644 --- a/src/servers/mod.rs +++ b/src/servers/mod.rs @@ -13,7 +13,6 @@ use protocol::tcp; #[derive(Debug)] pub(crate) struct Server { pub proxies: Vec>, - pub config: ParsedConfig, } #[derive(Debug, Clone)] @@ -31,7 +30,6 @@ impl Server { pub fn new(config: ParsedConfig) -> Self { let mut new_server = Server { proxies: Vec::new(), - config: config.clone(), }; for (name, proxy) in config.servers.iter() { diff --git a/src/servers/protocol/tcp.rs b/src/servers/protocol/tcp.rs index f33a7f8..f4565f6 100644 --- a/src/servers/protocol/tcp.rs +++ b/src/servers/protocol/tcp.rs @@ -72,21 +72,17 @@ async fn accept(inbound: TcpStream, proxy: Arc) -> Result<(), Box Result<(), Box> { match upstream { Upstream::Ban => { diff --git a/src/servers/upstream_address.rs b/src/servers/upstream_address.rs index 3220a12..5b07e04 100644 --- a/src/servers/upstream_address.rs +++ b/src/servers/upstream_address.rs @@ -19,6 +19,13 @@ impl Display for UpstreamAddress { } impl UpstreamAddress { + pub fn new(address: String) -> Self { + UpstreamAddress { + address, + ..Default::default() + } + } + pub fn is_valid(&self) -> bool { if let Some(resolved) = self.resolved_time { if let Some(ttl) = self.ttl { @@ -44,20 +51,24 @@ impl UpstreamAddress { pub async fn resolve(&mut self, mode: ResolutionMode) -> Result> { if self.is_resolved() && self.is_valid() { debug!( - "Already got address {:?}, still valid for {}", + "Already got address {:?}, still valid for {:.3}s", &self.resolved_addresses, - self.time_remaining() + self.time_remaining().as_seconds_f64() ); return Ok(self.resolved_addresses.clone()); } - debug!("Resolving addresses for {}", &self.address); + debug!( + "Resolving addresses for {} with mode {:?}", + &self.address, &mode + ); let lookup_result = tokio::net::lookup_host(&self.address).await; - let resolved_addresses = match lookup_result { - Ok(resolved_addresses) => resolved_addresses, + let resolved_addresses: Vec = match lookup_result { + Ok(resolved_addresses) => resolved_addresses.into_iter().collect(), Err(e) => { + debug!("Failed looking up {}: {}", &self.address, &e); // Protect against DNS flooding. Cache the result for 1 second. self.resolved_time = Some(Instant::now()); self.ttl = Some(Duration::seconds(3)); @@ -65,6 +76,8 @@ impl UpstreamAddress { } }; + debug!("Resolved addresses: {:?}", &resolved_addresses); + let addresses: Vec = match mode { ResolutionMode::Ipv4 => resolved_addresses .into_iter() @@ -76,10 +89,13 @@ impl UpstreamAddress { .filter(|a| a.is_ipv6()) .collect(), - _ => resolved_addresses.collect(), + _ => resolved_addresses, }; - debug!("Got addresses for {}: {:?}", &self.address, &addresses); + debug!( + "Got {} addresses for {}: {:?}", + &mode, &self.address, &addresses + ); debug!( "Resolved at {}", OffsetDateTime::now_utc() @@ -113,3 +129,13 @@ impl From<&str> for ResolutionMode { } } } + +impl Display for ResolutionMode { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + ResolutionMode::Ipv4 => write!(f, "IPv4Only"), + ResolutionMode::Ipv6 => write!(f, "IPv6Only"), + ResolutionMode::Ipv4AndIpv6 => write!(f, "IPv4 and IPv6"), + } + } +}