dns-uri 0.1.0

#![doc = include_str!("../README.md")]
#![no_std]

extern crate alloc;

use alloc::string::{String, ToString};
use core::fmt;
use core::net::{AddrParseError, IpAddr, SocketAddr};
use core::num::ParseIntError;
use core::str::FromStr;

#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
/// URI of a DNS server.
///
/// This implements parsing from the following string formats:
///
/// ```rust
/// # use dns_uri::Uri;
/// #
/// // Regular DNS, IPv4, without port (default port 53)
/// # let server: Uri = r#"
/// 8.8.8.8
/// # "#.trim().parse().unwrap();
/// # assert!(matches!(server, Uri::Regular { .. }));
/// // Regular DNS, IPv4, with port
/// # let server: Uri = r#"
/// 8.8.8.8:10053
/// # "#.trim().parse().unwrap();
/// # assert!(matches!(server, Uri::Regular { .. }));
/// // Regular DNS, IPv6, without port (default port 53)
/// # let server: Uri = r#"
/// [2001:4860:4860::8888]
/// # "#.trim().parse().unwrap();
/// # assert!(matches!(server, Uri::Regular { .. }));
/// // We don't accept bare IPv6 address without blanket.
/// # assert!(matches!(server, Uri::Regular { .. }));
/// # let server = r#"
/// 2001:4860:4860::8888
/// # "#.trim().parse::<Uri>().unwrap_err();
/// // Regular DNS, IPv6, with port
/// # let server: Uri = r#"
/// [2001:4860:4860::8888]:10053
/// # "#.trim().parse().unwrap();
/// # assert!(matches!(server, Uri::Regular { .. }));
/// // Regular DNS, in URI format.
/// # let server: Uri = r#"
/// udp://8.8.8.8
/// # "#.trim().parse().unwrap();
/// # let server: Uri = r#"
/// udp://8.8.8.8:10053
/// # "#.trim().parse().unwrap();
/// # assert!(matches!(server, Uri::Regular { .. }));
/// # let server: Uri = r#"
/// tcp://[2001:4860:4860::8888]
/// # "#.trim().parse().unwrap();
/// # assert!(matches!(server, Uri::Regular { .. }));
/// # let server: Uri = r#"
/// tcp://[2001:4860:4860::8888]:10053
/// # "#.trim().parse().unwrap();
/// # assert!(matches!(server, Uri::Regular { .. }));
/// // DNS over TLS.
/// # let server: Uri = r#"
/// tls://dns.google
/// # "#.trim().parse().unwrap();
/// # assert!(matches!(server, Uri::TLS { .. }));
/// # let server: Uri = r#"
/// tls://dns.google:10853
/// # "#.trim().parse().unwrap();
/// # assert!(matches!(server, Uri::TLS { .. }));
/// // DNS over HTTPS, without custom endpoint.
/// # let server: Uri = r#"
/// https://dns.google
/// # "#.trim().parse().unwrap();
/// # assert!(matches!(server, Uri::HTTPS { .. }));
/// # let server: Uri = r#"
/// https://dns.google:8443
/// # "#.trim().parse().unwrap();
/// # assert!(matches!(server, Uri::HTTPS { .. }));
/// // DNS over HTTPS, with custom endpoint.
/// # let server: Uri = r#"
/// https://dns.google/dns-query
/// # "#.trim().parse().unwrap();
/// # assert!(matches!(server, Uri::HTTPS { .. }));
/// // For DoH, a root path `/` is also considered as a custom endpoint, please pay attention to this.
/// # let server: Uri = r#"
/// https://dns.google/
/// # "#.trim().parse().unwrap();
/// // For DoH / DoQ / DoT, you can specify custom SNI via query parameter `sni`.
/// # let server: Uri = r#"
/// tls://8.8.8.8?sni=dns.google
/// # "#.trim().parse().unwrap();
/// ```
pub enum Uri {
    #[non_exhaustive]
    /// Regular DNS, over UDP or TCP.
    Regular {
        /// Server IP address
        addr: IpAddr,

        /// Server port
        port: u16,

        /// Prefer TCP over UDP.
        prefer_tcp: bool,
    },

    #[non_exhaustive]
    /// DNS over TLS
    TLS {
        /// Server [`Host`]
        host: Host,

        /// Server port
        port: u16,
    },

    #[non_exhaustive]
    /// DNS over HTTPS
    HTTPS {
        /// Server [`Host`]
        host: Host,

        /// Server port
        port: u16,

        /// The HTTP endpoint where the DNS `NameServer` provides service. Only
        /// relevant to DNS-over-HTTPS.
        custom_http_endpoint: Option<String>,

        /// Force HTTP/3 (aka., DNS over HTTP/3).
        force_http3: bool,
    },

    #[non_exhaustive]
    /// DNS over QUIC
    QUIC {
        /// Server [`Host`]
        host: Host,

        /// Server port
        port: u16,
    },
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
/// The parsed host of a DNS server.
pub enum Host {
    #[non_exhaustive]
    /// An IP address.
    IpAddr {
        /// The parsed IP address.
        addr: IpAddr,

        /// A custom server name (SNI)
        custom_server_name: Option<String>,
    },

