From 9796ebd7ef8c9a23ee8128273d925acae00e43b0 Mon Sep 17 00:00:00 2001
From: Matej Urbas <matej.urbas@gmail.com>
Date: Mon, 17 Apr 2023 12:08:42 -0400
Subject: [PATCH] Add `--process-ops` flag to `nix-daemon`

---
 src/nix/daemon.cc | 33 +++++++++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/nix/daemon.cc b/src/nix/daemon.cc
index 373dedf7c..4dada8e0e 100644
--- a/src/nix/daemon.cc
+++ b/src/nix/daemon.cc
@@ -443,16 +443,23 @@ static void processStdioConnection(ref<Store> store, TrustedFlag trustClient)
  *
  * @param forceTrustClientOpt See `daemonLoop()` and the parameter with
  * the same name over there for details.
+ *
+ * @param procesOps Whether to force processing ops even if the next
+ * store also is a remote store and could process it directly.
  */
-static void runDaemon(bool stdio, std::optional<TrustedFlag> forceTrustClientOpt)
+static void runDaemon(bool stdio, std::optional<TrustedFlag> forceTrustClientOpt, bool processOps)
 {
     if (stdio) {
         auto store = openUncachedStore();
 
+        std::shared_ptr<RemoteStore> remoteStore;
+
         // If --force-untrusted is passed, we cannot forward the connection and
         // must process it ourselves (before delegating to the next store) to
         // force untrusting the client.
-        if (auto remoteStore = store.dynamic_pointer_cast<RemoteStore>(); remoteStore && (!forceTrustClientOpt || *forceTrustClientOpt != NotTrusted))
+        processOps |= !forceTrustClientOpt || *forceTrustClientOpt != NotTrusted;
+
+        if (!processOps && (remoteStore = store.dynamic_pointer_cast<RemoteStore>()))
             forwardStdioConnection(*remoteStore);
         else
             // `Trusted` is passed in the auto (no override case) because we
@@ -468,6 +475,7 @@ static int main_nix_daemon(int argc, char * * argv)
     {
         auto stdio = false;
         std::optional<TrustedFlag> isTrustedOpt = std::nullopt;
+        auto processOps = false;
 
         parseCmdLine(argc, argv, [&](Strings::iterator & arg, const Strings::iterator & end) {
             if (*arg == "--daemon")
@@ -487,11 +495,14 @@ static int main_nix_daemon(int argc, char * * argv)
             } else if (*arg == "--default-trust") {
                 experimentalFeatureSettings.require(Xp::DaemonTrustOverride);
                 isTrustedOpt = std::nullopt;
+            } else if (*arg == "--process-ops") {
+                experimentalFeatureSettings.require(Xp::MountedSSHStore);
+                processOps = true;
             } else return false;
             return true;
         });
 
-        runDaemon(stdio, isTrustedOpt);
+        runDaemon(stdio, isTrustedOpt, processOps);
 
         return 0;
     }
@@ -503,6 +514,7 @@ struct CmdDaemon : StoreCommand
 {
     bool stdio = false;
     std::optional<TrustedFlag> isTrustedOpt = std::nullopt;
+    bool processOps = false;
 
     CmdDaemon()
     {
@@ -538,6 +550,19 @@ struct CmdDaemon : StoreCommand
             }},
             .experimentalFeature = Xp::DaemonTrustOverride,
         });
+
+        addFlag({
+            .longName = "process-ops",
+            .description = R"(
+              Forces the daemon to process received commands itself rather than forwarding the commands straight to the remote store.
+
+              This is useful for the `mounted-ssh://` store where some actions need to be performed on the remote end but as connected user, and not as the user of the underlying daemon on the remote end.
+            )",
+            .handler = {[&]() {
+                processOps = true;
+            }},
+            .experimentalFeature = Xp::MountedSSHStore,
+        });
     }
 
     std::string description() override
@@ -556,7 +581,7 @@ struct CmdDaemon : StoreCommand
 
     void run(ref<Store> store) override
     {
-        runDaemon(stdio, isTrustedOpt);
+        runDaemon(stdio, isTrustedOpt, processOps);
     }
 };