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(()) }