[internal-dns] Avoid 'picking ports' (#1233)

Fixes #1200

The internal-dns crate includes two servers: a "DNS server" and a "Dropshot server". This change refactors the launching of both to ensure that their ports are queryable, especially under tests conditions where they are now set to "port zero" (let the OS pick for us).

This PR also removes the dependency on the "port-picker" crate.

internal-dns/Cargo.toml

@@ -30,7 +30,6 @@ expectorate = "1.0.5"
 omicron-test-utils = { path = "../test-utils" }
 openapiv3 = "1.0"
 openapi-lint = { git = "https://github.com/oxidecomputer/openapi-lint", branch = "main" }
-portpicker = "0.1"
 serde_json = "1.0"
 subprocess = "0.2.9"
 trust-dns-resolver = "0.21"

internal-dns/src/bin/dns-server.rs

@@ -54,20 +54,19 @@ async fn main() -> Result<(), anyhow::Error> {
 
     let db = Arc::new(sled::open(&config.data.storage_path)?);
 
-    {
+    let _dns_server = {
         let db = db.clone();
         let log = log.clone();
         let dns_config = internal_dns::dns_server::Config {
             bind_address: dns_address.to_string(),
             zone: zone.to_string(),
         };
-        tokio::spawn(async move {
-            internal_dns::dns_server::run(log, db, dns_config).await
-        });
-    }
+        internal_dns::dns_server::run(log, db, dns_config).await?
+    };
 
-    let server = internal_dns::start_server(config, log, db).await?;
-    server
+    let dropshot_server =
+        internal_dns::start_dropshot_server(config, log, db).await?;
+    dropshot_server
         .await
         .map_err(|error_message| anyhow!("server exiting: {}", error_message))
 }

internal-dns/src/dns_server.rs

@@ -34,23 +34,43 @@ pub struct Config {
     pub zone: String,
 }
 
-pub async fn run(log: Logger, db: Arc<sled::Db>, config: Config) -> Result<()> {
+pub struct Server {
+    pub address: SocketAddr,
+    pub handle: tokio::task::JoinHandle<Result<()>>,
+}
+
+impl Drop for Server {
+    fn drop(&mut self) {
+        self.handle.abort()
+    }
+}
+
+pub async fn run(
+    log: Logger,
+    db: Arc<sled::Db>,
+    config: Config,
+) -> Result<Server> {
     let socket = Arc::new(UdpSocket::bind(config.bind_address).await?);
+    let address = socket.local_addr()?;
 
-    loop {
-        let mut buf = vec![0u8; 16384];
-        let (n, src) = socket.recv_from(&mut buf).await?;
-        buf.resize(n, 0);
+    let handle = tokio::task::spawn(async move {
+        loop {
+            let mut buf = vec![0u8; 16384];
+            let (n, src) = socket.recv_from(&mut buf).await?;
+            buf.resize(n, 0);
 
-        let socket = socket.clone();
-        let log = log.clone();
-        let db = db.clone();
-        let zone = config.zone.clone();
+            let socket = socket.clone();
+            let log = log.clone();
+            let db = db.clone();
+            let zone = config.zone.clone();
 
-        tokio::spawn(async move {
-            handle_req(log, db, socket, src, buf, zone).await
-        });
-    }
+            tokio::spawn(async move {
+                handle_req(log, db, socket, src, buf, zone).await
+            });
+        }
+    });
+
+    Ok(Server { address, handle })
 }
 
 async fn respond_nxdomain(

internal-dns/src/lib.rs

@@ -20,7 +20,7 @@ pub struct Config {
     pub data: dns_data::Config,
 }
 
-pub async fn start_server(
+pub async fn start_dropshot_server(
     config: Config,
     log: slog::Logger,
     db: Arc<sled::Db>,

internal-dns/tests/basic_test.rs

@@ -2,7 +2,7 @@
 // License, v. 2.0. If a copy of the MPL was not distributed with this
 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
 
-use std::net::{Ipv6Addr, SocketAddr, SocketAddrV6};
+use std::net::Ipv6Addr;
 use std::sync::Arc;
 
 use anyhow::{Context, Result};
@@ -280,13 +280,16 @@ pub async fn servfail() -> Result<(), anyhow::Error> {
 struct TestContext {
     client: Client,
     resolver: TokioAsyncResolver,
-    server: dropshot::HttpServer<Arc<internal_dns::dropshot_server::Context>>,
+    dns_server: internal_dns::dns_server::Server,
+    dropshot_server:
+        dropshot::HttpServer<Arc<internal_dns::dropshot_server::Context>>,
     tmp: tempdir::TempDir,
 }
 
 impl TestContext {
     async fn cleanup(self) {
-        self.server.close().await.expect("Failed to clean up server");
+        drop(self.dns_server);
+        self.dropshot_server.close().await.expect("Failed to clean up server");
         self.tmp.close().expect("Failed to clean up tmp directory");
     }
 }
@@ -295,7 +298,7 @@ async fn init_client_server(
     zone: String,
 ) -> Result<TestContext, anyhow::Error> {
     // initialize dns server config
-    let (tmp, config, dropshot_port, dns_port) = test_config()?;
+    let (tmp, config) = test_config()?;
     let log = config
         .log
         .to_logger("internal-dns")
@@ -305,17 +308,21 @@ async fn init_client_server(
     let db = Arc::new(sled::open(&config.data.storage_path)?);
     db.clear()?;
 
-    let client =
-        Client::new(&format!("http://[::1]:{}", dropshot_port), log.clone());
+    // launch a dns server
+    let dns_server = {
+        let db = db.clone();
+        let log = log.clone();
+        let dns_config = internal_dns::dns_server::Config {
+            bind_address: "[::1]:0".into(),
+            zone,
+        };
+
+        internal_dns::dns_server::run(log, db, dns_config).await?
+    };
 
     let mut rc = ResolverConfig::new();
     rc.add_name_server(NameServerConfig {
-        socket_addr: SocketAddr::V6(SocketAddrV6::new(
-            Ipv6Addr::LOCALHOST,
-            dns_port,
-            0,
-            0,
-        )),
+        socket_addr: dns_server.address,
         protocol: Protocol::Udp,
         tls_dns_name: None,
         trust_nx_responses: false,
@@ -325,33 +332,21 @@ async fn init_client_server(
     let resolver =
         TokioAsyncResolver::tokio(rc, ResolverOpts::default()).unwrap();
 
-    // launch a dns server
-    {
-        let db = db.clone();
-        let log = log.clone();
-        let dns_config = internal_dns::dns_server::Config {
-            bind_address: format!("[::1]:{}", dns_port),
-            zone,
-        };
-
-        tokio::spawn(async move {
-            internal_dns::dns_server::run(log, db, dns_config).await
-        });
-    }
-
     // launch a dropshot server
-    let server = internal_dns::start_server(config, log, db).await?;
+    let dropshot_server =
+        internal_dns::start_dropshot_server(config, log.clone(), db).await?;
 
     // wait for server to start
     tokio::time::sleep(tokio::time::Duration::from_millis(250)).await;
 
-    Ok(TestContext { client, resolver, server, tmp })
+    let client =
+        Client::new(&format!("http://{}", dropshot_server.local_addr()), log);
+
+    Ok(TestContext { client, resolver, dns_server, dropshot_server, tmp })
 }
 
 fn test_config(
-) -> Result<(tempdir::TempDir, internal_dns::Config, u16, u16), anyhow::Error> {
-    let dropshot_port = portpicker::pick_unused_port().expect("pick port");
-    let dns_port = portpicker::pick_unused_port().expect("pick port");
+) -> Result<(tempdir::TempDir, internal_dns::Config), anyhow::Error> {
     let tmp_dir = tempdir::TempDir::new("internal-dns-test")?;
     let mut storage_path = tmp_dir.path().to_path_buf();
     storage_path.push("test");
@@ -362,7 +357,7 @@ fn test_config(
             level: dropshot::ConfigLoggingLevel::Info,
         },
         dropshot: dropshot::ConfigDropshot {
-            bind_address: format!("[::1]:{}", dropshot_port).parse().unwrap(),
+            bind_address: format!("[::1]:0").parse().unwrap(),
             request_body_max_bytes: 1024,
             ..Default::default()
         },
@@ -372,5 +367,5 @@ fn test_config(
         },
     };
 
-    Ok((tmp_dir, config, dropshot_port, dns_port))
+    Ok((tmp_dir, config))
 }