    #[non_exhaustive]
    /// A domain name.
    ServerName {
        /// The parsed domain name.
        name: String,
    },
}

impl Host {
    #[doc(hidden)]
    #[must_use]
    pub fn new_ip_addr(addr: IpAddr, custom_server_name: Option<String>) -> Self {
        Self::IpAddr {
            addr,
            custom_server_name,
        }
    }

    #[doc(hidden)]
    #[must_use]
    pub fn new_server_name(name: String) -> Self {
        Self::ServerName { name }
    }
}

const DEFAULT_PORT_DNS: u16 = 53;
const DEFAULT_PORT_DNS_OVER_TLS: u16 = 853;
const DEFAULT_PORT_DNS_OVER_QUIC: u16 = 853;
const DEFAULT_PORT_DNS_OVER_HTTPS: u16 = 443;

const SCHEME_UDP: &str = "udp";
const SCHEME_TCP: &str = "tcp";
const SCHEME_TLS: &str = "tls";
const SCHEME_HTTPS: &str = "https";
const SCHEME_H3: &str = "h3";
const SCHEME_QUIC: &str = "quic";

impl fmt::Display for Uri {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Regular { addr, port, prefer_tcp } => {
                write!(
                    f,
                    "{}://{}",
                    if *prefer_tcp { SCHEME_TCP } else { SCHEME_UDP },
                    SocketAddr::new(*addr, *port)
                )
            }
            Self::TLS { host, port } => match host {
                Host::IpAddr {
                    addr,
                    custom_server_name: None,
                } => {
                    write!(f, "{SCHEME_TLS}://{}", SocketAddr::new(*addr, *port))
                }
                Host::IpAddr {
                    addr,
                    custom_server_name: Some(custom_server_name),
                } => write!(
                    f,
                    "{SCHEME_TLS}://{}?sni={custom_server_name}",
                    SocketAddr::new(*addr, *port),
                ),
                Host::ServerName { name } => write!(f, "{SCHEME_TLS}://{name}:{port}"),
            },
            Self::HTTPS {
                host,
                port,
                custom_http_endpoint,
                force_http3,
            } => match host {
                Host::IpAddr {
                    addr,
                    custom_server_name: None,
                } => {
                    write!(
                        f,
                        "{}://{}{}",
                        if *force_http3 { SCHEME_H3 } else { SCHEME_HTTPS },
                        SocketAddr::new(*addr, *port),
                        custom_http_endpoint.as_deref().unwrap_or(""),
                    )
                }
                Host::IpAddr {
                    addr,
                    custom_server_name: Some(custom_server_name),
                } => write!(
                    f,
                    "{}://{}{}?sni={}",
                    if *force_http3 { SCHEME_H3 } else { SCHEME_HTTPS },
                    SocketAddr::new(*addr, *port),
                    custom_http_endpoint.as_deref().unwrap_or(""),
                    custom_server_name,
                ),
                Host::ServerName { name } => write!(
                    f,
                    "{}://{}:{}{}",
                    if *force_http3 { SCHEME_H3 } else { SCHEME_HTTPS },
                    name,
                    port,
                    custom_http_endpoint.as_deref().unwrap_or(""),
                ),
            },
            Self::QUIC { host, port } => match host {
                Host::IpAddr {
                    addr,
                    custom_server_name: None,
                } => {
                    write!(f, "{SCHEME_QUIC}://{}", SocketAddr::new(*addr, *port))
                }
                Host::IpAddr {
                    addr,
                    custom_server_name: Some(custom_server_name),
                } => write!(
                    f,
                    "{SCHEME_QUIC}://{}?sni={custom_server_name}",
                    SocketAddr::new(*addr, *port),
                ),
                Host::ServerName { name } => write!(f, "{SCHEME_QUIC}://{name}:{port}"),
            },
        }
    }
}

impl FromStr for Uri {
    type Err = ParseError;

    #[allow(clippy::too_many_lines)]
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if !s.contains('/') {
            // Try to parse as a plain IP address or socket address.

            let (closing_blacket_index, colon_index) = if s.starts_with('[') {
                let r @ Some(closing_blacket_index) = memchr::memchr(b']', s.as_bytes()) else {
                    return Err(ParseError::InvalidIpAddr);
                };

                (
                    r,
                    s.as_bytes().get(closing_blacket_index + 1).and_then(|colon| {
                        if colon == &b':' {
                            Some(closing_blacket_index + 1)
                        } else {
                            None
                        }
                    }),
                )
            } else {
                (None, memchr::memchr(b':', s.as_bytes()))
            };

            match (closing_blacket_index, colon_index) {
                // An IPv6 address with port, e.g., [::1]:53
                (Some(closing_blacket_index), Some(colon)) => {
                    return Ok(Self::Regular {
                        addr: IpAddr::V6(s[1..closing_blacket_index].parse()?),
                        port: s[colon + 1..].parse()?,
                        prefer_tcp: false,
                    });
                }
                // An IPv4 address with port, e.g., 127.0.0.1:53
                (None, Some(colon)) => {
                    return Ok(Self::Regular {
                        addr: IpAddr::V4(s[..colon].parse()?),
                        port: s[colon + 1..].parse()?,
                        prefer_tcp: false,
                    });
                }
                // An IPv6 address (with blanket) without port, e.g., [::1]
                (Some(closing_blacket_index), None) => {
                    return Ok(Self::Regular {
                        addr: IpAddr::V6(s[1..closing_blacket_index].parse()?),
                        port: DEFAULT_PORT_DNS,
                        prefer_tcp: false,
                    });
                }
                // An IPv4 address without port, e.g., 127.0.0.1
                (None, None) => {
                    return Ok(Self::Regular {
                        addr: IpAddr::V4(s.parse()?),
                        port: DEFAULT_PORT_DNS,
                        prefer_tcp: false,
                    });
                }
            }
        }

