From 62d69b12da84909070b043221e248f364583ffa8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christina=20S=C3=B8rensen?= Date: Fri, 15 Mar 2024 23:50:59 +0100 Subject: [PATCH] fix: changes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Christina Sørensen --- Cargo.lock | 11 ++++ Cargo.toml | 2 + flake.nix | 2 +- src/cli.rs | 39 ++++++-------- src/main.rs | 153 +++++++++++++++++----------------------------------- src/net.rs | 26 +++++++++ src/nix.rs | 53 ++++++++++++++++++ 7 files changed, 160 insertions(+), 126 deletions(-) create mode 100644 src/net.rs create mode 100644 src/nix.rs diff --git a/Cargo.lock b/Cargo.lock index 0eaa9e2..f9a5c6e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -540,6 +540,16 @@ dependencies = [ "byteorder", ] +[[package]] +name = "gethostname" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0176e0459c2e4a1fe232f984bca6890e681076abb9934f6cea7c326f3fc47818" +dependencies = [ + "libc", + "windows-targets 0.48.5", +] + [[package]] name = "getopts" version = "0.2.21" @@ -870,6 +880,7 @@ dependencies = [ "dns-lookup", "domain", "futures", + "gethostname", "itertools", "log", "openssl", diff --git a/Cargo.toml b/Cargo.toml index 8de4931..c3a373e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,7 @@ clap = { version = "4.5.1", features = ["cargo"] } dns-lookup = "2.0.4" domain = { version = "0.9.3", features = ["tokio", "resolv"] } futures = "0.3.30" +gethostname = "0.4.3" itertools = "0.12.1" log = "0.4.21" openssl = { version = "0.10.63" } @@ -36,3 +37,4 @@ tokio = { version = "1.36.0", features = ["macros", "full"] } clap = { version = "4.5.1", features = ["cargo"] } clap_complete = "4" clap_mangen = "0.2.20" +gethostname = "0.4.3" diff --git a/flake.nix b/flake.nix index 8885d2a..a81351d 100644 --- a/flake.nix +++ b/flake.nix @@ -121,7 +121,7 @@ # For `nix develop`: devShells.default = pkgs.mkShell { inherit (self.checks.${system}.pre-commit-check) shellHook; - inputsFrom = [ self.packages.${system}.default]; + inputsFrom = [self.packages.${system}.default]; nativeBuildInputs = with pkgs; [rustup toolchain just zip reuse]; }; diff --git a/src/cli.rs b/src/cli.rs index 192c504..4dcaa45 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -3,36 +3,31 @@ // // SPDX-License-Identifier: AGPL-3.0-only -use clap::{arg, command, crate_authors, Arg, Command}; +use std::{cell::OnceCell, sync::OnceLock}; + +use clap::{arg, command, crate_authors, value_parser, Arg, Command}; + +const DEFAULT_CACHE: &str = "cache.nixos.org"; pub fn build_cli() -> Command { + use std::path::PathBuf; + command!() .author(crate_authors!("\n")) .arg( - Arg::new("all") - .short('a') - .long("all") - .help("Shows all fortunes, including unkind."), + arg!(--cache "check a specific cache") + .required(false) + .default_value(DEFAULT_CACHE), ) .arg( - Arg::new("unkind") - .short('o') - .short('u') - .long("unkind") - .help("Shows only unkind fortunes."), + arg!(-n --name "Hostname of machine.") + .required(false) + .value_parser(value_parser!(String)), ) .arg( - Arg::new("find") - .short('m') - .long("find") - .value_name("pattern") - .help("Finds fortunes matching regex query."), + arg!(-c --config "Path to NixOS config.") + .required(false) + .value_parser(value_parser!(PathBuf)), ) - .arg( - Arg::new("length") - .short('n') - .long("length") - .help("Finds a fortune that is shorter than provided number."), - ) - .arg(arg!(-s --short ... "Shows a short aporism.")) + .arg(arg!(-v --verbose ... "Verbosity level.")) } diff --git a/src/main.rs b/src/main.rs index acf1486..1652213 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,15 +3,12 @@ // // SPDX-License-Identifier: AGPL-3.0-only -use std::{ - io, - net::{IpAddr, SocketAddr}, -}; +use std::{io, net::SocketAddr, env, sync::OnceLock}; use dns_lookup::lookup_host; -use futures::{future::join_all, stream, StreamExt}; +use futures::future::join_all; use itertools::Itertools; -use rayon::prelude::*; +use gethostname::gethostname; #[allow(unused)] use log::{debug, error, info, trace, warn}; @@ -19,93 +16,16 @@ use log::{debug, error, info, trace, warn}; use crate::nix::get_requisites; mod cli; +mod net; +mod nix; -mod nix { +/// The initial time to wait on http 104, in milliseconds +const SLIDE: u64 = 100; - use serde_json::{Result, Value}; - use std::{ - path::Path, - process::{Command, Stdio}, - str::Lines, - }; +const DEFAULT_CACHE: &str = "cache.nixos.org"; - pub fn get_requisites(host: &str) -> String { - let get_drv_path = Command::new("nix") - .current_dir(Path::new("/home/ces/org/src/git/afk-nixos")) - .env("NIXPKGS_ALLOW_INSECURE", "1") - .args([ - "build", - "--impure", - "--quiet", - &format!( - "./#nixosConfigurations.{}.config.system.build.toplevel", - host - ), - "--dry-run", - "--json", - "--option", - "eval-cache", - "true", - ]) - .output() - .unwrap(); - - let drv_path_json: Value = - serde_json::from_str(&String::from_utf8(get_drv_path.stdout).unwrap()).unwrap(); - let drv_path = drv_path_json[0]["drvPath"].clone(); - - println!("drv_path: {}", &drv_path); - - let get_drv_requisites = Command::new("nix-store") - .args(["--query", "--requisites", drv_path.as_str().unwrap()]) - .stdout(Stdio::piped()) - .spawn() - .unwrap(); - let drv_requisites_remove_base = Command::new("cut") - .args(["-d", "/", "-f4"]) - .stdin(Stdio::from(get_drv_requisites.stdout.unwrap())) - .stdout(Stdio::piped()) - .spawn() - .unwrap(); - let drv_requisites_to_hash = Command::new("cut") - .args(["-d", "-", "-f1"]) - .stdin(Stdio::from(drv_requisites_remove_base.stdout.unwrap())) - .stdout(Stdio::piped()) - .spawn() - .unwrap(); - - String::from_utf8(drv_requisites_to_hash.wait_with_output().unwrap().stdout).unwrap() - } -} - -mod net { - use std::{net::SocketAddr, time::Duration}; - - use async_recursion::async_recursion; - use reqwest::{Client, ClientBuilder, StatusCode}; - use tokio::time::sleep; - - #[async_recursion] - pub async fn nar_exists(client: Client, domain: &str, hash: &str, slide: u64) -> usize { - let response = client - .head(format!("https://{domain}/{hash}.narinfo")) - .send() - .await; - - match response { - Ok(response) if response.status().as_u16() == 200 => 1, - Ok(response) if response.status().as_u16() == 404 => 0, - _ => { - // We're so fast now we get rate limited. - // - // Writng an actual sliding window seems kinda hard, - // so we do this instead. - sleep(Duration::from_millis(slide)).await; - nar_exists(client, domain, hash, slide * 2).await - } - } - } -} +const HOST_NAME: OnceLock = OnceLock::new(); +const CACHE_URL: OnceLock = OnceLock::new(); #[tokio::main(flavor = "multi_thread")] async fn main() -> io::Result<()> { @@ -113,35 +33,62 @@ async fn main() -> io::Result<()> { let matches = cli::build_cli().get_matches(); - let domain = "cache.nixos.org"; - let ips: Vec = lookup_host(domain).unwrap(); + match matches + .get_one::("verbose") + .expect("Count's are defaulted") + { + 0 => env::set_var("RUST_LOG", "error"), + 1 => env::set_var("RUST_LOG", "warn"), + 2 => env::set_var("RUST_LOG", "info"), + 3 => env::set_var("RUST_LOG", "debug"), + 4 => env::set_var("RUST_LOG", "trace"), + _ => { + trace!("More than four -v flags don't increase log level."); + env::set_var("RUST_LOG", "trace") + } + } + + if let Some(name) = matches.get_one::("name") { + HOST_NAME.get_or_init(|| name.to_owned()); + } + else { + HOST_NAME.get_or_init(|| gethostname().into_string().unwrap()); + } + + if let Some(cache) = matches.get_one::("cache") { + trace!("Got cache argument: {cache}"); + CACHE_URL.get_or_init(|| cache.to_owned()); + } + else { + trace!("No cache argument, using default: {}", DEFAULT_CACHE.to_string()); + CACHE_URL.get_or_init(|| DEFAULT_CACHE.to_string()); + } + + let domain = CACHE_URL.get().unwrap().to_owned(); + let ips: Vec = lookup_host(&domain).unwrap(); debug!("{:#?}", &ips); let domain_addr = SocketAddr::new(ips[0], 443); let client = reqwest::Client::builder() - .resolve(domain, domain_addr) + .resolve(&domain, domain_addr) .build() .unwrap(); - let binding = get_requisites("DBCAC"); - let connection_buffer = binding + let binding = get_requisites(HOST_NAME.get().unwrap()); + + let tasks = binding .lines() .map(|line| line.to_owned()) - .collect::>(); - - // FIXME make constant - let slide = 100; - - // FIXME we take ten just for testing - let tasks = connection_buffer + .collect::>() .into_iter() .map(|hash| { let client = client.clone(); + let domain = domain.clone(); tokio::spawn(async move { info!("connecting to {domain} {domain_addr:#?} for {hash}"); - net::nar_exists(client, domain, &hash, slide).await + net::nar_exists(client, &domain, &hash, SLIDE).await }) }) .collect_vec(); diff --git a/src/net.rs b/src/net.rs new file mode 100644 index 0000000..7364a15 --- /dev/null +++ b/src/net.rs @@ -0,0 +1,26 @@ +use std::time::Duration; + +use async_recursion::async_recursion; +use reqwest::Client; +use tokio::time::sleep; + +#[async_recursion] +pub async fn nar_exists(client: Client, domain: &str, hash: &str, slide: u64) -> usize { + let response = client + .head(format!("https://{domain}/{hash}.narinfo")) + .send() + .await; + + match response { + Ok(response) if response.status().as_u16() == 200 => 1, + Ok(response) if response.status().as_u16() == 404 => 0, + _ => { + // We're so fast now we get rate limited. + // + // Writng an actual sliding window seems kinda hard, + // so we do this instead. + sleep(Duration::from_millis(slide)).await; + nar_exists(client, domain, hash, slide * 2).await + } + } +} diff --git a/src/nix.rs b/src/nix.rs new file mode 100644 index 0000000..e25186c --- /dev/null +++ b/src/nix.rs @@ -0,0 +1,53 @@ +use serde_json::Value; +use std::{ + path::Path, + process::{Command, Stdio}, +}; + +pub fn get_requisites(host: &str) -> String { + let get_drv_path = Command::new("nix") + .current_dir(Path::new("/home/ces/org/src/git/afk-nixos")) + .env("NIXPKGS_ALLOW_INSECURE", "1") + .args([ + "build", + "--impure", + "--quiet", + &format!( + "./#nixosConfigurations.{}.config.system.build.toplevel", + host + ), + "--dry-run", + "--json", + "--option", + "eval-cache", + "true", + ]) + .output() + .unwrap(); + + let drv_path_json: Value = + serde_json::from_str(&String::from_utf8(get_drv_path.stdout).unwrap()).unwrap(); + let drv_path = drv_path_json[0]["drvPath"].clone(); + + println!("drv_path: {}", &drv_path); + + let get_drv_requisites = Command::new("nix-store") + .args(["--query", "--requisites", drv_path.as_str().unwrap()]) + .stdout(Stdio::piped()) + .spawn() + .unwrap(); + let drv_requisites_remove_base = Command::new("cut") + .args(["-d", "/", "-f4"]) + .stdin(Stdio::from(get_drv_requisites.stdout.unwrap())) + .stdout(Stdio::piped()) + .spawn() + .unwrap(); + let drv_requisites_to_hash = Command::new("cut") + .args(["-d", "-", "-f1"]) + .stdin(Stdio::from(drv_requisites_remove_base.stdout.unwrap())) + .stdout(Stdio::piped()) + .spawn() + .unwrap(); + + String::from_utf8(drv_requisites_to_hash.wait_with_output().unwrap().stdout).unwrap() +}