Skip to main content

std\sys\env/
windows.rs

1use crate::ffi::{OsStr, OsString};
2use crate::os::windows::prelude::*;
3use crate::sys::pal::{c, cvt, fill_utf16_buf, to_u16s};
4use crate::{fmt, io, ptr, slice};
5
6pub struct Env {
7    base: *mut c::WCHAR,
8    iter: EnvIterator,
9}
10
11impl fmt::Debug for Env {
12    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
13        let Self { base: _, iter } = self;
14        f.debug_list().entries(iter.clone()).finish()
15    }
16}
17
18impl Iterator for Env {
19    type Item = (OsString, OsString);
20
21    fn next(&mut self) -> Option<(OsString, OsString)> {
22        let Self { base: _, iter } = self;
23        iter.next()
24    }
25}
26
27#[derive(Clone)]
28struct EnvIterator(*mut c::WCHAR);
29
30impl Iterator for EnvIterator {
31    type Item = (OsString, OsString);
32
33    fn next(&mut self) -> Option<(OsString, OsString)> {
34        let Self(cur) = self;
35        loop {
36            unsafe {
37                if **cur == 0 {
38                    return None;
39                }
40                let p = *cur as *const u16;
41                let mut len = 0;
42                while *p.add(len) != 0 {
43                    len += 1;
44                }
45                let s = slice::from_raw_parts(p, len);
46                *cur = cur.add(len + 1);
47
48                // Windows allows environment variables to start with an equals
49                // symbol (in any other position, this is the separator between
50                // variable name and value). Since`s` has at least length 1 at
51                // this point (because the empty string terminates the array of
52                // environment variables), we can safely slice.
53                let pos = match s[1..].iter().position(|&u| u == b'=' as u16).map(|p| p + 1) {
54                    Some(p) => p,
55                    None => continue,
56                };
57                return Some((
58                    OsStringExt::from_wide(&s[..pos]),
59                    OsStringExt::from_wide(&s[pos + 1..]),
60                ));
61            }
62        }
63    }
64}
65
66impl Drop for Env {
67    fn drop(&mut self) {
68        unsafe {
69            c::FreeEnvironmentStringsW(self.base);
70        }
71    }
72}
73
74pub fn env() -> Env {
75    unsafe {
76        let ch = c::GetEnvironmentStringsW();
77        if ch.is_null() {
78            panic!("failure getting env string from OS: {}", io::Error::last_os_error());
79        }
80        Env { base: ch, iter: EnvIterator(ch) }
81    }
82}
83
84pub fn getenv(k: &OsStr) -> Option<OsString> {
85    let k = to_u16s(k).ok()?;
86    fill_utf16_buf(
87        |buf, sz| unsafe { c::GetEnvironmentVariableW(k.as_ptr(), buf, sz) },
88        OsStringExt::from_wide,
89    )
90    .ok()
91}
92
93pub unsafe fn setenv(k: &OsStr, v: &OsStr) -> io::Result<()> {
94    // SAFETY: We ensure that k and v are null-terminated wide strings.
95    unsafe {
96        let k = to_u16s(k)?;
97        let v = to_u16s(v)?;
98
99        cvt(c::SetEnvironmentVariableW(k.as_ptr(), v.as_ptr())).map(drop)
100    }
101}
102
103pub unsafe fn unsetenv(n: &OsStr) -> io::Result<()> {
104    // SAFETY: We ensure that v is a null-terminated wide strings.
105    unsafe {
106        let v = to_u16s(n)?;
107        cvt(c::SetEnvironmentVariableW(v.as_ptr(), ptr::null())).map(drop)
108    }
109}