Merge pull request #160 from DeterminateSystems/fh-fetch-command

Initial implementation of `fh fetch`

src/cli/cmd/apply/mod.rs

@@ -11,9 +11,14 @@ use std::{
 use clap::{Parser, Subcommand};
 use color_eyre::eyre::Context;
 use tempfile::{tempdir, TempDir};
-use tokio::io::AsyncWriteExt as _;
 
-use crate::cli::{cmd::nix_command, error::FhError};
+use crate::{
+    cli::{
+        cmd::{copy_closure, nix_command},
+        error::FhError,
+    },
+    shared::create_temp_netrc,
+};
 
 use self::{home_manager::HomeManager, nix_darwin::NixDarwin, nixos::NixOs};
 
@@ -122,63 +127,14 @@ impl CommandExecute for ApplySubcommand {
         match resolved_path.token {
             Some(token) => {
                 if self.use_scoped_token == TokenChoice::Always {
-                    let mut nix_args = vec![
-                        "copy".to_string(),
-                        "--option".to_string(),
-                        "narinfo-cache-negative-ttl".to_string(),
-                        "0".to_string(),
-                        "--from".to_string(),
-                        self.cache_addr.to_string(),
-                        resolved_path.store_path.clone(),
-                    ];
-
                     let dir = tempdir()?;
-                    let temp_netrc_path = dir.path().join("netrc");
-
-                    let mut f = tokio::fs::OpenOptions::new()
-                        .create(true)
-                        .truncate(true)
-                        .write(true)
-                        .mode(0o600)
-                        .open(&temp_netrc_path)
-                        .await?;
-
-                    let cache_netrc_contents = format!(
-                        "machine {} login flakehub password {}\n",
-                        self.cache_addr.host_str().expect("valid host"),
-                        token
-                    );
-                    f.write_all(cache_netrc_contents.as_bytes())
-                        .await
-                        .wrap_err("writing restricted netrc file")?;
+
+                    let temp_netrc_path =
+                        create_temp_netrc(dir.path(), &self.cache_addr, &token).await?;
 
                     let display = temp_netrc_path.display().to_string();
-                    nix_args.extend_from_slice(&["--netrc-file".to_string(), display]);
-
-                    // NOTE(cole-h): Theoretically, this could be garbage collected immediately after we
-                    // copy it. There's no good way to prevent this at this point in time because:
-                    //
-                    // 0. We want to be able to use the scoped token to talk to FlakeHub Cache, which we
-                    // do via `--netrc-file`, and we want to be able to run this on any user -- trusted
-                    // or otherwise
-                    //
-                    // 1. `nix copy` substitutes on the client, so `--netrc-file` works just fine (it
-                    // won't be sent to the daemon, which will say "no" if you're not a trusted user),
-                    // but it doesn't have a `--profile` or `--out-link` argument, so we can't GC
-                    // root it that way
-                    //
-                    // 2. `nix build --max-jobs 0` does have `--profile` and `--out-link`, but passing
-                    // `--netrc-file` will send it to the daemon which doesn't work if you're not a
-                    // trusted user
-                    //
-                    // 3. Manually making a symlink somewhere doesn't work because adding that symlink
-                    // to gcroots/auto requires root, stashing it in a process's environment is so ugly
-                    // I will not entertain it, and holding a handle to it requires it to exist in the
-                    // first place (so there's still a small window of time where it can be GC'd)
-                    //
-                    // This will be resolved when https://github.com/NixOS/nix/pull/11657 makes it into
-                    // a Nix release.
-                    nix_command(&nix_args, false)
+
+                    copy_closure(self.cache_addr.as_str(), &resolved_path.store_path, display)
                         .await
                         .wrap_err("failed to copy resolved store path with Nix")?;
 

src/cli/cmd/fetch.rs

@@ -0,0 +1,74 @@
+use std::process::ExitCode;
+
+use clap::Parser;
+use color_eyre::eyre;
+use color_eyre::Result;
+
+use crate::cli::cmd::copy_closure_with_gc_root;
+use crate::shared::create_temp_netrc;
+
+use super::{CommandExecute, FlakeHubClient};
+
+#[derive(Parser)]
+pub(crate) struct FetchSubcommand {
+    /// The FlakeHub flake reference to fetch.
+    /// References must be of this form: {org}/{flake}/{version_req}#{attr_path}
+    flake_ref: String,
+
+    /// Target link to use as a Nix garbage collector root
+    target: String,
+
+    #[clap(from_global)]
+    api_addr: url::Url,
+
+    #[clap(from_global)]
+    cache_addr: url::Url,
+
+    #[clap(from_global)]
+    frontend_addr: url::Url,
+}
+
+#[async_trait::async_trait]
+impl CommandExecute for FetchSubcommand {
+    #[tracing::instrument(skip_all)]
+    async fn execute(self) -> Result<ExitCode> {
+        let parsed = super::parse_flake_output_ref(&self.frontend_addr, &self.flake_ref)?;
+
+        let resolved_path = FlakeHubClient::resolve(
+            self.api_addr.as_str(),
+            &parsed,
+            /* use scoped token */ true,
+        )
+        .await?;
+
+        tracing::info!(
+            "Resolved {} to {}",
+            self.flake_ref,
+            resolved_path.store_path
+        );
+
+        let token = match resolved_path.token {
+            Some(token) => token,
+            None => eyre::bail!("Did not receive a scoped token from FlakeHub!"),
+        };
+
+        let dir = tempfile::tempdir()?;
+
+        let netrc_path = create_temp_netrc(dir.path(), &self.cache_addr, &token).await?;
+        let token_path = netrc_path.display().to_string();
+
+        copy_closure_with_gc_root(
+            self.cache_addr.as_str(),
+            &resolved_path.store_path,
+            token_path,
+            &self.target,
+        )
+        .await?;
+
+        tracing::info!("Copied {} to {}", resolved_path.store_path, self.target);
+
+        dir.close()?;
+
+        Ok(ExitCode::SUCCESS)
+    }
+}

src/cli/cmd/mod.rs

@@ -3,6 +3,7 @@ pub(crate) mod apply;
 pub(crate) mod completion;
 pub(crate) mod convert;
 pub(crate) mod eject;
+pub(crate) mod fetch;
 pub(crate) mod init;
 pub(crate) mod list;
 pub(crate) mod login;
@@ -12,7 +13,7 @@ pub(crate) mod status;
 
 use std::{fmt::Display, process::Stdio};
 
-use color_eyre::eyre::WrapErr;
+use color_eyre::eyre::{self, WrapErr};
 use once_cell::sync::Lazy;
 use reqwest::{
     header::{HeaderMap, HeaderValue, ACCEPT, AUTHORIZATION},
@@ -23,6 +24,7 @@ use tabled::settings::{
     style::{HorizontalLine, On, VerticalLineIter},
     Style,
 };
+use tokio::process::Command;
 use url::Url;
 
 use self::{
@@ -69,6 +71,7 @@ pub(crate) enum FhSubcommands {
     Completion(completion::CompletionSubcommand),
     Convert(convert::ConvertSubcommand),
     Eject(eject::EjectSubcommand),
+    Fetch(fetch::FetchSubcommand),
     Init(init::InitSubcommand),
     List(list::ListSubcommand),
     Login(login::LoginSubcommand),
@@ -494,6 +497,136 @@ fn validate_segment(s: &str) -> Result<(), FhError> {
     Ok(())
 }
 
+/// Copy a Nix closure from a given host into the store.
+pub async fn copy_closure(
+    cache_host: impl Into<String>,
+    store_path: impl Into<String>,
+    token_path: impl Into<String>,
+) -> color_eyre::Result<()> {
+    let args = vec![
+        "copy".into(),
+        "--option".into(),
+        "narinfo-cache-negative-ttl".into(),
+        "0".into(),
+        "--from".into(),
+        cache_host.into(),
+        store_path.into(),
+        "--netrc-file".into(),
+        token_path.into(),
+    ];
+
+    nix_command(&args, false)
+        .await
+        .wrap_err("Failed to copy resolved store path with Nix")?;
+
+    Ok(())
+}
+
+async fn copy_supports_out_link() -> color_eyre::Result<bool> {
+    const OUT_LINK_NOT_SUPPORTED: &[u8] = b"error: unrecognised flag '--out-link'";
+
+    // Not using nix_command() here because we need to read the stderr of the resulting command
+    let output = Command::new("nix")
+        .args(["copy", "--out-link"])
+        .output()
+        .await
+        .wrap_err("Could not run nix")?;
+
+    // Grab only the first line of output from nix since it's the one we care about (the problem it encountered)
+    let error_line = output.stderr.split(|&c| c == b'\n').next();
+    let error_line = match error_line {
+        Some(line) => line,
+        None => {
+            tracing::warn!("Could not determine if `nix copy` supports --out-link; falling back to manual links");
+            return Ok(false);
+        }
+    };
+
+    let supported = error_line != OUT_LINK_NOT_SUPPORTED;
+    tracing::debug!(supported, "Setting support for nix copy --out-link");
+
+    Ok(supported)
+}
+
+async fn copy_closure_with_out_link(
+    cache_host: impl Into<String>,
+    store_path: impl Into<String>,
+    token_path: impl Into<String>,
+    out_path: impl Into<String>,
+) -> color_eyre::Result<()> {
+    let args = vec![
+        "copy".into(),
+        "--option".into(),
+        "narinfo-cache-negative-ttl".into(),
+        "0".into(),
+        "--from".into(),
+        cache_host.into(),
+        store_path.into(),
+        "--out-link".into(),
+        out_path.into(),
+        "--netrc-file".into(),
+        token_path.into(),
+    ];
+
+    nix_command(&args, false)
+        .await
+        .wrap_err("Failed to copy resolved store path with Nix")
+}
+
+async fn copy_closure_with_realise(
+    cache_host: impl Into<String>,
+    store_path: impl Into<String>,
+    token_path: impl Into<String>,
+    out_path: impl Into<String>,
+) -> color_eyre::Result<()> {
+    let cache_host = cache_host.into();
+    let store_path = store_path.into();
+    let token_path = token_path.into();
+    let out_path = out_path.into();
+
+    // First, copy the closure down into the user's Nix store
+    copy_closure(cache_host, &store_path, &token_path).await?;
+
+    // Now we can use a plain `nix-store --realise` on it
+    let mut command = Command::new("nix-store");
+    let output = command
+        .arg("--realise")
+        .arg(&store_path)
+        .arg("--add-root")
+        .arg(&out_path)
+        .spawn()
+        .wrap_err("Failed to spawn nix-store command")?
+        .wait_with_output()
+        .await?;
+
+    eyre::ensure!(
+        output.status.success(),
+        "Could not use nix-store --realise to copy {store_path} to {out_path}; consider upgrading to Nix version 2.26 or greater which is immune to this problem"
+    );
+
+    Ok(())
+}
+
+/// Copy a Nix closure like [`copy_closure`], but with a GC root. The bool that
+/// is returned indicates if `nix copy --out-link` (supported with version 2.26)
+/// was used.
+pub async fn copy_closure_with_gc_root(
+    cache_host: impl Into<String>,
+    store_path: impl Into<String>,
+    token_path: impl Into<String>,
+    out_path: impl Into<String>,
+) -> color_eyre::Result<bool> {
+    let use_out_link = copy_supports_out_link().await?;
+
+    if use_out_link {
+        copy_closure_with_out_link(cache_host, store_path, token_path, out_path).await?;
+    } else {
+        copy_closure_with_realise(cache_host, store_path, token_path, out_path).await?;
+    }
+
+    Ok(use_out_link)
+}
+
 #[cfg(test)]
 mod tests {
     #[test]

src/main.rs

@@ -39,6 +39,7 @@ async fn main() -> color_eyre::Result<std::process::ExitCode> {
         FhSubcommands::Completion(completion) => completion.execute().await,
         FhSubcommands::Convert(convert) => convert.execute().await,
         FhSubcommands::Eject(eject) => eject.execute().await,
+        FhSubcommands::Fetch(fetch) => fetch.execute().await,
         FhSubcommands::Init(init) => init.execute().await,
         FhSubcommands::List(list) => list.execute().await,
         FhSubcommands::Login(login) => login.execute().await,

src/shared/mod.rs

@@ -1,7 +1,10 @@
-use std::path::Path;
+use std::path::{Path, PathBuf};
 
 use color_eyre::eyre::Context as _;
 use serde::{Deserialize, Serialize};
+use tokio::fs::OpenOptions;
+use tokio::io::AsyncWriteExt;
+use url::Url;
 
 #[derive(Deserialize, Serialize)]
 pub struct DaemonInfoReponse {
@@ -181,3 +184,27 @@ pub fn merge_nix_configs(
         .unwrap_or(&new_config)
         .to_owned()
 }
+
+pub async fn create_temp_netrc(
+    dir: &Path,
+    host_url: &Url,
+    token: &str,
+) -> color_eyre::Result<PathBuf> {
+    let host = host_url.host_str().expect("Malformed URL: missing host");
+
+    let path = dir.join("netrc");
+
+    let mut file = OpenOptions::new()
+        .create(true)
+        .truncate(true)
+        .write(true)
+        .mode(0o600)
+        .open(&path)
+        .await?;
+
+    let contents = format!("machine {host} login flakehub password {token}\n");
+
+    file.write_all(contents.as_bytes()).await?;
+
+    Ok(path)
+}