From c56db2fafa2350ab0b0046bbafa00e3b68871205 Mon Sep 17 00:00:00 2001
From: Dean Li <deantvv@gmail.com>
Date: Mon, 18 Oct 2021 21:03:02 +0800
Subject: [PATCH] os: fix putenv

Check key and value before calling `env::set_var`

From doc of `env::set_var`,
```

This function may panic if `key` is empty, contains an ASCII equals sign
`'='` or the NUL character `'\0'`, or when the value contains the NUL
character.
```
---
 Lib/test/test_posix.py |  2 +-
 vm/src/stdlib/os.rs    | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/Lib/test/test_posix.py b/Lib/test/test_posix.py
index 7a4d59e18..e8957620e 100644
--- a/Lib/test/test_posix.py
+++ b/Lib/test/test_posix.py
@@ -983,7 +983,7 @@ class PosixTester(unittest.TestCase):
             self.assertEqual(type(k), item_type)
             self.assertEqual(type(v), item_type)
 
-    @unittest.skip("TODO: RUSTPYTHON, thread 'main' panicked at 'failed to set environment variable'")
+    @unittest.skipUnless(os.name == 'posix', "TODO: RUSTPYTHON, only works on Unix")
     def test_putenv(self):
         with self.assertRaises(ValueError):
             os.putenv('FRUIT\0VEGETABLE', 'cabbage')
diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs
index 82bd69873..59de9a18a 100644
--- a/vm/src/stdlib/os.rs
+++ b/vm/src/stdlib/os.rs
@@ -681,6 +681,25 @@ pub(super) mod _os {
             Either::A(ref s) => s.as_str().as_ref(),
             Either::B(ref b) => super::bytes_as_osstr(b.as_bytes(), vm)?,
         };
+
+        #[cfg(unix)]
+        {
+            use super::ffi_ext::OsStrExt;
+            let check_embedded_null = |s: &ffi::OsStr| ffi::CString::new(s.as_bytes()).is_ok();
+
+            if !check_embedded_null(key) || !check_embedded_null(value) {
+                return Err(vm.new_value_error("embedded null byte".to_string()));
+            }
+        }
+
+        let check_key = |s: &ffi::OsStr| {
+            s.to_str()
+                .map(|vs| !vs.is_empty() && !vs.contains('='))
+                .unwrap_or(false)
+        };
+        if !check_key(key) {
+            return Err(vm.new_value_error("illegal environment variable name".to_string()));
+        }
         env::set_var(key, value);
         Ok(())
     }