1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

//! The main server which combines the DNS and HTTP(S) servers.

use anyhow::Result;
use iroh_metrics::metrics::start_metrics_server;
use tracing::info;

use crate::{
    config::Config,
    dns::{DnsHandler, DnsServer},
    http::HttpServer,
    state::AppState,
    store::ZoneStore,
};

/// Spawn the server and run until the `Ctrl-C` signal is received, then shutdown.
pub async fn run_with_config_until_ctrl_c(config: Config) -> Result<()> {
    let mut store = ZoneStore::persistent(Config::signed_packet_store_path()?)?;
    if let Some(bootstrap) = config.mainline_enabled() {
        info!("mainline fallback enabled");
        store = store.with_mainline_fallback(bootstrap);
    };
    let server = Server::spawn(config, store).await?;
    tokio::signal::ctrl_c().await?;
    info!("shutdown");
    server.shutdown().await?;
    Ok(())
}

/// The iroh-dns server.
pub struct Server {
    http_server: HttpServer,
    dns_server: DnsServer,
    metrics_task: tokio::task::JoinHandle<anyhow::Result<()>>,
}

impl Server {
    /// Spawn the server.
    ///
    /// This will spawn several background tasks:
    /// * A DNS server task
    /// * A HTTP server task, if `config.http` is not empty
    /// * A HTTPS server task, if `config.https` is not empty
    pub async fn spawn(config: Config, store: ZoneStore) -> Result<Self> {
        let dns_handler = DnsHandler::new(store.clone(), &config.dns)?;

        let state = AppState { store, dns_handler };

        let metrics_addr = config.metrics_addr();
        let metrics_task = tokio::task::spawn(async move {
            if let Some(addr) = metrics_addr {
                start_metrics_server(addr).await?;
            }
            Ok(())
        });
        let http_server = HttpServer::spawn(
            config.http,
            config.https,
            config.pkarr_put_rate_limit,
            state.clone(),
        )
        .await?;
        let dns_server = DnsServer::spawn(config.dns, state.dns_handler.clone()).await?;
        Ok(Self {
            http_server,
            dns_server,
            metrics_task,
        })
    }

    /// Cancel the server tasks and wait for all tasks to complete.
    pub async fn shutdown(self) -> Result<()> {
        self.metrics_task.abort();
        let (res1, res2) = tokio::join!(self.dns_server.shutdown(), self.http_server.shutdown(),);
        res1?;
        res2?;
        Ok(())
    }

    /// Wait for all tasks to complete.
    ///
    /// This will run forever unless all tasks close with an error, or `Self::cancel` is called.
    pub async fn run_until_error(self) -> Result<()> {
        tokio::select! {
            res = self.dns_server.run_until_done() => res?,
            res = self.http_server.run_until_done() => res?,
        }
        self.metrics_task.abort();
        Ok(())
    }

    /// Spawn a server suitable for testing.
    ///
    /// This will run the DNS and HTTP servers, but not the HTTPS server.
    ///
    /// It returns the server handle, the [`SocketAddr`] of the DNS server and the [`Url`] of the
    /// HTTP server.
    #[cfg(test)]
    pub async fn spawn_for_tests() -> Result<(Self, std::net::SocketAddr, url::Url)> {
        Self::spawn_for_tests_with_mainline(None).await
    }

    /// Spawn a server suitable for testing, while optionally enabling mainline with custom
    /// bootstrap addresses.
    #[cfg(test)]
    pub async fn spawn_for_tests_with_mainline(
        mainline: Option<crate::config::BootstrapOption>,
    ) -> Result<(Self, std::net::SocketAddr, url::Url)> {
        use std::net::{IpAddr, Ipv4Addr};

        use crate::config::MetricsConfig;

        let mut config = Config::default();
        config.dns.port = 0;
        config.dns.bind_addr = Some(IpAddr::V4(Ipv4Addr::LOCALHOST));
        config.http.as_mut().unwrap().port = 0;
        config.http.as_mut().unwrap().bind_addr = Some(IpAddr::V4(Ipv4Addr::LOCALHOST));
        config.https = None;
        config.metrics = Some(MetricsConfig::disabled());

        let mut store = ZoneStore::in_memory()?;
        if let Some(bootstrap) = mainline {
            info!("mainline fallback enabled");
            store = store.with_mainline_fallback(bootstrap);
        }
        let server = Self::spawn(config, store).await?;
        let dns_addr = server.dns_server.local_addr();
        let http_addr = server.http_server.http_addr().expect("http is set");
        let http_url = format!("http://{http_addr}").parse()?;
        Ok((server, dns_addr, http_url))
    }
}