From 198bc22e3b856bf2a86225c2ce5b3a7394e3ac0c Mon Sep 17 00:00:00 2001
From: Robert Hensing <robert@roberthensing.nl>
Date: Fri, 12 May 2023 19:54:54 +0200
Subject: [PATCH] nix: Add command baseDir to parse --expr relative to shebang
 script

---
 doc/manual/src/release-notes/rl-next.md |  1 +
 src/libcmd/installables.cc              |  3 ++-
 src/libutil/args.cc                     |  9 +++++++++
 src/libutil/args.hh                     | 20 ++++++++++++++++++++
 tests/functional/flakes/flakes.sh       | 19 +++++++++++++++++++
 5 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/doc/manual/src/release-notes/rl-next.md b/doc/manual/src/release-notes/rl-next.md
index 4bff3c685..28b6d75f5 100644
--- a/doc/manual/src/release-notes/rl-next.md
+++ b/doc/manual/src/release-notes/rl-next.md
@@ -4,6 +4,7 @@
   contents of any `#! nix` lines and the script's location to a single call.
 
   Verbatim strings may be passed in double backtick (```` `` ````) quotes.
+  `--expr` resolves relative paths based on the shebang script location.
 
   Some examples:
   ```
diff --git a/src/libcmd/installables.cc b/src/libcmd/installables.cc
index e7f58556f..528643dc5 100644
--- a/src/libcmd/installables.cc
+++ b/src/libcmd/installables.cc
@@ -445,7 +445,8 @@ Installables SourceExprCommand::parseInstallables(
         else if (file)
             state->evalFile(lookupFileArg(*state, *file), *vFile);
         else {
-            auto e = state->parseExprFromString(*expr, state->rootPath(CanonPath::fromCwd()));
+            CanonPath dir(CanonPath::fromCwd(getCommandBaseDir()));
+            auto e = state->parseExprFromString(*expr, state->rootPath(dir));
             state->eval(e, *vFile);
         }
 
diff --git a/src/libutil/args.cc b/src/libutil/args.cc
index ab6e0e266..0012b3f47 100644
--- a/src/libutil/args.cc
+++ b/src/libutil/args.cc
@@ -276,6 +276,7 @@ void Args::parseCmdline(const Strings & _cmdline, bool allowShebang)
                     cmdline.push_back(word);
                 }
                 cmdline.push_back(script);
+                commandBaseDir = dirOf(script);
                 for (auto pos = savedArgs.begin(); pos != savedArgs.end();pos++)
                     cmdline.push_back(*pos);
             }
@@ -336,6 +337,14 @@ void Args::parseCmdline(const Strings & _cmdline, bool allowShebang)
         d.completer(*completions, d.n, d.prefix);
 }
 
+Path Args::getCommandBaseDir() const
+{
+    if (parent)
+        return parent->getCommandBaseDir();
+    else
+        return commandBaseDir;
+}
+
 bool Args::processFlag(Strings::iterator & pos, Strings::iterator end)
 {
     assert(pos != end);
diff --git a/src/libutil/args.hh b/src/libutil/args.hh
index e753dcaf6..9c942606e 100644
--- a/src/libutil/args.hh
+++ b/src/libutil/args.hh
@@ -24,6 +24,16 @@ class AddCompletions;
 
 class Args
 {
+    /**
+     * @brief The command's "working directory", but only set when top level.
+     *
+     * Use getCommandBaseDir() to get the directory regardless of whether this
+     * is a top-level command or subcommand.
+     *
+     * @see getCommandBaseDir()
+     */
+    Path commandBaseDir = ".";
+
 public:
 
     /**
@@ -44,6 +54,16 @@ public:
      */
     virtual std::string doc() { return ""; }
 
+    /**
+     * @brief Get the base directory for the command.
+     *
+     * @return Generally the working directory, but in case of a shebang
+     *         interpreter, returns the directory of the script.
+     *
+     * This only returns the correct value after parseCmdline() has run.
+     */
+    Path getCommandBaseDir() const;
+
 protected:
 
     /**
diff --git a/tests/functional/flakes/flakes.sh b/tests/functional/flakes/flakes.sh
index a0a34ffa9..76f3495dd 100644
--- a/tests/functional/flakes/flakes.sh
+++ b/tests/functional/flakes/flakes.sh
@@ -105,6 +105,24 @@ foo
 EOF
 chmod +x $nonFlakeDir/shebang-reject.sh
 
+cat > $nonFlakeDir/shebang-inline-expr.sh <<EOF
+#! $(type -P env) nix
+EOF
+cat >> $nonFlakeDir/shebang-inline-expr.sh <<"EOF"
+#! nix --offline shell
+#! nix --impure --expr ``
+#! nix let flake = (builtins.getFlake (toString ../flake1)).packages;
+#! nix     fooScript = flake.${builtins.currentSystem}.fooScript;
+#! nix     /* just a comment !@#$%^&*()__+ # */
+#! nix  in fooScript
+#! nix ``
+#! nix --no-write-lock-file --command bash
+set -ex
+foo
+echo "$@"
+EOF
+chmod +x $nonFlakeDir/shebang-inline-expr.sh
+
 # Construct a custom registry, additionally test the --registry flag
 nix registry add --registry "$registry" flake1 "git+file://$flake1Dir"
 nix registry add --registry "$registry" flake2 "git+file://$percentEncodedFlake2Dir"
@@ -552,4 +570,5 @@ expectStderr 1 nix flake metadata "$flake2Dir" --no-allow-dirty --reference-lock
 [[ $($nonFlakeDir/shebang.sh) = "foo" ]]
 [[ $($nonFlakeDir/shebang.sh "bar") = "foo"$'\n'"bar" ]]
 [[ $($nonFlakeDir/shebang-comments.sh ) = "foo" ]]
+[[ $($nonFlakeDir/shebang-inline-expr.sh baz) = "foo"$'\n'"baz" ]]
 expect 1 $nonFlakeDir/shebang-reject.sh 2>&1 | grepQuiet -F 'error: unsupported unquoted character in nix shebang: *. Use double backticks to escape?'