matrix-oracle 0.2.0

//! Resolution for the server-server API

use std::net::{IpAddr, SocketAddr};

use hickory_resolver::{TokioResolver, net::NetError, proto::rr::RData};
use reqwest_middleware::ClientWithMiddleware;
use serde::{Deserialize, Serialize};
use tracing::{debug, info, instrument};

use crate::cache;

pub mod error;

/// well-known information about the delegated server for server-server
/// communication.
///
/// See [the specification] for more information.
///
/// [the specification]: https://matrix.org/docs/spec/server_server/latest#get-well-known-matrix-server
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ServerWellKnown {
	/// The server name to delegate server-server communications to, with
	/// optional port
	#[serde(rename = "m.server")]
	pub server: String,
}

/// Client for server-server well-known lookups.
#[derive(Debug, Clone)]
pub struct Resolver {
	/// HTTP client.
	http: ClientWithMiddleware,
	/// DNS resolver.
	resolver: TokioResolver,
}

/// Resolved server name
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Server {
	/// IP address with implicit default port (8448)
	Ip(IpAddr),
	/// IP address and explicit port
	Socket(SocketAddr),
	/// Host string with implicit default port (8448)
	Host(String),
	/// Host string with explicit port.
	HostPort(String),
	/// Address from srv record, hostname from server name.
	Srv(String, String),
}

impl Server {
	/// The value to use for the `Host` HTTP header.
	#[must_use]
	pub fn host_header(&self) -> String {
		match self {
			Server::Ip(addr) => addr.to_string(),
			Server::Socket(addr) => addr.to_string(),
			Server::Host(host) => host.clone(),
			Server::HostPort(host) => host.clone(),
			Server::Srv(_, host) => host.to_string(),
		}
	}

	/// The address to connect to.
	#[must_use]
	pub fn address(&self) -> String {
		match self {
			Server::Ip(addr) => format!("{}:8448", addr),
			Server::Socket(addr) => addr.to_string(),
			Server::Host(host) => format!("{}:8448", host),
			Server::HostPort(host) => host.clone(),
			Server::Srv(host, _) => host.clone(),
		}
	}
}

impl Resolver {
	/// Constructs a new client.
	pub fn new() -> Result<Self, NetError> {
		Ok(Self {
			http: reqwest_middleware::ClientBuilder::new(reqwest::Client::new())
				.with(cache())
				.build(),
			resolver: TokioResolver::builder_tokio()?.build()?,
		})
	}

	/// Constructs a new client with the given HTTP client and DNS resolver
	/// instances.
	#[must_use]
	pub fn with(http: reqwest::Client, resolver: TokioResolver) -> Self {
		Self { http: reqwest_middleware::ClientBuilder::new(http).with(cache()).build(), resolver }
	}

	/// Resolve the given server name
	#[instrument(skip(self, port), err)]
	pub async fn resolve(
		&self,
		name: &str,
		#[cfg(test)] port: Option<u16>,
	) -> error::Result<Server> {
		// 1. The host is an ip literal
		debug!("Parsing socket literal");
		if let Ok(addr) = name.parse::<SocketAddr>() {
			info!("The server name is a socket literal");
			return Ok(Server::Socket(addr));
		}
		debug!("Parsing IP literal");
		if let Ok(addr) = name.parse::<IpAddr>() {
			info!("The server name is an IP literal");
			return Ok(Server::Ip(addr));
		}
		// 2. The host is not an ip literal, but includes a port
		debug!("Parsing host with port");
		if split_port(name).is_some() {
			info!("The servername is a host with port");
			return Ok(Server::HostPort(name.to_owned()));
		}
		// 3. Query the .well-known endpoint
		debug!("Querying well known");
		if let Some(well_known) = self
			.well_known(
				name,
				#[cfg(test)]
				port,
			)
			.await?
		{
			debug!("Well-known received: {:?}", &well_known);
			// 3.1 delegated_hostname is an ip literal
			debug!("Parsing delegated socket literal");
			if let Ok(addr) = well_known.server.parse::<SocketAddr>() {
				info!("The server name is a delegated IP literal");
				return Ok(Server::Socket(addr));
			}
			debug!("Parsing delegated IP literal");
			if let Ok(addr) = well_known.server.parse::<IpAddr>() {
				info!("The server name is a delegated socket literal");
				return Ok(Server::Ip(addr));
			}
			// 3.2 delegated_hostname includes a port
			debug!("Parsing delegated hostname with port");
			if split_port(&well_known.server).is_some() {
				info!("The server name is a delegated hostname with port");
				return Ok(Server::HostPort(well_known.server));
			}
			// 3.3 Look up SRV record (fed record, then deprecated)
			debug!("Looking up SRV record for delegated hostname");
			if let Some(name) = self.srv_lookup_with_fallback(&well_known.server).await {
				info!("The server name is a delegated SRV record");
				return Ok(Server::Srv(name, well_known.server));
			}
			// 3.4 Use hostname in .well-known
			debug!("Using delegated hostname directly");
			return Ok(Server::Host(well_known.server));
		}
		// 4. The .well-known lookup failed, query SRV (fed record, then deprecated)
		debug!("Looking up SRV record for hostname");
		if let Some(srv) = self.srv_lookup_with_fallback(name).await {
			info!("The server name is an SRV record");
			return Ok(Server::Srv(srv, name.to_owned()));
		}
		// 5. No SRV record found, use hostname
		debug!("Using provided hostname directly");
		Ok(Server::Host(name.to_owned()))
	}

