kapiti 0.0.3

use std::convert::TryFrom;
use std::fmt;
use std::iter::Iterator;
use std::net::SocketAddr;
use std::time::Duration;
use std::vec::Vec;

use anyhow::{bail, Context, Result};
use hyper::Uri;
use tracing::{self, trace};

use crate::cache;
use crate::client::{https, system, tcp, tls, udp, DnsClient};
use crate::resolver::Resolver;

static DEFAULT_UDP_TCP_PORT: u16 = 53;
static DEFAULT_HTTPS_PORT: u16 = 443;
static DEFAULT_TLS_PORT: u16 = 853;

#[derive(Clone, Debug, PartialEq)]
enum UpstreamType {
    /// System resolver
    System,
    /// UDP-only
    Udp,
    /// TCP-only
    Tcp,
    /// DoH
    Https,
    /// DoT
    Tls,
}

struct UpstreamInfo {
    uri: Uri,
    upstream_type: UpstreamType,
}

impl fmt::Debug for UpstreamInfo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.uri.to_string().as_str())
    }
}

impl UpstreamInfo {
    fn new(upstream_orig: &String) -> Result<(UpstreamInfo, Option<UpstreamInfo>)> {
        let upstream = upstream_orig.trim();
        let uri = Uri::try_from(upstream)
            .with_context(|| format!("Failed to parse upstream URI: {}", upstream))?;
        match uri.scheme_str() {
            None => {
                if let Some("system") = uri.host() {
                    // URI as 'system'
                    Ok((UpstreamInfo {
                        uri,
                        upstream_type: UpstreamType::System,
                    }, None))
                } else {
                    // URI as '<ip>': treat as 'udp+tcp'
                    Ok((
                        UpstreamInfo {
                            uri: uri.clone(),
                            upstream_type: UpstreamType::Udp,
                        },
                        Some(UpstreamInfo {
                            uri,
                            upstream_type: UpstreamType::Tcp,
                        })
                    ))
                }
            },
            Some(scheme) => match scheme {
                "udp+tcp" => Ok((
                    UpstreamInfo {
                        uri: uri.clone(),
                        upstream_type: UpstreamType::Udp,
                    },
                    Some(UpstreamInfo {
                        uri,
                        upstream_type: UpstreamType::Tcp,
                    })
                )),
                "udp" => Ok((UpstreamInfo {
                    uri,
                    upstream_type: UpstreamType::Udp,
                }, None)),
                "tcp" => Ok((UpstreamInfo {
                    uri,
                    upstream_type: UpstreamType::Tcp,
                }, None)),
                "https" => Ok((UpstreamInfo {
                    uri,
                    upstream_type: UpstreamType::Https,
                }, None)),
                "tls" => Ok((UpstreamInfo {
                    uri,
                    upstream_type: UpstreamType::Tls,
                }, None)),
                other => bail!(format!("Unsupported upstream URI scheme '{}', expected one of: 'system', '<ip>', 'udp+tcp://<ip>', 'udp://<ip>', 'tcp://<ip>', 'https://<ip/host>', or 'tls://<ip/host>'", other)),
            },
        }
    }
}

