fh apply: support requesting restricted tokens

src/cli/cmd/apply/mod.rs

@@ -11,6 +11,7 @@ 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};
 
@@ -24,11 +25,18 @@ pub(crate) struct ApplySubcommand {
     #[clap(subcommand)]
     system: System,
 
+    /// Use a scoped token generated by FlakeHub that allows substituting the given output _only_.
+    #[clap(long, default_value_t = true)]
+    use_scoped_token: bool,
+
     #[clap(from_global)]
     api_addr: url::Url,
 
     #[clap(from_global)]
     frontend_addr: url::Url,
+
+    #[clap(from_global)]
+    cache_addr: url::Url,
 }
 
 #[derive(Subcommand)]
@@ -77,16 +85,96 @@ impl CommandExecute for ApplySubcommand {
 
         tracing::info!("Resolving {}", output_ref);
 
-        let resolved_path = FlakeHubClient::resolve(self.api_addr.as_ref(), &output_ref).await?;
+        let resolved_path =
+            FlakeHubClient::resolve(self.api_addr.as_ref(), &output_ref, self.use_scoped_token)
+                .await?;
 
         tracing::debug!(
             "Successfully resolved reference {} to path {}",
             &output_ref,
             &resolved_path.store_path
         );
 
+        let profile_path = applyer.profile_path();
+
+        match resolved_path.token {
+            Some(token) => {
+                if self.use_scoped_token {
+                    let mut nix_args = vec![
+                        "copy".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 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)
+                        .await
+                        .wrap_err("failed to copy resolved store path with Nix")?;
+
+                    dir.close()?;
+                } else {
+                    tracing::warn!(
+                        "Received a scoped token from FlakeHub, but we didn't request one! Ignoring."
+                    );
+                }
+            }
+            None => {
+                if self.use_scoped_token {
+                    return Err(color_eyre::eyre::eyre!(
+                        "FlakeHub did not return a restricted token!"
+                    ));
+                }
+            }
+        }
+
         let (profile_path, _tempdir) = apply_path_to_profile(
-            applyer.profile_path(),
+            profile_path,
             &resolved_path.store_path,
             applyer.requires_root(),
         )
@@ -212,18 +300,21 @@ async fn apply_path_to_profile(
 
     nix_command(
         &[
-            "build",
+            "build".to_string(),
             // Don't create a result symlink in the current directory for the profile being installed.
             // This is verified to not introduce a race condition against an eager garbage collection.
-            "--no-link",
-            "--print-build-logs",
+            "--no-link".to_string(),
+            "--print-build-logs".to_string(),
             // `--max-jobs 0` ensures that `nix build` doesn't really *build* anything
             // and acts more as a fetch operation
-            "--max-jobs",
-            "0",
-            "--profile",
-            profile_path.to_str().ok_or(FhError::InvalidProfile)?,
-            store_path,
+            "--max-jobs".to_string(),
+            "0".to_string(),
+            "--profile".to_string(),
+            profile_path
+                .to_str()
+                .ok_or(FhError::InvalidProfile)?
+                .to_string(),
+            store_path.to_string(),
         ],
         sudo_if_necessary,
     )

src/cli/cmd/convert.rs

@@ -85,7 +85,7 @@ impl CommandExecute for ConvertSubcommand {
 
             tracing::debug!("Running: nix flake lock");
 
-            nix_command(&["flake", "lock"], false)
+            nix_command(&["flake".to_string(), "lock".to_string()], false)
                 .await
                 .wrap_err("failed to create missing lock file entries")?;
         }

src/cli/cmd/mod.rs

@@ -157,15 +157,19 @@ impl FlakeHubClient {
         Ok(res)
     }
 
-    async fn resolve(api_addr: &str, output_ref: &FlakeOutputRef) -> Result<ResolvedPath, FhError> {
+    async fn resolve(
+        api_addr: &str,
+        output_ref: &FlakeOutputRef,
+        include_token: bool,
+    ) -> Result<ResolvedPath, FhError> {
         let FlakeOutputRef {
             ref org,
             project: ref flake,
             ref version_constraint,
             ref attr_path,
         } = output_ref;
 
-        let url = flakehub_url!(
+        let mut url = flakehub_url!(
             api_addr,
             "f",
             org,
@@ -175,6 +179,10 @@ impl FlakeHubClient {
             attr_path
         );
 
+        if include_token {
+            url.set_query(Some("include_token=true"));
+        }
+
         get(url, true).await
     }
 
@@ -386,7 +394,7 @@ fn is_root_user() -> bool {
     nix::unistd::getuid().is_root()
 }
 
-async fn nix_command(args: &[&str], sudo_if_necessary: bool) -> Result<(), FhError> {
+async fn nix_command(args: &[String], sudo_if_necessary: bool) -> Result<(), FhError> {
     command_exists("nix")?;
 
     let use_sudo = sudo_if_necessary && !is_root_user();

src/cli/cmd/resolve.rs

@@ -29,6 +29,8 @@ pub(crate) struct ResolvedPath {
     attribute_path: String,
     // The resolved store path
     pub(crate) store_path: String,
+    // A JWT that can only substitute the closure of this store path
+    pub(crate) token: Option<String>,
 }
 
 #[async_trait::async_trait]
@@ -37,7 +39,8 @@ impl CommandExecute for ResolveSubcommand {
     async fn execute(self) -> color_eyre::Result<ExitCode> {
         let output_ref = parse_flake_output_ref(&self.frontend_addr, &self.flake_ref)?;
 
-        let resolved_path = FlakeHubClient::resolve(self.api_addr.as_ref(), &output_ref).await?;
+        let resolved_path =
+            FlakeHubClient::resolve(self.api_addr.as_ref(), &output_ref, false).await?;
 
         tracing::debug!(
             "Successfully resolved reference {} to path {}",