use std::cmp;
use std::collections::BTreeMap;
use std::env;
use std::ffi::OsStr;
use std::ffi::OsString;
use std::ffi::c_void;
use std::fmt;
use std::io;
use std::os::windows::ffi::OsStrExt;
use windows_sys::Win32::Foundation::TRUE;
use windows_sys::Win32::Globalization::CSTR_EQUAL;
use windows_sys::Win32::Globalization::CSTR_GREATER_THAN;
use windows_sys::Win32::Globalization::CSTR_LESS_THAN;
use windows_sys::Win32::Globalization::CompareStringOrdinal;
#[derive(Clone, Default)]
pub struct CommandEnv {
clear: bool,
saw_path: bool,
vars: BTreeMap<EnvKey, Option<OsString>>,
}
impl fmt::Debug for CommandEnv {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut debug_command_env = f.debug_struct("CommandEnv");
debug_command_env
.field("clear", &self.clear)
.field("vars", &self.vars);
debug_command_env.finish()
}
}
impl CommandEnv {
pub fn capture(&self) -> BTreeMap<EnvKey, OsString> {
let mut result = BTreeMap::<EnvKey, OsString>::new();
if !self.clear {
for (k, v) in env::vars_os() {
result.insert(k.into(), v);
}
}
for (k, maybe_v) in &self.vars {
if let Some(v) = maybe_v {
result.insert(k.clone(), v.clone());
} else {
result.remove(k);
}
}
result
}
pub fn is_unchanged(&self) -> bool {
!self.clear && self.vars.is_empty()
}
pub fn capture_if_changed(&self) -> Option<BTreeMap<EnvKey, OsString>> {
if self.is_unchanged() {
None
} else {
Some(self.capture())
}
}
pub fn set(&mut self, key: &OsStr, value: &OsStr) {
let key = EnvKey::from(key);
self.maybe_saw_path(&key);
self.vars.insert(key, Some(value.to_owned()));
}
pub fn clear(&mut self) {
self.clear = true;
self.vars.clear();
}
pub fn have_changed_path(&self) -> bool {
self.saw_path || self.clear
}
fn maybe_saw_path(&mut self, key: &EnvKey) {
if !self.saw_path && key == "PATH" {
self.saw_path = true;
}
}
}
#[derive(Clone, Debug, Eq)]
#[doc(hidden)]
pub struct EnvKey {
os_string: OsString,
utf16: Vec<u16>,
}
impl EnvKey {
pub fn new<T: Into<OsString>>(key: T) -> Self {
EnvKey::from(key.into())
}
}
impl Ord for EnvKey {
fn cmp(&self, other: &Self) -> cmp::Ordering {
unsafe {
let result = CompareStringOrdinal(
self.utf16.as_ptr(),
self.utf16.len() as _,
other.utf16.as_ptr(),
other.utf16.len() as _,
TRUE,
);
match result {
CSTR_LESS_THAN => cmp::Ordering::Less,
CSTR_EQUAL => cmp::Ordering::Equal,
CSTR_GREATER_THAN => cmp::Ordering::Greater,
_ => panic!(
"comparing environment keys failed: {}",
std::io::Error::last_os_error()
),
}
}
}
}
impl PartialOrd for EnvKey {
fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
Some(self.cmp(other))
}
}
impl PartialEq for EnvKey {
fn eq(&self, other: &Self) -> bool {
if self.utf16.len() != other.utf16.len() {
false
} else {
self.cmp(other) == cmp::Ordering::Equal
}
}
}
impl PartialOrd<str> for EnvKey {
fn partial_cmp(&self, other: &str) -> Option<cmp::Ordering> {
Some(self.cmp(&EnvKey::new(other)))
}
}
impl PartialEq<str> for EnvKey {
fn eq(&self, other: &str) -> bool {
if self.os_string.len() != other.len() {
false
} else {
self.cmp(&EnvKey::new(other)) == cmp::Ordering::Equal
}
}
}
impl From<OsString> for EnvKey {
fn from(k: OsString) -> Self {
EnvKey {
utf16: k.encode_wide().collect(),
os_string: k,
}
}
}
impl From<&OsStr> for EnvKey {
fn from(k: &OsStr) -> Self {
Self::from(k.to_os_string())
}
}
pub fn ensure_no_nuls<T: AsRef<OsStr>>(s: T) -> io::Result<T> {
if s.as_ref().encode_wide().any(|b| b == 0) {
Err(io::Error::new(
io::ErrorKind::InvalidInput,
"nul byte found in provided data",
))
} else {
Ok(s)
}
}
pub fn make_envp(
maybe_env: Option<BTreeMap<EnvKey, OsString>>,
) -> io::Result<(*mut c_void, Vec<u16>)> {
if let Some(env) = maybe_env {
let mut blk = Vec::new();
if env.is_empty() {
blk.push(0);
}
for (k, v) in env {
ensure_no_nuls(k.os_string)?;
blk.extend(k.utf16);
blk.push('=' as u16);
blk.extend(ensure_no_nuls(v)?.encode_wide());
blk.push(0);
}
blk.push(0);
Ok((blk.as_mut_ptr() as *mut c_void, blk))
} else {
Ok((std::ptr::null_mut(), Vec::new()))
}
}