/// Convert the config-provided upstream strings into a list of prepared client instances.
pub fn parse_upstreams(
    cache_tx: async_channel::Sender<cache::task::CacheMsg>,
    config_upstreams: &Vec<String>,
    client_timeout: &Duration,
) -> Result<Resolver> {
    if config_upstreams.is_empty() {
        bail!("Upstreams list is empty, at least one upstream is required");
    }
    let mut upstreams = Vec::new();
    // Parse upstreams, check if URIs are provided as an IP or a hostname
    for config_upstream in config_upstreams {
        match UpstreamInfo::new(config_upstream)? {
            (upstream_info1, Some(upstream_info2)) => {
                trace!(
                    "Parsed upstreams: {:?} + {:?}",
                    upstream_info1,
                    upstream_info2
                );
                upstreams.push(upstream_info1);
                upstreams.push(upstream_info2);
            }
            (upstream_info, None) => {
                trace!("Parsed upstream: {:?}", upstream_info);
                upstreams.push(upstream_info);
            }
        }
    }

    let mut clients: Vec<Box<dyn DnsClient + Send + 'static>> = Vec::new();
    for upstream_info in &upstreams {
        clients.push(match upstream_info.upstream_type {
            UpstreamType::Https => {
                // Check if the DoH upstream is provided as a hostname or as an IP.
                if let Some(_upstream_addr) = to_addr(&upstream_info.uri, DEFAULT_HTTPS_PORT)? {
                    // It's an IP, so no "bootstrap" resolution should be necessary.
                    Box::new(https::Client::new_ip(upstream_info.uri.clone(), client_timeout.clone())?)
                } else {
                    // It's a hostname, so we need a separate "bootstrap" resolver to look that up.
                    let bootstrap_clients = to_bootstrap_clients(&upstreams, client_timeout)?;
                    if bootstrap_clients.is_empty() {
                        bail!(
                            r#"At least one upstream server must be specified as an IP.
For example, given a primary DoH upstream of 'https://example.com', add a secondary upstream of e.g. 'system' or '1.1.1.1' which can then be used to resolve 'example.com'.
Configured upstreams are: {:?}"#,
                            upstreams
                        );
                    }
                    let internal_resolver = Resolver::new(cache_tx.clone(), bootstrap_clients);
                    Box::new(https::Client::new_hostname(
                        upstream_info.uri.clone(),
                        internal_resolver,
                        client_timeout.clone(),
                    )?)
                }
            },
            UpstreamType::System | UpstreamType::Tcp | UpstreamType::Udp | UpstreamType::Tls => {
                // For these types, any endpoint is normally provided as an IP,
                // so no bootstrap lookup of the upstream host is needed (or supported).
                to_client(upstream_info, &client_timeout)?
            },
        });
    }
    Ok(Resolver::new(cache_tx, clients))
}

fn to_bootstrap_clients(
    upstreams: &Vec<UpstreamInfo>,
    client_timeout: &Duration,
) -> Result<Vec<Box<dyn DnsClient + Send>>> {
    let mut bootstrap_clients: Vec<Box<dyn DnsClient + Send + 'static>> = Vec::new();
    for upstream_info in upstreams {
        match upstream_info.upstream_type {
            UpstreamType::System | UpstreamType::Tcp | UpstreamType::Udp | UpstreamType::Tls => {
                bootstrap_clients.push(to_client(upstream_info, client_timeout)?);
            }
            UpstreamType::Https => {}
        }
    }
    Ok(bootstrap_clients)
}

fn to_client(
    upstream_info: &UpstreamInfo,
    client_timeout: &Duration,
) -> Result<Box<dyn DnsClient + Send>> {
    match upstream_info.upstream_type {
        UpstreamType::System => Ok(Box::new(system::Client::new())),
        UpstreamType::Udp => Ok(Box::new(udp::Client::new(
            to_addr(&upstream_info.uri, DEFAULT_UDP_TCP_PORT)?.with_context(|| {
                format!(
                    "UDP upstreams must be specified as an IP address: {:?}",
                    upstream_info.uri
                )
            })?,
            client_timeout.clone(),
        ))),
        UpstreamType::Tcp => Ok(Box::new(tcp::Client::new(
            to_addr(&upstream_info.uri, DEFAULT_UDP_TCP_PORT)?.with_context(|| {
                format!(
                    "TCP upstreams must be specified as an IP address: {:?}",
                    upstream_info.uri
                )
            })?,
            client_timeout.clone(),
        ))),
        UpstreamType::Tls => Ok(Box::new(tls::Client::new(
            to_addr(&upstream_info.uri, DEFAULT_TLS_PORT)?.with_context(|| {
                format!(
                    "TLS upstreams must be specified as an IP address: {:?}",
                    upstream_info.uri
                )
            })?,
            client_timeout.clone(),
        ))),
        UpstreamType::Https => bail!(
            "{:?} client requires a bootstrap resolver",
            upstream_info.upstream_type
        ),
    }
}

/// If the URI is provided as an IP, extract it, otherwise return `None`.
fn to_addr(uri: &Uri, default_port: u16) -> Result<Option<SocketAddr>> {
    let authority = uri
        .authority()
        .with_context(|| format!("Missing authority in upstream: {}", uri))?;
    let ip_addr = authority.host().parse();
    Ok(ip_addr.map_or(None, |ip| {
        Some(SocketAddr::new(
            ip,
            authority.port_u16().unwrap_or(default_port),
        ))
    }))
}