qsub,slurm: fix string escaping bug

escapeShellArgs calls toString on the passed item, which for paths
converts it to a string containing the full path outside the nix
store. Avoid this by calling escapeShellArg on strings only.

2 files changed, 90 insertions, 87 deletions

diff --git a/lib/qsub.nix b/lib/qsub.nix
index f5a00ec..a8827a8 100644
--- a/lib/qsub.nix
+++ b/lib/qsub.nix
@@ -1,74 +1,79 @@
-{stdenv, lib, writeScript}:
+{ stdenv, lib, writeScript }:
 
 with lib;
 
-{ ppn, mem, walltime, queue ? null, qsubFlags ? null, tmpDir, sleepTime, qsubPath ? "/usr/bin" }:
+let escape = x: if builtins.typeOf x == "string" then escapeShellArg x else x;
+
+in { ppn, mem, walltime, queue ? null, qsubFlags ? null, tmpDir, sleepTime
+, qsubPath ? "/usr/bin" }:
 drv:
-  let ppnReified = if drv.multicore then ppn else 1;
-  in lib.overrideDerivation drv ({ args, builder, name, ... }: {
-    builder = "/bin/bash";
-    args = let
-      script = writeScript "qsub-script" ''
-        #!${stdenv.shell}
-        while [ ! -e ${tmpDir}/qsub-$PBS_JOBID ] ; do
-          sleep ${toString sleepTime}
-        done
-        set -a
-        . ${tmpDir}/qsub-$PBS_JOBID/nix-set
-        set +a
-        TMPDIR=${tmpDir}/qsub-$PBS_JOBID
-        TEMP=$TMPDIR
-        TMP=$TMPDIR
-        NIX_BUILD_TOP=$TMPDIR
-        cd $TMPDIR
-        ${builder} ${lib.escapeShellArgs args} &> qsub-log
-        echo $? > qsub-exit
-      '';
+let ppnReified = if drv.multicore then ppn else 1;
+in overrideDerivation drv ({ args, builder, name, ... }: {
+  builder = "/bin/bash";
+  args = let
+    script = writeScript "qsub-script" ''
+      #!${stdenv.shell}
+      while [ ! -e ${tmpDir}/qsub-$PBS_JOBID ] ; do
+        sleep ${toString sleepTime}
+      done
+      set -a
+      . ${tmpDir}/qsub-$PBS_JOBID/nix-set
+      set +a
+      TMPDIR=${tmpDir}/qsub-$PBS_JOBID
+      TEMP=$TMPDIR
+      TMP=$TMPDIR
+      NIX_BUILD_TOP=$TMPDIR
+      cd $TMPDIR
+      ${builder} ${concatMapStringsSep " " escape args} &> qsub-log
+      echo $? > qsub-exit
+    '';
 
-      qsub = writeScript "qsub" ''
-        #!${stdenv.shell}
-        PATH=${qsubPath}
-        SHELL=/bin/sh
-        NIX_BUILD_CORES=${toString ppnReified}
+    qsub = writeScript "qsub" ''
+      #!${stdenv.shell}
+      PATH=${qsubPath}
+      SHELL=/bin/sh
+      NIX_BUILD_CORES=${toString ppnReified}
 
-        while : ; do
-          qsub -l nodes=1:ppn=${toString ppnReified},mem=${toString mem}gb,walltime=${walltime} \
-            -N "${name}" \
-            ${optionalString (queue != null) "-q ${queue}"} \
-            ${optionalString (qsubFlags != null) qsubFlags} \
-            ${script} 2>&1 > id
-          if [ $? -eq 0 ] ; then
-            break
-          fi
-          if ! grep "Please retry" id > /dev/null ; then
-            cat id >&2
-            exit 1
-          fi
-          sleep ${toString sleepTime}
-        done
-        id=$(cat id)
-        echo $id
+      while : ; do
+        qsub -l nodes=1:ppn=${toString ppnReified},mem=${
+          toString mem
+        }gb,walltime=${walltime} \
+          -N "${name}" \
+          ${optionalString (queue != null) "-q ${queue}"} \
+          ${optionalString (qsubFlags != null) qsubFlags} \
+          ${script} 2>&1 > id
+        if [ $? -eq 0 ] ; then
+          break
+        fi
+        if ! grep "Please retry" id > /dev/null ; then
+          cat id >&2
+          exit 1
+        fi
+        sleep ${toString sleepTime}
+      done
+      id=$(cat id)
+      echo $id
 
-        function cleanup {
-          qdel $id 2>/dev/null || true
-          sleep ${toString sleepTime}
-          rm -rf ${tmpDir}/qsub-$id
-        }
-        trap cleanup INT TERM EXIT
+      function cleanup {
+        qdel $id 2>/dev/null || true
+        sleep ${toString sleepTime}
+        rm -rf ${tmpDir}/qsub-$id
+      }
+      trap cleanup INT TERM EXIT
 
-        cp -r $TMPDIR ${tmpDir}/qsub-$id
-        set > ${tmpDir}/qsub-$id/nix-set
-        until qstat -f ''${id%%.} 2>&1 | grep "\(Unknown Job\|job_state = C\)" > /dev/null ; do
-          sleep ${toString sleepTime}
-        done
-        cat ${tmpDir}/qsub-$id/qsub-log
-        if [ -e ${tmpDir}/qsub-$id/qsub-exit ]; then
-          exitCode=$(cat ${tmpDir}/qsub-$id/qsub-exit)
-        else
-          exitCode=1
-        fi
-        exit $exitCode
-      '';
+      cp -r $TMPDIR ${tmpDir}/qsub-$id
+      set > ${tmpDir}/qsub-$id/nix-set
+      until qstat -f ''${id%%.} 2>&1 | grep "\(Unknown Job\|job_state = C\)" > /dev/null ; do
+        sleep ${toString sleepTime}
+      done
+      cat ${tmpDir}/qsub-$id/qsub-log
+      if [ -e ${tmpDir}/qsub-$id/qsub-exit ]; then
+        exitCode=$(cat ${tmpDir}/qsub-$id/qsub-exit)
+      else
+        exitCode=1
+      fi
+      exit $exitCode
+    '';
 
-      in [ "-c" qsub ];
-  })
+  in [ "-c" qsub ];
+})
diff --git a/lib/slurm.nix b/lib/slurm.nix
index 30a5f5e..4ad41ab 100644
--- a/lib/slurm.nix
+++ b/lib/slurm.nix
@@ -1,28 +1,26 @@
-{stdenv, lib, writeScript, coreutils}:
+{ stdenv, lib, writeScript, coreutils }:
 
 with lib;
 