	/// Query the .well-known information for a host.
	#[cfg_attr(test, allow(unused_variables))]
	#[instrument(skip(self, name, port), err)]
	async fn well_known(
		&self,
		name: &str,
		#[cfg(test)] port: Option<u16>,
	) -> error::Result<Option<ServerWellKnown>> {
		#[cfg(not(test))]
		let response = self.http.get(format!("https://{}/.well-known/matrix/server", name)).send().await;

		#[cfg(test)]
		#[allow(clippy::expect_used)]
		let response = self
			.http
			.get(format!(
				"http://{name}:{port}/.well-known/matrix/server",
				port = port.expect("port needed for test env")
			))
			.send()
			.await;

		// Only return Err on connection failure, skip to next step for other errors.
		let response = match response {
			Ok(response) => response,
			Err(reqwest_middleware::Error::Reqwest(e)) if e.is_connect() => return Err(e.into()),
			Err(_) => return Ok(None),
		};
		let well_known = response.json::<ServerWellKnown>().await.ok();
		Ok(well_known)
	}

	/// Query a single SRV record set and return `host:port` of the
	/// lowest-priority record, if any.
	#[instrument(skip(self))]
	async fn srv_lookup(&self, query: &str) -> Option<String> {
		let srv = self.resolver.srv_lookup(query).await.ok()?;
		// Get a record with the lowest priority value
		match srv
			.answers()
			.iter()
			.filter_map(|record| match &record.data {
				RData::SRV(srv) => Some(srv),
				_ => None,
			})
			.min_by_key(|srv| srv.priority)
		{
			Some(srv) => {
				let target = srv.target.to_ascii();
				let host = target.trim_end_matches('.');
				Some(format!("{}:{}", host, srv.port))
			}
			None => None,
		}
	}

	/// Look up the federation SRV record, falling back to the deprecated
	/// `_matrix._tcp` record per the server-server resolution spec.
	#[instrument(skip(self))]
	async fn srv_lookup_with_fallback(&self, name: &str) -> Option<String> {
		if let Some(found) = self.srv_lookup(&fed_srv_name(name)).await {
			debug!("Found _matrix-fed._tcp SRV record");
			return Some(found);
		}
		debug!("No _matrix-fed._tcp record; trying deprecated _matrix._tcp");
		self.srv_lookup(&deprecated_srv_name(name)).await
	}

	/// Get the [`SocketAddr`] of an address
	pub async fn socket(&self, server: &Server) -> Result<SocketAddr, NetError> {
		let (host, port) = match *server {
			Server::Ip(ip) => return Ok(SocketAddr::new(ip, 8448)),
			Server::Socket(socket) => return Ok(socket),
			Server::Host(ref host) => (host.as_str(), 8448),
			#[allow(clippy::expect_used)]
			Server::HostPort(ref host) => split_port(host).expect("HostPort was constructed with port"),
			#[allow(clippy::expect_used)]
			Server::Srv(ref addr, _) => split_port(addr).expect("The SRV record includes the port"),
		};
		let record = self.resolver.lookup_ip(host).await?;
		// We naively get the first IP.
		let socket =
			SocketAddr::new(record.iter().next().ok_or(NetError::Message("No records"))?, port);
		Ok(socket)
	}
}

/// The current (spec v1.8+) federation SRV record name for a server name.
fn fed_srv_name(name: &str) -> String {
	format!("_matrix-fed._tcp.{name}")
}

/// The deprecated SRV record name, kept as a fallback per the spec.
fn deprecated_srv_name(name: &str) -> String {
	format!("_matrix._tcp.{name}")
}

/// Get the port at the end of a host string if there is one.
fn split_port(host: &str) -> Option<(&str, u16)> {
	match &host.split(':').collect::<Vec<_>>()[..] {
		[host, port] => match port.parse() {
			Ok(port) => Some((host, port)),
			Err(_) => None,
		},
		_ => None,
	}
}

#[cfg(test)]
mod tests {
	use std::net::{IpAddr, SocketAddr};

	use hickory_resolver::TokioResolver;
	use proptest::prelude::*;
	use wiremock::{
		Mock, MockServer, ResponseTemplate,
		matchers::{method, path},
	};