        let uri = fluent_uri::Uri::parse(s)?;

        let authority = uri.authority().ok_or(ParseError::MissingHost)?;

        let server_ip_addr = match authority.host_parsed() {
            fluent_uri::component::Host::Ipv4(ipv4_addr) => Some(IpAddr::V4(ipv4_addr)),
            fluent_uri::component::Host::Ipv6(ipv6_addr) => Some(IpAddr::V6(ipv6_addr)),
            _ => None,
        };
        let server_port = authority.port_to_u16()?;

        match uri.scheme().as_str() {
            r @ (SCHEME_UDP | SCHEME_TCP) => Ok(Self::Regular {
                addr: server_ip_addr.ok_or(ParseError::InvalidIpAddr)?,
                port: server_port.unwrap_or(DEFAULT_PORT_DNS),
                prefer_tcp: r == SCHEME_TCP,
            }),
            SCHEME_TLS => Ok(Self::TLS {
                host: server_ip_addr.map_or_else(
                    || Host::ServerName {
                        name: authority.host().to_string(),
                    },
                    |addr| Host::IpAddr {
                        addr,
                        custom_server_name: custom_sni_from_query(&uri).map(ToString::to_string),
                    },
                ),
                port: server_port.unwrap_or(DEFAULT_PORT_DNS_OVER_TLS),
            }),
            r @ (SCHEME_HTTPS | SCHEME_H3) => Ok(Self::HTTPS {
                host: server_ip_addr.map_or_else(
                    || Host::ServerName {
                        name: authority.host().to_string(),
                    },
                    |addr| Host::IpAddr {
                        addr,
                        custom_server_name: custom_sni_from_query(&uri).map(ToString::to_string),
                    },
                ),
                port: server_port.unwrap_or(DEFAULT_PORT_DNS_OVER_HTTPS),
                custom_http_endpoint: (!uri.path().is_empty()).then_some(uri.path().to_string()),
                force_http3: r == SCHEME_H3,
            }),
            SCHEME_QUIC => Ok(Self::QUIC {
                host: server_ip_addr.map_or_else(
                    || Host::ServerName {
                        name: authority.host().to_string(),
                    },
                    |addr| Host::IpAddr {
                        addr,
                        custom_server_name: custom_sni_from_query(&uri).map(ToString::to_string),
                    },
                ),
                port: server_port.unwrap_or(DEFAULT_PORT_DNS_OVER_QUIC),
            }),
            _ => Err(ParseError::UnsupportedScheme),
        }
    }
}

#[inline(always)]
fn custom_sni_from_query<'a>(uri: &fluent_uri::Uri<&'a str>) -> Option<&'a str> {
    uri.query()
        .and_then(|query| query.split('&').find_map(|query| query.as_str().strip_prefix("sni=")))
}

#[derive(Debug)]
/// Error parsing a DNS server address.
pub enum ParseError {
    /// The URI is invalid.
    InvalidUri(fluent_uri::error::ParseError),

    /// The host is missing.
    MissingHost,

    /// The IP address is invalid or missing.
    InvalidIpAddr,

    /// The port is invalid.
    InvalidPort,

    /// The scheme is unsupported.
    UnsupportedScheme,
}

impl fmt::Display for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidUri(err) => write!(f, "Invalid URI: {err}"),
            Self::MissingHost => write!(f, "Missing host"),
            Self::InvalidIpAddr => write!(f, "Invalid or missing IP address"),
            Self::InvalidPort => write!(f, "Invalid port"),
            Self::UnsupportedScheme => write!(f, "Unsupported scheme"),
        }
    }
}

impl From<fluent_uri::error::ParseError> for ParseError {
    fn from(err: fluent_uri::error::ParseError) -> Self {
        Self::InvalidUri(err)
    }
}

impl From<AddrParseError> for ParseError {
    fn from(_: AddrParseError) -> Self {
        Self::InvalidIpAddr
    }
}

impl From<ParseIntError> for ParseError {
    fn from(_: ParseIntError) -> Self {
        Self::InvalidPort
    }
}