Merge pull request #219 from K900/smaller-oof

fix(split-paths): use standard library tools, fix shell quoting better

checks/rustfmt.nix

@@ -0,0 +1,9 @@
+{ runCommand
+, cargo
+, rustfmt
+, ...
+}:
+runCommand "check-rustfmt" { nativeBuildInputs = [ cargo rustfmt ]; } ''
+  cargo fmt --manifest-path=${./../scripts/native-utils}/Cargo.toml --check
+  touch $out
+''

flake.nix

@@ -54,11 +54,16 @@
             {
               nixpkgs-fmt = pkgs.callPackage ./checks/nixpkgs-fmt.nix args;
               shfmt = pkgs.callPackage ./checks/shfmt.nix args;
+              rustfmt = pkgs.callPackage ./checks/rustfmt.nix args;
               side-effects = pkgs.callPackage ./checks/side-effects.nix args;
               username = pkgs.callPackage ./checks/username.nix args;
+              test-native-utils = self.packages.${system}.utils;
             };
 
-          packages.staticShim = pkgs.pkgsStatic.callPackage ./scripts/native-utils { };
+          packages = {
+            utils = pkgs.callPackage ./scripts/native-utils { };
+            staticUtils = pkgs.pkgsStatic.callPackage ./scripts/native-utils { };
+          };
 
           devShell = pkgs.mkShell {
             RUST_SRC_PATH = "${pkgs.rust.packages.stable.rustPlatform.rustLibSrc}";

scripts/native-utils/src/split_path.rs

@@ -1,38 +1,142 @@
-use std::env;
+use std::{
+    env,
+    ffi::{OsStr, OsString},
+    io::{self, Write},
+    os::unix::prelude::{OsStrExt, OsStringExt},
+    path::{Path, PathBuf},
+};
 
 use clap::Parser;
 
 #[derive(Parser, Debug)]
 struct Args {
-  #[arg(long)]
-  automount_root: String,
+    #[arg(long)]
+    automount_root: PathBuf,
 
-  #[arg(long)]
-  include_interop: bool,
+    #[arg(long)]
+    include_interop: bool,
+}
+
+const SINGLE_QUOTE: u8 = b'\'';
+const DOUBLE_QUOTE: u8 = b'"';
+
+fn shell_escape(s: &OsStr) -> OsString {
+    // a shameless ripoff of the Python algorithm:
+    // https://github.com/python/cpython/blob/f1f3af7b8245e61a2e0abef03b2c6c5902ed7df8/Lib/shlex.py#L323
+    let mut result = Vec::new();
+
+    result.push(SINGLE_QUOTE);
+
+    for &byte in s.as_bytes() {
+        result.push(byte);
+        if byte == SINGLE_QUOTE {
+            result.push(DOUBLE_QUOTE);
+            result.push(SINGLE_QUOTE);
+            result.push(DOUBLE_QUOTE);
+            result.push(SINGLE_QUOTE);
+        }
+    }
+
+    result.push(SINGLE_QUOTE);
+
+    OsString::from_vec(result)
+}
+
+fn build_export(var: &str, paths: &[PathBuf]) -> OsString {
+    let mut result = OsString::new();
+    result.push("export ");
+    result.push(var);
+    result.push("=");
+    result.push(shell_escape(
+        &env::join_paths(paths).expect("paths must be valid"),
+    ));
+    result.push("\n");
+    result
+}
+
+fn do_split_paths(path: &OsStr, automount_root: &Path, include_interop: bool) -> OsString {
+    let mut native = vec![];
+    let mut interop = vec![];
+
+    for part in env::split_paths(&path) {
+        if part.starts_with(automount_root) {
+            interop.push(part);
+        } else {
+            native.push(part);
+        }
+    }
+
+    if include_interop {
+        native.extend(interop.clone());
+    };
+
+    let mut result = OsString::new();
+    result.push(build_export("PATH", &native));
+    result.push(build_export("WSLPATH", &interop));
+    result
 }
 
 fn main() -> anyhow::Result<()> {
-  let args = Args::parse();
+    let args = Args::parse();
+
+    let path = env::var_os("PATH").expect("PATH is not set, aborting");
 
-  let path = env::var("PATH")?;
+    io::stdout()
+        .lock()
+        .write_all(do_split_paths(&path, &args.automount_root, args.include_interop).as_bytes())?;
 
-  let mut native = vec![];
-  let mut interop = vec![];
+    Ok(())
+}
 
-  for part in path.split(':') {
-    if part.starts_with(&args.automount_root) {
-      interop.push(part);
-    } else {
-      native.push(part);
+#[cfg(test)]
+mod tests {
+    use std::{ffi::OsString, path::Path};
+
+    use crate::do_split_paths;
+
+    #[test]
+    fn simple() {
+        assert_eq!(
+            do_split_paths(
+                &OsString::from("/good/foo:/bad/foo"),
+                Path::new("/bad"),
+                false
+            ),
+            OsString::from("export PATH='/good/foo'\nexport WSLPATH='/bad/foo'\n")
+        );
     }
-  }
 
-  if args.include_interop {
-    native.extend(&interop);
-  };
+    #[test]
+    fn exactly_one() {
+        assert_eq!(
+            do_split_paths(&OsString::from("/good/foo"), Path::new("/bad"), true),
+            OsString::from("export PATH='/good/foo'\nexport WSLPATH=''\n")
+        );
+    }
 
-  println!("export PATH='{}'", native.join(":"));
-  println!("export WSLPATH='{}'", interop.join(":"));
+    #[test]
+    fn include_interop() {
+        assert_eq!(
+            do_split_paths(
+                &OsString::from("/good/foo:/bad/foo"),
+                Path::new("/bad"),
+                true
+            ),
+            OsString::from("export PATH='/good/foo:/bad/foo'\nexport WSLPATH='/bad/foo'\n")
+        );
+    }
 
-  Ok(())
+    #[test]
+    fn spicy_escapes() {
+        assert_eq!(
+            do_split_paths(
+                &OsString::from("/good/foo'bar:/bad/foo"),
+                Path::new("/bad"),
+                true
+            ),
+            OsString::from(
+                "export PATH='/good/foo'\"'\"'bar:/bad/foo'\nexport WSLPATH='/bad/foo'\n"
+            )
+        );
+    }
 }