Add shell-wrapper to inject env variables into non-login shells

modules/systemd/native/default.nix

@@ -1,10 +1,13 @@
 { config, pkgs, lib, ... }:
 with lib; {
 
+  imports = [
+    ./wrap-shell.nix
+  ];
+
   config =
     let
       cfg = config.wsl;
-      nativeUtils = pkgs.callPackage ../../../utils { };
 
       bashWrapper = pkgs.writeShellScriptBin "sh" ''
         export PATH="$PATH:${lib.makeBinPath [ pkgs.systemd pkgs.gnugrep ]}"
@@ -13,6 +16,8 @@ with lib; {
     in
     mkIf (cfg.enable && cfg.nativeSystemd) {
 
+      system.build.nativeUtils = pkgs.callPackage ../../../utils { };
+
       wsl = {
         binShPkg = bashWrapper;
         wslConf = {
@@ -25,7 +30,7 @@ with lib; {
         shimSystemd = stringAfter [ ] ''
           echo "setting up /sbin/init shim..."
           mkdir -p /sbin
-          ln -sf ${nativeUtils}/bin/systemd-shim /sbin/init
+          ln -sf ${config.system.build.nativeUtils}/bin/systemd-shim /sbin/init
         '';
         setupLogin = lib.mkIf cfg.populateBin (stringAfter [ ] ''
           echo "setting up /bin/login..."
@@ -38,7 +43,7 @@ with lib; {
         # preserve $PATH from parent
         variables.PATH = [ "$PATH" ];
         extraInit = ''
-          eval $(${nativeUtils}/bin/split-path --automount-root="${cfg.wslConf.automount.root}" ${lib.optionalString cfg.interop.includePath "--include-interop"})
+          eval $(${config.system.build.nativeUtils}/bin/split-path --automount-root="${cfg.wslConf.automount.root}" ${lib.optionalString cfg.interop.includePath "--include-interop"})
         '';
       };
     };

modules/systemd/native/wrap-shell.nix

@@ -0,0 +1,35 @@
+{ config, lib, utils, pkgs, modulesPath, ... }:
+
+with lib;
+
+let
+  cfg = config.wsl;
+
+  users-groups-module = import "${modulesPath}/config/users-groups.nix" {
+    inherit lib utils pkgs;
+    config = recursiveUpdate config {
+      users.users = mapAttrs
+        (n: v: v // {
+          shell =
+            let
+              shellPath = utils.toShellPath v.shell;
+              wrapper = pkgs.stdenvNoCC.mkDerivation {
+                name = "wrapped-${last (splitString "/" (shellPath))}";
+                buildCommand = ''
+                  mkdir -p $out/bin
+                  cp ${config.system.build.nativeUtils}/bin/shell-wrapper $out/wrapper
+                  ln -s ${shellPath} $out/shell
+                '';
+              };
+            in
+            wrapper.outPath + "/wrapper";
+        })
+        config.users.users;
+    };
+  };
+in
+{
+  config = mkIf (cfg.enable && cfg.nativeSystemd) {
+    system.activationScripts.users = users-groups-module.config.system.activationScripts.users;
+  };
+}

utils/Cargo.toml

@@ -24,3 +24,7 @@ path = "src/shim.rs"
 [[bin]]
 name = "split-path"
 path = "src/split_path.rs"
+
+[[bin]]
+name = "shell-wrapper"
+path = "src/shell_wrapper.rs"

utils/src/shell_wrapper.rs

@@ -0,0 +1,56 @@
+use anyhow::{anyhow, Context};
+use std::os::unix::process::CommandExt;
+use std::process::Command;
+use std::{env, fs::read_link};
+
+fn real_main() -> anyhow::Result<()> {
+    let exe = read_link("/proc/self/exe").context("when locating the wrapper binary")?;
+    let exe_dir = exe.parent().ok_or(anyhow!(
+        "could not locate the wrapper binary's parent directory"
+    ))?;
+
+    // Some binaries behave differently depending on the file name they are called with (arg[0]).
+    // Therefore we dereference our symlink to get whatever it was originally.
+    let shell = read_link(exe_dir.join("shell")).context("when locating the wrapped shell")?;
+
+    // Skip if environment was already set
+    if env::var_os("__NIXOS_SET_ENVIRONMENT_DONE") != Some("1".into()) {
+        // Load the environment from /etc/set-environment
+        let output = Command::new("/bin/sh")
+            .args(&["-c", ". /etc/set-environment && /usr/bin/env -0"])
+            .output()
+            .context("when reading /etc/set-environment")?;
+
+        // Parse the output
+        let output_string =
+            String::from_utf8(output.stdout).context("when decoding the output of env")?;
+        let env = output_string
+            .split('\0')
+            .filter(|entry| !entry.is_empty())
+            .map(|entry| {
+                entry
+                    .split_once("=")
+                    .ok_or(anyhow!("invalid env entry: {}", entry))
+            })
+            .collect::<Result<Vec<_>, _>>()
+            .context("when parsing the output of env")?;
+
+        // Apply the environment variables
+        for &(key, val) in &env {
+            env::set_var(key, val);
+        }
+    }
+
+    Err(anyhow!(Command::new(&shell)
+        .arg0(shell)
+        .args(env::args_os().skip(1))
+        .exec())
+    .context("when trying to exec the wrapped shell"))
+}
+
+fn main() {
+    let result = real_main();
+
+    env::set_var("RUST_BACKTRACE", "1");
+    eprintln!("[shell-wrapper] Error: {:?}", result.unwrap_err());
+}