From 93122446d6001f9789d05e565f73bebfa3f53b50 Mon Sep 17 00:00:00 2001
From: Wojtek Czekalski <wczekalski@me.com>
Date: Thu, 31 Oct 2024 14:44:52 +0100
Subject: [PATCH] Add support for remote sudo

---
 nixos-modules/microvm/ssh-deploy.nix | 48 +++++++++++++++++++++++-----
 1 file changed, 40 insertions(+), 8 deletions(-)

diff --git a/nixos-modules/microvm/ssh-deploy.nix b/nixos-modules/microvm/ssh-deploy.nix
index 55980fa0..bccfd1ea 100644
--- a/nixos-modules/microvm/ssh-deploy.nix
+++ b/nixos-modules/microvm/ssh-deploy.nix
@@ -97,12 +97,25 @@ in
     installOnHost = pkgs.writeShellScriptBin "microvm-install-on-host" ''
       set -eou pipefail
 
+      USAGE="Usage: $0 root@<host> [--use-remote-sudo]"
+
       HOST="$1"
       if [[ -z "$HOST" ]]; then
-        echo "Usage: $0 root@<host>"
+        echo $USAGE
         exit 1
       fi
       shift
+      SSH_CMD="bash"
+      if [ $# -gt 0 ]; then
+        if [ "$1" == "--use-remote-sudo" ]; then
+          SSH_CMD="sudo bash"
+          shift
+        else
+          echo "$USAGE"
+          exit 1
+        fi
+      fi
+
 
       echo "Copying derivations to $HOST"
       nix copy --no-check-sigs --to "ssh-ng://$HOST" \
@@ -110,7 +123,7 @@ in
         "${paths.closureInfoDrv}^out" \
         "${paths.runnerDrv}^out"
 
-      ssh "$HOST" -- bash -e <<__SSH__
+      ssh "$HOST" -- $SSH_CMD -e <<__SSH__
       set -eou pipefail
 
       echo "Initializing MicroVM ${hostName} if necessary"
@@ -160,14 +173,26 @@ in
       pkgs.writeShellScriptBin "microvm-switch" ''
         set -eou pipefail
 
+        USAGE="Usage: $0 root@<target> [--use-remote-sudo]"
+
         TARGET="$1"
         if [[ -z "$TARGET" ]]; then
-          echo "Usage: $0 root@<target>"
+          echo "$USAGE"
           exit 1
         fi
         shift
+        SSH_CMD="bash"
+        if [ $# -gt 0 ]; then
+          if [ "$1" == "--use-remote-sudo" ]; then
+            SSH_CMD="sudo bash"
+            shift
+          else
+            echo "$USAGE"
+            exit 1
+          fi
+        fi
 
-        ssh "$TARGET" bash -e <<__SSH__
+        ssh "$TARGET" $SSH_CMD -e <<__SSH__
         set -eou pipefail
 
         hostname=\$(cat /etc/hostname)
@@ -193,14 +218,21 @@ in
       shift
       TARGET="$1"
       shift
-      if [[ -z "$HOST" || -z "$TARGET" ]]; then
-        echo "Usage: $0 root@<host> root@<target> switch"
+      OPTS="$@"
+      if [ $# -gt 0 ]; then
+        if [ "$1" == "--use-remote-sudo" ]; then
+          OPTS="$1"
+          shift
+        fi
+      fi
+      if [[ -z "$HOST" || -z "$TARGET" || $# -gt 0 ]]; then
+        echo "Usage: $0 root@<host> root@<target> [--use-remote-sudo] switch"
         exit 1
       fi
 
-      ${lib.getExe installOnHost} "$HOST"
+      ${lib.getExe installOnHost} "$HOST" $OPTS
       ${if canSwitchViaSsh
-        then ''${lib.getExe sshSwitch} "$TARGET"''
+        then ''${lib.getExe sshSwitch} "$TARGET" $OPTS''
         else ''ssh "$HOST" -- systemctl restart "microvm@${hostName}.service"''
        }
     '';