	use super::{Resolver, Server};

	/// Validates the SRV query-name construction for the fed record and the
	/// deprecated fallback record.
	#[test]
	fn srv_query_names() {
		assert_eq!(super::fed_srv_name("example.test"), "_matrix-fed._tcp.example.test");
		assert_eq!(super::deprecated_srv_name("example.test"), "_matrix._tcp.example.test");
	}

	/// Validates correct parsing of IP literals and server name with port
	#[tokio::test]
	async fn literals() -> Result<(), Box<dyn std::error::Error>> {
		let resolver = Resolver::new()?;
		assert_eq!(
			resolver.resolve("127.0.0.1", None).await?,
			Server::Ip(IpAddr::from([127, 0, 0, 1])),
			"1. IP literal"
		);
		assert_eq!(
			resolver.resolve("127.0.0.1:4884", None).await?,
			Server::Socket(SocketAddr::new(IpAddr::from([127, 0, 0, 1]), 4884)),
			"1. Socket literal"
		);
		assert_eq!(
			resolver.resolve("example.test:1234", None).await?,
			Server::HostPort(String::from("example.test:1234")),
			"2. Host with port"
		);
		Ok(())
	}

	/// Validates correct handing of the .well-known http endpoint.
	#[tokio::test]
	async fn http() -> Result<(), Box<dyn std::error::Error>> {
		let mock_server = MockServer::start().await;

		let client = reqwest::Client::builder()
			.resolve("example.test", *mock_server.address())
			.resolve("destination.test", *mock_server.address())
			.build()?;
		let resolver = Resolver::with(client, TokioResolver::builder_tokio()?.build()?);

		let addr = mock_server.address();

		Mock::given(method("GET"))
			.and(path("/.well-known/matrix/server"))
			.respond_with(
				ResponseTemplate::new(200).set_body_raw(
					format!(r#"{{"m.server": "{}"}}"#, addr.ip()),
					"application/json",
				),
			)
			.up_to_n_times(1)
			.expect(1)
			.mount(&mock_server)
			.await;

		assert_eq!(
			resolver.resolve("example.test", Some(addr.port())).await?,
			Server::Ip(addr.ip()),
			"3.1 delegated_hostname is an IP literal"
		);

		Mock::given(method("GET"))
			.and(path("/.well-known/matrix/server"))
			.respond_with(
				ResponseTemplate::new(200)
					.set_body_raw(format!(r#"{{"m.server": "{}"}}"#, addr), "application/json"),
			)
			.up_to_n_times(1)
			.expect(1)
			.mount(&mock_server)
			.await;

		assert_eq!(
			resolver.resolve("example.test", Some(addr.port())).await?,
			Server::Socket(*mock_server.address()),
			"3.1 delegated_hostname is a socket literal"
		);

		Mock::given(method("GET"))
			.and(path("/.well-known/matrix/server"))
			.respond_with(ResponseTemplate::new(200).set_body_raw(
				format!(r#"{{"m.server": "destination.test:{}"}}"#, addr.port()),
				"application/json",
			))
			.expect(1)
			.up_to_n_times(1)
			.mount(&mock_server)
			.await;

		assert_eq!(
			resolver.resolve("example.test", Some(addr.port())).await?,
			Server::HostPort(format!("destination.test:{}", addr.port())),
			"3.2 delegated_hostname includes a port"
		);
		Ok(())
	}

	proptest! {
		/// `split_port` round-trips a `host:port` string for any valid u16 port
		/// and non-colon host.
		#[test]
		fn split_port_roundtrip(host in "[a-z][a-z0-9-]{0,20}", port in any::<u16>()) {
			let input = format!("{host}:{port}");
			prop_assert_eq!(super::split_port(&input), Some((host.as_str(), port)));
		}

		/// A host with no colon never parses as host:port.
		#[test]
		fn split_port_rejects_bare_host(host in "[a-z][a-z0-9-]{0,20}") {
			prop_assert_eq!(super::split_port(&host), None);
		}

		/// SRV query builders always produce the spec-prefixed names.
		#[test]
		fn srv_names_prefixed(name in "[a-z][a-z0-9.-]{0,40}") {
			prop_assert!(super::fed_srv_name(&name).starts_with("_matrix-fed._tcp."));
			prop_assert!(super::deprecated_srv_name(&name).starts_with("_matrix._tcp."));
			prop_assert!(super::fed_srv_name(&name).ends_with(&name));
		}

		/// `Host` always derives the default federation port 8448 in `address`,
		/// and the bare host in `host_header`.
		#[test]
		fn host_derivations(host in "[a-z][a-z0-9.-]{0,40}") {
			let server = Server::Host(host.clone());
			prop_assert_eq!(server.address(), format!("{host}:8448"));
			prop_assert_eq!(server.host_header(), host);
		}
	}
}