From b4c001de2f921b78997407034bf54b83e10a2c3e Mon Sep 17 00:00:00 2001 From: Gabriel Souza Franco Date: Fri, 16 Apr 2021 00:27:26 -0300 Subject: [PATCH] chore: cleanup string-based code in find_actual_destination --- src/server_server.rs | 101 ++++++++++++++++++++++++++++++------------- 1 file changed, 70 insertions(+), 31 deletions(-) diff --git a/src/server_server.rs b/src/server_server.rs index 1bee9d5..e85532e 100644 --- a/src/server_server.rs +++ b/src/server_server.rs @@ -45,6 +45,35 @@ use std::{ #[cfg(feature = "conduit_bin")] use rocket::{get, post, put}; +#[derive(Clone, Debug, PartialEq)] +enum FederationDestination { + Literal(SocketAddr), + Named(String, String), +} + +impl FederationDestination { + fn into_url(self) -> String { + match self { + Self::Literal(addr) => format!("https://{}", addr), + Self::Named(host, port) => format!("https://{}{}", host, port), + } + } + + fn into_uri(self) -> String { + match self { + Self::Literal(addr) => addr.to_string(), + Self::Named(host, ref port) => host + port, + } + } + + fn host(&self) -> String { + match &self { + Self::Literal(addr) => addr.ip().to_string(), + Self::Named(host, _) => host.clone() + } + } +} + #[tracing::instrument(skip(globals))] pub async fn send_request( globals: &crate::database::globals::Globals, @@ -69,17 +98,17 @@ where result } else { let result = find_actual_destination(globals, &destination).await; + let (actual_destination, host) = result.clone(); + let result = (result.0.into_url(), result.1.into_uri()); globals .actual_destination_cache .write() .unwrap() .insert(Box::::from(destination), result.clone()); - let actual_destination = result.0.strip_prefix("https://").unwrap().splitn(2, ':').next().unwrap(); - let host = result.1.splitn(2, ':').next().unwrap_or(&result.1); if actual_destination != host { globals.tls_name_override.write().unwrap().insert( - actual_destination.to_owned(), - webpki::DNSNameRef::try_from_ascii_str(&host) + actual_destination.host(), + webpki::DNSNameRef::try_from_ascii_str(&host.host()) .unwrap() .to_owned(), ); @@ -212,22 +241,23 @@ where } #[tracing::instrument] -fn get_ip_with_port(destination_str: String) -> Option { - if destination_str.parse::().is_ok() { - Some(destination_str) +fn get_ip_with_port(destination_str: &str) -> Option { + if let Ok(destination) = destination_str.parse::() { + Some(FederationDestination::Literal(destination)) } else if let Ok(ip_addr) = destination_str.parse::() { - Some(SocketAddr::new(ip_addr, 8448).to_string()) + Some(FederationDestination::Literal(SocketAddr::new(ip_addr, 8448))) } else { None } } #[tracing::instrument] -fn add_port_to_hostname(destination_str: String) -> String { - match destination_str.find(':') { - None => destination_str.to_owned() + ":8448", - Some(_) => destination_str.to_string(), - } +fn add_port_to_hostname(destination_str: &str) -> FederationDestination { + let (host, port) = match destination_str.find(':') { + None => (destination_str, ":8448"), + Some(pos) => destination_str.split_at(pos), + }; + FederationDestination::Named(host.to_string(), port.to_string()) } /// Returns: actual_destination, host header @@ -237,36 +267,37 @@ fn add_port_to_hostname(destination_str: String) -> String { async fn find_actual_destination( globals: &crate::database::globals::Globals, destination: &'_ ServerName, -) -> (String, String) { +) -> (FederationDestination, FederationDestination) { let destination_str = destination.as_str().to_owned(); - let mut host = destination_str.clone(); - let actual_destination = "https://".to_owned() - + &match get_ip_with_port(destination_str.clone()) { + let mut hostname = destination_str.clone(); + let actual_destination = match get_ip_with_port(&destination_str) { Some(host_port) => { // 1: IP literal with provided or default port host_port } None => { - if destination_str.find(':').is_some() { + if let Some(pos) = destination_str.find(':') { // 2: Hostname with included port - destination_str + let (host, port) = destination_str.split_at(pos); + FederationDestination::Named(host.to_string(), port.to_string()) } else { match request_well_known(globals, &destination.as_str()).await { // 3: A .well-known file is available Some(delegated_hostname) => { - host = delegated_hostname.clone(); - match get_ip_with_port(delegated_hostname.clone()) { + hostname = delegated_hostname.clone(); + match get_ip_with_port(&delegated_hostname) { Some(host_and_port) => host_and_port, // 3.1: IP literal in .well-known file None => { - if destination_str.find(':').is_some() { + if let Some(pos) = destination_str.find(':') { // 3.2: Hostname with port in .well-known file - destination_str + let (host, port) = destination_str.split_at(pos); + FederationDestination::Named(host.to_string(), port.to_string()) } else { match query_srv_record(globals, &delegated_hostname).await { // 3.3: SRV lookup successful Some(hostname) => hostname, // 3.4: No SRV records, just use the hostname from .well-known - None => add_port_to_hostname(delegated_hostname), + None => add_port_to_hostname(&delegated_hostname), } } } @@ -278,7 +309,7 @@ async fn find_actual_destination( // 4: SRV record found Some(hostname) => hostname, // 5: No SRV record found - None => add_port_to_hostname(destination_str.to_string()), + None => add_port_to_hostname(&destination_str), } } } @@ -286,24 +317,32 @@ async fn find_actual_destination( } }; - (actual_destination, host) + let hostname = get_ip_with_port(&hostname).unwrap_or_else(|| { + match hostname.find(':') { + Some(pos) => { + let (host, port) = hostname.split_at(pos); + FederationDestination::Named(host.to_string(), port.to_string()) + } + None => FederationDestination::Named(hostname, "".to_string()) + } + }); + (actual_destination, hostname) } #[tracing::instrument(skip(globals))] async fn query_srv_record( globals: &crate::database::globals::Globals, hostname: &'_ str, -) -> Option { +) -> Option { if let Ok(Some(host_port)) = globals .dns_resolver() .srv_lookup(format!("_matrix._tcp.{}", hostname)) .await .map(|srv| { srv.iter().next().map(|result| { - format!( - "{}:{}", - result.target().to_string().trim_end_matches('.'), - result.port().to_string() + FederationDestination::Named( + result.target().to_string().trim_end_matches('.').to_string(), + format!(":{}", result.port()) ) }) })