Fix env capture (#5205)

* Fix env capture

* Add test for env capture

crates/nu-cli/src/util.rs

@@ -1,7 +1,7 @@
 use crate::CliError;
 use log::trace;
 use nu_engine::eval_block;
-use nu_parser::{lex, parse, trim_quotes, Token, TokenContents};
+use nu_parser::{lex, parse, unescape_unquote_string, Token, TokenContents};
 use nu_protocol::engine::StateWorkingSet;
 use nu_protocol::{
     ast::Call,
@@ -98,6 +98,10 @@ pub fn print_pipeline_data(
 // env vars into it (in a "NAME"="value" format, quite similar to the output of the Unix 'env'
 // tool), then uses the file to get the spans. The file stays in memory, no filesystem IO is done.
 pub fn gather_parent_env_vars(engine_state: &mut EngineState) {
+    gather_env_vars(std::env::vars(), engine_state);
+}
+
+fn gather_env_vars(vars: impl Iterator<Item = (String, String)>, engine_state: &mut EngineState) {
     fn report_capture_error(engine_state: &EngineState, env_str: &str, msg: &str) {
         let working_set = StateWorkingSet::new(engine_state);
         report_error(
@@ -110,20 +114,35 @@ pub fn gather_parent_env_vars(engine_state: &mut EngineState) {
     }
 
     fn put_env_to_fake_file(name: &str, val: &str, fake_env_file: &mut String) {
-        fake_env_file.push('`');
-        fake_env_file.push_str(name);
-        fake_env_file.push('`');
+        fn push_string_literal(s: &str, fake_env_file: &mut String) {
+            fake_env_file.push('"');
+            for c in s.chars() {
+                if c == '\\' || c == '"' {
+                    fake_env_file.push('\\');
+                }
+                fake_env_file.push(c);
+            }
+            fake_env_file.push('"');
+        }
+
+        push_string_literal(name, fake_env_file);
         fake_env_file.push('=');
-        fake_env_file.push('`');
-        fake_env_file.push_str(val);
-        fake_env_file.push('`');
+        push_string_literal(val, fake_env_file);
         fake_env_file.push('\n');
     }
 
     let mut fake_env_file = String::new();
+    let mut has_pwd = false;
 
-    // Make sure we always have PWD
-    if std::env::var("PWD").is_err() {
+    // Write all the env vars into a fake file
+    for (name, val) in vars {
+        if name == "PWD" {
+            has_pwd = true;
+        }
+        put_env_to_fake_file(&name, &val, &mut fake_env_file);
+    }
+
+    if !has_pwd {
         match std::env::current_dir() {
             Ok(cwd) => {
                 put_env_to_fake_file("PWD", &cwd.to_string_lossy(), &mut fake_env_file);
@@ -142,11 +161,6 @@ pub fn gather_parent_env_vars(engine_state: &mut EngineState) {
         }
     }
 
-    // Write all the env vars into a fake file
-    for (name, val) in std::env::vars() {
-        put_env_to_fake_file(&name, &val, &mut fake_env_file);
-    }
-
     // Lex the fake file, assign spans to all environment variables and add them
     // to stack
     let span_offset = engine_state.next_span_start();
@@ -184,8 +198,19 @@ pub fn gather_parent_env_vars(engine_state: &mut EngineState) {
                     continue;
                 }
 
-                let bytes = trim_quotes(bytes);
-                String::from_utf8_lossy(bytes).to_string()
+                let (bytes, parse_error) = unescape_unquote_string(bytes, *span);
+
+                if parse_error.is_some() {
+                    report_capture_error(
+                        engine_state,
+                        &String::from_utf8_lossy(contents),
+                        "Got unparsable name.",
+                    );
+
+                    continue;
+                }
+
+                bytes
             } else {
                 report_capture_error(
                     engine_state,
@@ -213,10 +238,20 @@ pub fn gather_parent_env_vars(engine_state: &mut EngineState) {
                     continue;
                 }
 
-                let bytes = trim_quotes(bytes);
+                let (bytes, parse_error) = unescape_unquote_string(bytes, *span);
+
+                if parse_error.is_some() {
+                    report_capture_error(
+                        engine_state,
+                        &String::from_utf8_lossy(contents),
+                        "Got unparsable value.",
+                    );
+
+                    continue;
+                }
 
                 Value::String {
-                    val: String::from_utf8_lossy(bytes).to_string(),
+                    val: bytes,
                     span: *span,
                 }
             } else {
@@ -350,3 +385,32 @@ pub fn get_init_cwd() -> PathBuf {
         },
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_gather_env_vars() {
+        let mut engine_state = EngineState::new();
+        let symbols = r##" !"#$%&'()*+,-./:;<=>?@[\]^_`{|}~"##;
+
+        gather_env_vars(
+            [
+                ("FOO".into(), "foo".into()),
+                ("SYMBOLS".into(), symbols.into()),
+                (symbols.into(), "symbols".into()),
+            ]
+            .into_iter(),
+            &mut engine_state,
+        );
+
+        let env = engine_state.env_vars;
+
+        assert!(matches!(env.get("FOO"), Some(Value::String { val, .. }) if val == "foo"));
+        assert!(matches!(env.get("SYMBOLS"), Some(Value::String { val, .. }) if val == symbols));
+        assert!(matches!(env.get(symbols), Some(Value::String { val, .. }) if val == "symbols"));
+        assert!(env.get("PWD").is_some());
+        assert_eq!(env.len(), 4);
+    }
+}

crates/nu-parser/src/lib.rs

@@ -15,7 +15,7 @@ pub use lite_parse::{lite_parse, LiteBlock};
 
 pub use parser::{
     is_math_expression_like, parse, parse_block, parse_duration_bytes, parse_external_call,
-    trim_quotes, Import,
+    trim_quotes, unescape_unquote_string, Import,
 };
 
 #[cfg(feature = "plugin")]