feat(args): ipv4 and ipv6 flags

Allow forcing use of either IPv6 or IPv4 via -4/-6 CLI flags.

Fixes: #14.
Signed-off-by: Maximilian Marx <[email protected]>
Signed-off-by: Christina Sørensen <[email protected]>

crates/nix-weather/src/cli.rs

@@ -1,9 +1,10 @@
 // SPDX-FileCopyrightText: 2023-2024 Christina Sørensen
 // SPDX-FileContributor: Christina Sørensen
+// SPDX-FileContributor: Maximilian Marx
 //
 // SPDX-License-Identifier: EUPL-1.2
 
-use clap::{arg, command, crate_authors, value_parser, ArgAction, Command};
+use clap::{arg, command, crate_authors, value_parser, ArgAction, ArgGroup, Command};
 
 const DEFAULT_CACHE: &str = "cache.nixos.org";
 
@@ -38,4 +39,11 @@ pub fn build_cli() -> Command {
         .long("print-build-logs")
         .conflicts_with("verbose"),
     )
+    .arg(arg!(-'4' --"only-ipv4" "Use IPv4 addresses only.").action(ArgAction::SetTrue))
+    .arg(arg!(-'6' --"only-ipv6" "Use IPv6 addresses only.").action(ArgAction::SetTrue))
+    .group(
+      ArgGroup::new("address_family")
+        .args(["only-ipv4", "only-ipv6"])
+        .required(false),
+    )
 }

crates/nix-weather/src/main.rs

@@ -1,15 +1,17 @@
 // SPDX-FileCopyrightText: 2024 Christina Sørensen
 // SPDX-FileContributor: Christina Sørensen
+// SPDX-FileContributor: Maximilian Marx
 //
 // SPDX-License-Identifier: EUPL-1.2
 
+use std::sync::Arc;
 use std::time::Instant;
 use std::{env, io, net::SocketAddr};
 
-use dns_lookup::lookup_host;
 use futures::future::join_all;
 use gethostname::gethostname;
 use itertools::Itertools;
+use net::AddressFamilyFilter;
 
 use crate::nix::get_requisites;
 
@@ -91,14 +93,26 @@ async fn main() -> io::Result<()> {
     config_dir = DEFAULT_CONFIG_DIR.to_string();
   }
 
+  let address_family_filter = if matches.get_flag("only-ipv4") {
+    AddressFamilyFilter::OnlyIPv4
+  } else if matches.get_flag("only-ipv6") {
+    AddressFamilyFilter::OnlyIPv6
+  } else {
+    Default::default()
+  };
+
   let domain = cache_url.to_owned();
-  let ips: Vec<std::net::IpAddr> = lookup_host(&domain).unwrap();
+  let ips: Vec<std::net::IpAddr> = address_family_filter
+    .lookup_host(&domain)
+    .unwrap()
+    .collect();
 
   log::debug!("{:#?}", &ips);
 
   let domain_addr = SocketAddr::new(ips[0], 443);
 
   let client = reqwest::Client::builder()
+    .dns_resolver(Arc::new(address_family_filter))
     .resolve(&domain, domain_addr)
     .build()
     .unwrap();

crates/nix-weather/src/net.rs

@@ -1,15 +1,59 @@
 // SPDX-FileCopyrightText: 2024 Christina Sørensen
 // SPDX-FileContributor: Christina Sørensen
+// SPDX-FileContributor: Maximilian Marx
 //
 // SPDX-License-Identifier: EUPL-1.2
 
-use std::time::Duration;
+use std::{
+  io,
+  net::{IpAddr, SocketAddr},
+  time::Duration,
+};
 
-use reqwest::Client;
+use reqwest::{dns::Resolve, Client};
 use tokio::time::sleep;
 
 const MAX_SLIDE: u64 = 1000;
 
+#[derive(Debug, Copy, Clone, Default)]
+pub enum AddressFamilyFilter {
+  #[default]
+  Both,
+  OnlyIPv4,
+  OnlyIPv6,
+}
+
+impl AddressFamilyFilter {
+  pub fn lookup_host(self, host: &str) -> io::Result<impl Iterator<Item = IpAddr>> {
+    let addresses = dns_lookup::lookup_host(host)?;
+    Ok(self.filter_addresses(addresses))
+  }
+
+  fn filter_addresses<T>(self, addresses: T) -> impl Iterator<Item = IpAddr>
+  where
+    T: IntoIterator<Item = IpAddr>,
+  {
+    addresses.into_iter().filter(move |address| match self {
+      Self::Both => true,
+      Self::OnlyIPv4 => matches!(address, IpAddr::V4(_)),
+      Self::OnlyIPv6 => matches!(address, IpAddr::V6(_)),
+    })
+  }
+}
+
+impl Resolve for AddressFamilyFilter {
+  fn resolve(&self, name: reqwest::dns::Name) -> reqwest::dns::Resolving {
+    let filter = *self;
+    Box::pin(async move {
+      let addresses = filter.lookup_host(name.as_str())?;
+      let socket_addresses: Box<dyn Iterator<Item = SocketAddr> + Send> =
+        Box::new(addresses.map(|ip| SocketAddr::new(ip, 0)));
+
+      Ok(socket_addresses)
+    })
+  }
+}
+
 pub async fn nar_exists(client: Client, domain: &str, hash: &str, slide: u64) -> usize {
   let response = client
     .head(format!("https://{domain}/{hash}.narinfo"))