fix env on directory change (#13)

Fixes #12

e2e/gopath/.tool-versions

@@ -0,0 +1 @@
+golang 1.19.5

e2e/gopath/18/.tool-versions

@@ -0,0 +1 @@
+golang 1.18.10

e2e/gopath/test_gopath

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+
+set -e
+eval "$(rtx activate -s bash)"
+_rtx_hook
+
+assert_gopath() {
+  local expected="$1"
+  if [[ "$GOPATH" != "$expected" ]]; then
+    echo "Invalid GOPATH: $GOPATH, expected: $expected"
+    exit 1
+  fi
+}
+
+rtx i [email protected] [email protected] && _rtx_hook
+assert_gopath "$RTX_DATA_DIR/installs/golang/1.19.5/packages"
+cd 18 && _rtx_hook
+assert_gopath "$RTX_DATA_DIR/installs/golang/1.18.10/packages"
+cd .. && _rtx_hook
+assert_gopath "$RTX_DATA_DIR/installs/golang/1.19.5/packages"

src/cli/env.rs

@@ -94,4 +94,12 @@ mod test {
                 .as_ref()
         ));
     }
+
+    #[test]
+    fn test_env_golang() {
+        assert_cli!("plugin", "add", "golang");
+        assert_cli!("install", "golang");
+        let Output { stdout, .. } = assert_cli!("env", "golang", "-s", "bash");
+        assert!(stdout.content.contains("GOROOT="));
+    }
 }

src/cli/hook_env.rs

@@ -1,5 +1,4 @@
 use std::collections::HashMap;
-use std::env;
 
 use color_eyre::eyre::Result;
 
@@ -10,7 +9,7 @@ use crate::env_diff::{EnvDiff, EnvDiffOperation, EnvDiffPatches};
 use crate::hook_env::HookEnvWatches;
 use crate::output::Output;
 use crate::shell::{get_shell, ShellType};
-use crate::{dirs, hook_env};
+use crate::{dirs, env, hook_env};
 
 /// [internal] called by activate hook to update env vars directory change
 #[derive(Debug, clap::Args)]
