From 3e8ec98a264d0ae2e1d304e84eadc6c302b6e438 Mon Sep 17 00:00:00 2001
From: Zhaofeng Li <hello@zhaofeng.li>
Date: Sun, 18 Sep 2022 17:27:46 -0600
Subject: [PATCH] ssh: Support using the experimental `nix copy` command to
 copy

The new SSH store protocol (ssh-ng://) provides better performance.
This is now enabled for flake deployments.
---
 src/nix/hive/mod.rs |  5 +++
 src/nix/host/ssh.rs | 96 ++++++++++++++++++++++++++++++++++-----------
 2 files changed, 78 insertions(+), 23 deletions(-)

diff --git a/src/nix/hive/mod.rs b/src/nix/hive/mod.rs
index e5b9b43..26d0556 100644
--- a/src/nix/hive/mod.rs
+++ b/src/nix/hive/mod.rs
@@ -205,6 +205,11 @@ impl Hive {
                 if let Some(ssh_config) = &ssh_config {
                     host.set_ssh_config(ssh_config.clone());
                 }
+
+                if self.is_flake() {
+                    host.set_use_nix3_copy(true);
+                }
+
                 host.upcast()
             });
             let ssh_host = host.is_some();
diff --git a/src/nix/host/ssh.rs b/src/nix/host/ssh.rs
index 0dffbda..061946d 100644
--- a/src/nix/host/ssh.rs
+++ b/src/nix/host/ssh.rs
@@ -32,6 +32,9 @@ pub struct Ssh {
     /// Command to elevate privileges with.
     privilege_escalation_command: Vec<String>,
 
+    /// Whether to use the experimental `nix copy` command.
+    use_nix3_copy: bool,
+
     job: Option<JobHandle>,
 }
 
@@ -187,6 +190,7 @@ impl Ssh {
             port: None,
             ssh_config: None,
             privilege_escalation_command: Vec::new(),
+            use_nix3_copy: false,
             job: None,
         }
     }
@@ -203,6 +207,10 @@ impl Ssh {
         self.privilege_escalation_command = command;
     }
 
+    pub fn set_use_nix3_copy(&mut self, enable: bool) {
+        self.use_nix3_copy = enable;
+    }
+
     pub fn upcast(self) -> Box<dyn Host> {
         Box::new(self)
     }
@@ -252,31 +260,73 @@ impl Ssh {
         let ssh_options = self.ssh_options();
         let ssh_options_str = ssh_options.join(" ");
 
-        let mut command = Command::new("nix-copy-closure");
-        match direction {
-            CopyDirection::ToRemote => {
-                command.arg("--to");
-            }
-            CopyDirection::FromRemote => {
-                command.arg("--from");
-            }
-        }
+        let mut command = if self.use_nix3_copy {
+            // experimental `nix copy` command with ssh-ng://
+            let mut command = Command::new("nix");
 
-        // FIXME: Host-agnostic abstraction
-        if options.include_outputs {
-            command.arg("--include-outputs");
-        }
-        if options.use_substitutes {
-            command.arg("--use-substitutes");
-        }
-        if options.gzip {
-            command.arg("--gzip");
-        }
+            command.args(&[
+                "--extra-experimental-features",
+                "nix-command",
+                "copy",
+                "--no-check-sigs",
+            ]);
 
-        command
-            .arg(&self.ssh_target())
-            .arg(path.as_path())
-            .env("NIX_SSHOPTS", ssh_options_str);
+            if options.use_substitutes {
+                command.arg("--substitute-on-destination");
+            }
+
+            if path.ends_with(".drv") {
+                command.arg("--derivation");
+            }
+
+            match direction {
+                CopyDirection::ToRemote => {
+                    command.arg("--to");
+                }
+                CopyDirection::FromRemote => {
+                    command.arg("--from");
+                }
+            }
+
+            let mut store_uri = format!("ssh-ng://{}", self.ssh_target());
+            if options.gzip {
+                store_uri += "?compress=true";
+            }
+            command.arg(store_uri);
+
+            command.arg(path.as_path());
+
+            command
+        } else {
+            // nix-copy-closure (ssh://)
+            let mut command = Command::new("nix-copy-closure");
+
+            match direction {
+                CopyDirection::ToRemote => {
+                    command.arg("--to");
+                }
+                CopyDirection::FromRemote => {
+                    command.arg("--from");
+                }
+            }
+
+            // FIXME: Host-agnostic abstraction
+            if options.include_outputs {
+                command.arg("--include-outputs");
+            }
+            if options.use_substitutes {
+                command.arg("--use-substitutes");
+            }
+            if options.gzip {
+                command.arg("--gzip");
+            }
+
+            command.arg(&self.ssh_target()).arg(path.as_path());
+
+            command
+        };
+
+        command.env("NIX_SSHOPTS", ssh_options_str);
 
         command
     }