lune_utils/process/
env.rs

1#![allow(clippy::missing_panics_doc)]
2
3use std::{
4    collections::BTreeMap,
5    env::vars_os,
6    ffi::{OsStr, OsString},
7    sync::{Arc, Mutex},
8};
9
10use mlua::prelude::*;
11use os_str_bytes::{OsStrBytes, OsStringBytes};
12
13// Inner (shared) struct
14
15#[derive(Debug, Default)]
16struct ProcessEnvInner {
17    values: BTreeMap<OsString, OsString>,
18}
19
20impl FromIterator<(OsString, OsString)> for ProcessEnvInner {
21    fn from_iter<T: IntoIterator<Item = (OsString, OsString)>>(iter: T) -> Self {
22        Self {
23            values: iter.into_iter().collect(),
24        }
25    }
26}
27
28/**
29    A struct that can be easily shared, stored in Lua app data,
30    and that also guarantees the pairs are valid OS strings
31    that can be used for process environment variables.
32
33    Usable directly from Lua, implementing both `FromLua` and `LuaUserData`.
34
35    Also provides convenience methods for working with the variables
36    as either `OsString` or `Vec<u8>`, where using the latter implicitly
37    converts to an `OsString` and fails if the conversion is not possible.
38*/
39#[derive(Debug, Clone)]
40pub struct ProcessEnv {
41    inner: Arc<Mutex<ProcessEnvInner>>,
42}
43
44impl ProcessEnv {
45    #[must_use]
46    pub fn empty() -> Self {
47        Self {
48            inner: Arc::new(Mutex::new(ProcessEnvInner::default())),
49        }
50    }
51
52    #[must_use]
53    pub fn current() -> Self {
54        Self {
55            inner: Arc::new(Mutex::new(vars_os().collect())),
56        }
57    }
58
59    #[must_use]
60    pub fn len(&self) -> usize {
61        let inner = self.inner.lock().unwrap();
62        inner.values.len()
63    }
64
65    #[must_use]
66    pub fn is_empty(&self) -> bool {
67        let inner = self.inner.lock().unwrap();
68        inner.values.is_empty()
69    }
70
71    // OS strings
72
73    #[must_use]
74    pub fn get_all(&self) -> Vec<(OsString, OsString)> {
75        let inner = self.inner.lock().unwrap();
76        inner.values.clone().into_iter().collect()
77    }
78
79    #[must_use]
80    pub fn get_value(&self, key: impl AsRef<OsStr>) -> Option<OsString> {
81        let key = key.as_ref();
82
83        super::validate_os_key(key).ok()?;
84
85        let inner = self.inner.lock().unwrap();
86        inner.values.get(key).cloned()
87    }
88
89    pub fn set_value(&self, key: impl Into<OsString>, val: impl Into<OsString>) {
90        let key = key.into();
91        let val = val.into();
92
93        if super::validate_os_pair((&key, &val)).is_err() {
94            return;
95        }
96
97        let mut inner = self.inner.lock().unwrap();
98        inner.values.insert(key, val);
99    }
100
101    pub fn remove_value(&self, key: impl AsRef<OsStr>) {
102        let key = key.as_ref();
103
104        if super::validate_os_key(key).is_err() {
105            return;
106        }
107
108        let mut inner = self.inner.lock().unwrap();
109        inner.values.remove(key);
110    }
111
112    // Bytes wrappers
113
114    #[must_use]
115    pub fn get_all_bytes(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
116        self.get_all()
117            .into_iter()
118            .filter_map(|(k, v)| Some((k.into_io_vec()?, v.into_io_vec()?)))
119            .collect()
120    }
121
122    #[must_use]
123    pub fn get_value_bytes(&self, key: impl AsRef<[u8]>) -> Option<Vec<u8>> {
124        let key = OsStr::from_io_bytes(key.as_ref())?;
125        let val = self.get_value(key)?;
126        val.into_io_vec()
127    }
128
129    pub fn set_value_bytes(&self, key: impl AsRef<[u8]>, val: impl Into<Vec<u8>>) {
130        let key = OsStr::from_io_bytes(key.as_ref());
131        let val = OsString::from_io_vec(val.into());
132        if let (Some(key), Some(val)) = (key, val) {
133            self.set_value(key, val);
134        }
135    }
136
137    pub fn remove_value_bytes(&self, key: impl AsRef<[u8]>) {
138        let key = OsStr::from_io_bytes(key.as_ref());
139        if let Some(key) = key {
140            self.remove_value(key);
141        }
142    }
143
144    // Plain lua table conversion
145
146    #[doc(hidden)]
147    #[allow(clippy::missing_errors_doc)]
148    pub fn into_plain_lua_table(&self, lua: Lua) -> LuaResult<LuaTable> {
149        let all = self.get_all_bytes();
150        let tab = lua.create_table_with_capacity(0, all.len())?;
151
152        for (key, val) in all {
153            let key = lua.create_string(key)?;
154            let val = lua.create_string(val)?;
155            tab.set(key, val)?;
156        }
157
158        Ok(tab)
159    }
160}
161
162// Iterator implementations
163
164impl IntoIterator for ProcessEnv {
165    type Item = (OsString, OsString);
166    type IntoIter = std::collections::btree_map::IntoIter<OsString, OsString>;
167
168    fn into_iter(self) -> Self::IntoIter {
169        let inner = self.inner.lock().unwrap();
170        inner.values.clone().into_iter()
171    }
172}
173
174impl<K: Into<OsString>, V: Into<OsString>> FromIterator<(K, V)> for ProcessEnv {
175    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
176        Self {
177            inner: Arc::new(Mutex::new(
178                iter.into_iter()
179                    .map(|(k, v)| (k.into(), v.into()))
180                    .filter(|(k, v)| super::validate_os_pair((k, v)).is_ok())
181                    .collect(),
182            )),
183        }
184    }
185}
186
187impl<K: Into<OsString>, V: Into<OsString>> Extend<(K, V)> for ProcessEnv {
188    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
189        let mut inner = self.inner.lock().unwrap();
190        inner.values.extend(
191            iter.into_iter()
192                .map(|(k, v)| (k.into(), v.into()))
193                .filter(|(k, v)| super::validate_os_pair((k, v)).is_ok()),
194        );
195    }
196}
197
198// Lua implementations
199
200impl FromLua for ProcessEnv {
201    fn from_lua(value: LuaValue, _: &Lua) -> LuaResult<Self> {
202        if let LuaValue::Nil = value {
203            Ok(Self::from_iter([] as [(OsString, OsString); 0]))
204        } else if let LuaValue::Boolean(true) = value {
205            Ok(Self::current())
206        } else if let Some(u) = value.as_userdata().and_then(|u| u.borrow::<Self>().ok()) {
207            Ok(u.clone())
208        } else if let LuaValue::Table(arr) = value {
209            let mut args = Vec::new();
210            for pair in arr.pairs::<LuaValue, LuaValue>() {
211                let (key_res, val_res) = match pair {
212                    Ok((key, val)) => (Ok(key), Ok(val)),
213                    Err(err) => (Err(err.clone()), Err(err)),
214                };
215
216                let key = super::lua_value_to_os_string(key_res, "ProcessEnv")?;
217                let val = super::lua_value_to_os_string(val_res, "ProcessEnv")?;
218
219                super::validate_os_pair((&key, &val))?;
220
221                args.push((key, val));
222            }
223            Ok(Self::from_iter(args))
224        } else {
225            Err(LuaError::FromLuaConversionError {
226                from: value.type_name(),
227                to: String::from("ProcessEnv"),
228                message: Some(format!(
229                    "Invalid type for process env - expected table or nil, got '{}'",
230                    value.type_name()
231                )),
232            })
233        }
234    }
235}
236
237impl LuaUserData for ProcessEnv {
238    fn add_methods<M: LuaUserDataMethods<Self>>(methods: &mut M) {
239        methods.add_meta_method(LuaMetaMethod::Len, |_, this, (): ()| Ok(this.len()));
240        methods.add_meta_method(LuaMetaMethod::Index, |_, this, key: LuaValue| {
241            let key = super::lua_value_to_os_string(Ok(key), "OsString")?;
242            Ok(this.get_value(key))
243        });
244        methods.add_meta_method(
245            LuaMetaMethod::NewIndex,
246            |_, this, (key, val): (LuaValue, Option<LuaValue>)| {
247                let key = super::lua_value_to_os_string(Ok(key), "OsString")?;
248                if let Some(val) = val {
249                    let val = super::lua_value_to_os_string(Ok(val), "OsString")?;
250                    this.set_value(key, val);
251                } else {
252                    this.remove_value(key);
253                }
254                Ok(())
255            },
256        );
257        methods.add_meta_method(LuaMetaMethod::Iter, |lua, this, (): ()| {
258            let mut vars = this
259                .clone()
260                .into_iter()
261                .filter_map(|(key, val)| Some((key.into_io_vec()?, val.into_io_vec()?)));
262            lua.create_function_mut(move |lua, (): ()| match vars.next() {
263                None => Ok((LuaValue::Nil, LuaValue::Nil)),
264                Some((key, val)) => Ok((
265                    LuaValue::String(lua.create_string(key)?),
266                    LuaValue::String(lua.create_string(val)?),
267                )),
268            })
269        });
270    }
271}
lune_utils/process/env.rs

lune_utils/process/
env.rs