@@ -28,14 +27,14 @@ impl Command for HookEnv {
         config.settings.missing_runtime_behavior = Ignore;
         config.ensure_installed()?;
 
-        let current_env = self.clear_old_env(env::vars().collect(), out)?;
+        self.clear_old_env(out);
         let mut env: HashMap<String, String> = config
             .env()?
             .iter()
             .map(|(k, v)| (k.to_string_lossy().into(), v.to_string_lossy().into()))
             .collect();
         env.insert("__RTX_DIR".into(), dirs::CURRENT.to_string_lossy().into());
-        let diff = EnvDiff::new(&current_env, &env);
+        let diff = EnvDiff::new(&env::PRISTINE_ENV, &env);
         let mut patches = diff.to_patches();
         patches.push(EnvDiffOperation::Add(
             "__RTX_DIFF".into(),
@@ -81,46 +80,13 @@ impl HookEnv {
         output
     }
 
-    fn clear_old_env(
-        &self,
-        env: HashMap<String, String>,
-        out: &mut Output,
-    ) -> Result<HashMap<String, String>> {
-        let patches = get_env_diff(&env)?.reverse().to_patches();
+    fn clear_old_env(&self, out: &mut Output) {
+        let patches = env::__RTX_DIFF.reverse().to_patches();
         let output = self.build_env_commands(&patches);
         out.stdout.write(output);
-
-        Ok(apply_patches(&env, &patches))
-    }
-}
-
-fn get_env_diff(env: &HashMap<String, String>) -> Result<EnvDiff> {
-    let json = env.get("__RTX_DIFF").cloned();
-    match json {
-        Some(json) => Ok(EnvDiff::deserialize(&json)?),
-        None => Ok(EnvDiff::default()),
     }
 }
 
-fn apply_patches(
-    env: &HashMap<String, String>,
-    patches: &EnvDiffPatches,
-) -> HashMap<String, String> {
-    let mut new_env = env.clone();
-    for patch in patches {
-        match patch {
-            EnvDiffOperation::Add(k, v) | EnvDiffOperation::Change(k, v) => {
-                new_env.insert(k.into(), v.into());
-            }
-            EnvDiffOperation::Remove(k) => {
-                new_env.remove(k);
-            }
-        }
-    }
-
-    new_env
-}
-
 #[cfg(test)]
 mod test {
     use crate::assert_cli;

src/env.rs

@@ -1,9 +1,13 @@
+use std::collections::HashMap;
 pub use std::env::*;
 use std::ffi::OsString;
 use std::path::PathBuf;
 
 use lazy_static::lazy_static;
 
+use crate::env_diff::EnvDiff;
+use crate::hook_env::get_pristine_env;
+
 lazy_static! {
     pub static ref ARGS: Vec<String> = args().collect();
     pub static ref HOME: PathBuf = if cfg!(test) {
@@ -72,13 +76,27 @@ lazy_static! {
         var("RTX_MISSING_RUNTIME_BEHAVIOR").ok()
     };
     pub static ref __RTX_DIR: Option<PathBuf> = var_os("__RTX_DIR").map(PathBuf::from);
+    pub static ref __RTX_DIFF: EnvDiff = get_env_diff();
+    pub static ref PRISTINE_ENV: HashMap<String, String> =
+        get_pristine_env(&__RTX_DIFF, vars().collect());
     pub static ref RTX_DEFAULT_TOOL_VERSIONS_FILENAME: String = if cfg!(test) {
         ".tool-versions".into()
     } else {
         var("RTX_DEFAULT_TOOL_VERSIONS_FILENAME").unwrap_or_else(|_| ".tool-versions".into())
     };
 }
 
+fn get_env_diff() -> EnvDiff {
+    let env = vars().collect::<HashMap<_, _>>();
+    match env.get("__RTX_DIFF") {
+        Some(raw) => EnvDiff::deserialize(raw).unwrap_or_else(|err| {
+            warn!("Failed to deserialize __RTX_DIFF: {}", err);
+            EnvDiff::default()
+        }),
+        None => EnvDiff::default(),
+    }
+}
+
 fn var_is_true(key: &str) -> bool {
     match var(key) {
         Ok(v) => v == "true" || v == "1",

src/env_diff.rs

@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::fmt::Debug;
 use std::io::prelude::*;
 use std::path::Path;
 
@@ -7,11 +8,12 @@ use color_eyre::eyre::Result;
 use flate2::write::GzDecoder;
 use flate2::write::GzEncoder;
 use flate2::Compression;
+use itertools::Itertools;
 use serde_derive::{Deserialize, Serialize};
 
 use crate::{cmd, env};
 
-#[derive(Debug, Default, Serialize, Deserialize)]
+#[derive(Default, Serialize, Deserialize)]
 pub struct EnvDiff {
     #[serde(default)]
     pub old: HashMap<String, String>,
@@ -50,35 +52,38 @@ impl EnvDiff {
     }
 
     pub fn from_bash_script(script: &Path, env: HashMap<String, String>) -> Result<Self> {
-        let mut cmd = cmd!(
+        let out = cmd!(
             "bash",
             "-c",
             indoc::formatdoc! {"
                 set -e
                 . {script}
                 env
             ", script = script.display()}
-        );
-        for (k, v) in env.iter() {
-            cmd = cmd.env(k, v);
-        }
-        let out = cmd.read()?;
+        )
+        .full_env(&env)
+        .read()?;
 
         let mut additions = HashMap::new();
         for line in out.lines() {
             let (k, v) = line.split_once('=').unwrap_or_default();
-            if k == "_" || k == "SHLVL" || k == "PATH" || env.contains_key(k) {
+            if k == "_" || k == "SHLVL" || k == "PATH" {
                 continue;
             }
+            if let Some(orig) = env.get(k) {
+                if v == orig {
+                    continue;
+                }
+            }
             additions.insert(k.into(), v.into());
         }
         Ok(Self::new(&env::vars().collect(), &additions))
     }
 
-    pub fn deserialize(json: &str) -> Result<EnvDiff> {
+    pub fn deserialize(raw: &str) -> Result<EnvDiff> {
         let mut writer = Vec::new();
         let mut decoder = GzDecoder::new(writer);
-        let bytes = BASE64_STANDARD_NO_PAD.decode(json)?;
+        let bytes = BASE64_STANDARD_NO_PAD.decode(raw)?;
         decoder.write_all(&bytes[..])?;
         writer = decoder.finish()?;
         Ok(rmp_serde::from_slice(&writer[..])?)
@@ -117,6 +122,22 @@ impl EnvDiff {
     }
 }
 
+impl Debug for EnvDiff {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let print_sorted = |hashmap: &HashMap<String, String>| {
+            hashmap
+                .iter()
+                .map(|(k, v)| format!("{k}={v}"))
+                .sorted()
+                .collect::<Vec<_>>()
+        };
+        f.debug_struct("EnvDiff")
+            .field("old", &print_sorted(&self.old))
+            .field("new", &print_sorted(&self.new))
+            .finish()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use insta::assert_debug_snapshot;
@@ -191,7 +212,14 @@ mod tests {
     #[test]
     fn test_from_bash_script() {
         let path = dirs::HOME.join("fixtures/exec-env");
-        let ed = EnvDiff::from_bash_script(path.as_path(), HashMap::new()).unwrap();
-        assert_debug_snapshot!(ed.to_patches());
+        let orig = HashMap::from(
+            [
+                ("UNMODIFIED_VAR", "unmodified"),
+                ("MODIFIED_VAR", "original"),
+            ]
+            .map(|(k, v)| (k.into(), v.into())),
+        );
+        let ed = EnvDiff::from_bash_script(path.as_path(), orig).unwrap();
+        assert_debug_snapshot!(ed);
     }
 }

src/hook_env.rs

@@ -9,6 +9,7 @@ use flate2::write::GzDecoder;
 use flate2::write::GzEncoder;
 use flate2::Compression;
 
+use crate::env_diff::{EnvDiff, EnvDiffOperation, EnvDiffPatches};
 use crate::{dirs, env};
 
 /// this function will early-exit the application if hook-env is being
@@ -26,6 +27,16 @@ pub fn should_exit_early(current_env: HashMap<String, String>) -> bool {
     true
 }
 
+/// this returns the environment as if __RTX_DIFF was reversed
+/// putting the shell back into a state before hook-env was run
+pub fn get_pristine_env(
+    rtx_diff: &EnvDiff,
+    orig_env: HashMap<String, String>,
+) -> HashMap<String, String> {
+    let patches = rtx_diff.reverse().to_patches();
+    apply_patches(&orig_env, &patches)
+}
+
 fn has_rf_path_changed(env: &HashMap<String, String>) -> bool {
     if let Some(prev) = env.get("__RTX_DIR").map(PathBuf::from) {
         if prev == dirs::CURRENT.as_path() {
@@ -70,6 +81,27 @@ pub fn deserialize_watches(raw: String) -> Result<HookEnvWatches> {
     Ok(rmp_serde::from_slice(&writer[..])?)
 }
 
+fn apply_patches(
+    env: &HashMap<String, String>,
+    patches: &EnvDiffPatches,
+) -> HashMap<String, String> {
+    let mut new_env = env.clone();
+    for patch in patches {
+        match patch {
+            EnvDiffOperation::Add(k, v) | EnvDiffOperation::Change(k, v) => {
+                eprintln!("adding {k}");
+                new_env.insert(k.into(), v.into());
+            }
+            EnvDiffOperation::Remove(k) => {
+                eprintln!("removing {k}");
+                new_env.remove(k);
+            }
+        }
+    }
+
+    new_env
+}
+
 #[cfg(test)]
 mod test {
     use std::time::UNIX_EPOCH;
@@ -122,4 +154,20 @@ mod test {
             &UNIX_EPOCH
         );
     }
+
+    #[test]
+    fn test_apply_patches() {
+        let mut env = HashMap::new();
+        env.insert("foo".into(), "bar".into());
+        env.insert("baz".into(), "qux".into());
+        let patches = vec![
+            EnvDiffOperation::Add("foo".into(), "bar".into()),
+            EnvDiffOperation::Change("baz".into(), "qux".into()),
+            EnvDiffOperation::Remove("quux".into()),
+        ];
+        let new_env = apply_patches(&env, &patches);
+        assert_eq!(new_env.len(), 2);
+        assert_eq!(new_env.get("foo").unwrap(), "bar");
+        assert_eq!(new_env.get("baz").unwrap(), "qux");
+    }
 }

src/runtimes.rs

@@ -203,7 +203,9 @@ impl RuntimeVersion {
         if !self.is_installed() || !script.exists() {
             return Ok(HashMap::new());
         }
-        let ed = EnvDiff::from_bash_script(&script, self.script_env())?;
+        let mut env: HashMap<String, String> = env::PRISTINE_ENV.clone();
+        env.extend(self.script_env());
+        let ed = EnvDiff::from_bash_script(&script, env)?;
         let env = ed
             .to_patches()
             .into_iter()

src/snapshots/rtx__env_diff__tests__from_bash_script.snap

@@ -1,10 +1,11 @@
 ---
 source: src/env_diff.rs
-expression: ed.to_patches()
+expression: ed
 ---
-[
-    Add(
-        "CUSTOM_ENV",
-        "rtx_env",
-    ),
-]
+EnvDiff {
+    old: [],
+    new: [
+        "ADDED_VAR=added",
+        "MODIFIED_VAR=modified",
+    ],
+}

test/fixtures/exec-env

@@ -1,3 +1,5 @@
 #!/usr/bin/env bash
 
-export CUSTOM_ENV=rtx_env
+export UNMODIFIED_VAR="unmodified"
+export MODIFIED_VAR="modified"
+export ADDED_VAR="added"