-{ ppn, mem, walltime, partition ? null, slurmFlags ? null, salloc ? "/usr/bin/salloc", srun ? "/usr/bin/srun" }:
-drv:
-  let ppnReified = if drv.multicore then ppn else 1;
-  in lib.overrideDerivation drv ({ args, builder, name, ... }: {
-    builder = stdenv.shell;
-    args = let
-      script = writeScript "slurm-script" ''
-        #!${stdenv.shell}
-        ${builder} ${lib.escapeShellArgs args}
-      '';
+let escape = x: if builtins.typeOf x == "string" then escapeShellArg x else x;
 
-      slurm = writeScript "slurm" ''
-        #!${stdenv.shell}
-        NIX_BUILD_CORES=${toString ppnReified}
+in { ppn, mem, walltime, partition ? null, slurmFlags ? null
+, salloc ? "/usr/bin/salloc", srun ? "/usr/bin/srun" }:
+drv:
+let ppnReified = if drv.multicore then ppn else 1;
+in overrideDerivation drv ({ args, builder, name, ... }: {
+  builder = stdenv.shell;
+  args = let
+    slurm = writeScript "slurm" ''
+      #!${stdenv.shell}
+      NIX_BUILD_CORES=${toString ppnReified}
 
-        ${salloc} -c $NIX_BUILD_CORES --mem=${toString mem}G -t ${walltime} \
-          -J "${name}" \
-          ${optionalString (partition != null) "-p ${partition}"} \
-          ${optionalString (slurmFlags != null) slurmFlags} \
-          ${srun} ${script}
-      '';
+      ${salloc} -c $NIX_BUILD_CORES --mem=${toString mem}G -t ${walltime} \
+        -J "${name}" \
+        ${optionalString (partition != null) "-p ${partition}"} \
+        ${optionalString (slurmFlags != null) slurmFlags} \
+        ${srun} ${builder} ${concatMapStringsSep " " escape args}
+    '';
 
-      in [ "-c" slurm ];
-  })
+  in [ "-c" slurm ];
+})