envfile/
lib.rs

1//! Libary for parsing environment files into an in-memory map.
2//!
3//! ```rust
4//! extern crate envfile;
5//!
6//! use envfile::EnvFile;
7//! use std::io;
8//! use std::path::Path;
9//!
10//! fn main() -> io::Result<()> {
11//!     let mut envfile = EnvFile::new(&Path::new("examples/test.env"))?;
12//!
13//!     for (key, value) in &envfile.store {
14//!         println!("{}: {}", key, value);
15//!     }
16//!
17//!     envfile.update("ID", "example");
18//!     println!("ID: {}", envfile.get("ID").unwrap_or(""));
19//!
20//!     // envfile.write()?;
21//!
22//!     Ok(())
23//! }
24//! ```
25
26extern crate snailquote;
27
28use std::collections::BTreeMap;
29use std::fs::File;
30use std::io::{self, Read, Write};
31use std::path::{Path, PathBuf};
32use std::str;
33
34use snailquote::{unescape, escape};
35
36
37/// An opened environment file, whose contents are buffered into memory.
38pub struct EnvFile {
39    /// Where the environment file exists in memory.
40    pub path:  PathBuf,
41    /// The data that was parsed from the file.
42    pub store: BTreeMap<String, String>,
43}
44
45fn parse_line(entry: &[u8]) -> Option<(String, String)> {
46    str::from_utf8(entry).ok().and_then(|l| {
47        let line = l.trim();
48        // Ignore comment line
49        if line.starts_with('#') {
50            return None;
51        }
52        let vline = line.as_bytes();
53        vline.iter().position(|&x| x == b'=').and_then(|pos| {
54            str::from_utf8(&vline[..pos]).ok().and_then(|x| {
55                str::from_utf8(&vline[pos+1..]).ok().and_then(|right| {
56                    // The right hand side value can be a quoted string
57                    unescape(right).ok().map(|y| (x.to_owned(), y))
58                })
59            })
60        })
61    })
62}
63
64impl EnvFile {
65    /// Open and parse an environment file.
66    pub fn new<P: Into<PathBuf>>(path: P) -> io::Result<Self> {
67        let path = path.into();
68        let data = read(&path)?;
69        let mut store = BTreeMap::new();
70
71        let values = data.split(|&x| x == b'\n').flat_map(parse_line);
72
73        for (key, value) in values {
74            store.insert(key, value);
75        }
76
77        Ok(EnvFile { path, store })
78    }
79
80    /// Update or insert a key into the map.
81    pub fn update(&mut self, key: &str, value: &str) -> &mut Self {
82        self.store.insert(key.into(), value.into());
83        self
84    }
85
86    /// Fetch a key from the map.
87    pub fn get(&self, key: &str) -> Option<&str> {
88        self.store.get(key).as_ref().map(|x| x.as_str())
89    }
90
91    /// Write the map back to the original file.
92    ///
93    /// # Notes
94    /// The keys are written in ascending order.
95    pub fn write(&mut self) -> io::Result<()> {
96        let mut buffer = Vec::with_capacity(1024);
97        for (key, value) in &self.store {
98            buffer.extend_from_slice(key.as_bytes());
99            buffer.push(b'=');
100            // The value may contain space and need to be quoted
101            let v = escape(value.as_str()).into_owned();
102            buffer.extend_from_slice(v.as_bytes());
103            buffer.push(b'\n');
104        }
105
106        write(&self.path, &buffer)
107    }
108}
109
110fn open<P: AsRef<Path>>(path: P) -> io::Result<File> {
111    File::open(&path).map_err(|why| io::Error::new(
112        io::ErrorKind::Other,
113        format!("unable to open file at {:?}: {}", path.as_ref(), why)
114    ))
115}
116
117fn create<P: AsRef<Path>>(path: P) -> io::Result<File> {
118    File::create(&path).map_err(|why| io::Error::new(
119        io::ErrorKind::Other,
120        format!("unable to create file at {:?}: {}", path.as_ref(), why)
121    ))
122}
123
124fn read<P: AsRef<Path>>(path: P) -> io::Result<Vec<u8>> {
125    open(path).and_then(|mut file| {
126        let mut buffer = Vec::with_capacity(file.metadata().ok().map_or(0, |x| x.len()) as usize);
127        file.read_to_end(&mut buffer).map(|_| buffer)
128    })
129}
130
131fn write<P: AsRef<Path>, C: AsRef<[u8]>>(path: P, contents: C) -> io::Result<()> {
132    create(path).and_then(|mut file| file.write_all(contents.as_ref()))
133}
134
135#[cfg(test)]
136mod tests {
137    extern crate tempdir;
138    use super::*;
139    use self::tempdir::TempDir;
140    use std::collections::BTreeMap;
141    use std::io::Write;
142
143    const SAMPLE: &str = r#"DOUBLE_QUOTED_STRING="This is a 'double-quoted' string"
144EFI_UUID=DFFD-D047
145HOSTNAME=pop-testing
146KBD_LAYOUT=us
147KBD_MODEL=
148KBD_VARIANT=
149 LANG=en_US.UTF-8
150OEM_MODE=0
151# Intentional blank line
152
153# Should ignore = operator in comment
154RECOVERY_UUID=PARTUUID=asdfasd7asdf7sad-asdfa
155ROOT_UUID=2ef950c2-5ce6-4ae0-9fb9-a8c7468fa82c
156SINGLE_QUOTED_STRING='This is a single-quoted string'
157"#;
158
159    const SAMPLE_CLEANED: &str = r#"DOUBLE_QUOTED_STRING="This is a 'double-quoted' string"
160EFI_UUID=DFFD-D047
161HOSTNAME=pop-testing
162KBD_LAYOUT=us
163KBD_MODEL=
164KBD_VARIANT=
165LANG=en_US.UTF-8
166OEM_MODE=0
167RECOVERY_UUID=PARTUUID=asdfasd7asdf7sad-asdfa
168ROOT_UUID=2ef950c2-5ce6-4ae0-9fb9-a8c7468fa82c
169SINGLE_QUOTED_STRING='This is a single-quoted string'
170"#;
171
172    #[test]
173    fn env_file_read() {
174        let tempdir = TempDir::new("distinst_test").unwrap();
175        let path = &tempdir.path().join("recovery.conf");
176
177        {
178            let mut file = create(path).unwrap();
179            file.write_all(SAMPLE.as_bytes()).unwrap();
180        }
181
182        let env = EnvFile::new(path).unwrap();
183        assert_eq!(&env.store, &{
184            let mut map = BTreeMap::new();
185            map.insert("HOSTNAME".into(), "pop-testing".into());
186            map.insert("LANG".into(), "en_US.UTF-8".into());
187            map.insert("KBD_LAYOUT".into(), "us".into());
188            map.insert("KBD_MODEL".into(), "".into());
189            map.insert("KBD_VARIANT".into(), "".into());
190            map.insert("EFI_UUID".into(), "DFFD-D047".into());
191            map.insert("RECOVERY_UUID".into(), "PARTUUID=asdfasd7asdf7sad-asdfa".into());
192            map.insert("ROOT_UUID".into(), "2ef950c2-5ce6-4ae0-9fb9-a8c7468fa82c".into());
193            map.insert("OEM_MODE".into(), "0".into());
194            map.insert("DOUBLE_QUOTED_STRING".into(), "This is a 'double-quoted' string".into());
195            map.insert("SINGLE_QUOTED_STRING".into(), "This is a single-quoted string".into());
196            map
197        });
198    }
199
200    #[test]
201    fn env_file_write() {
202        let tempdir = TempDir::new("distinst_test").unwrap();
203        let path = &tempdir.path().join("recovery.conf");
204
205        {
206            let mut file = create(path).unwrap();
207            file.write_all(SAMPLE.as_bytes()).unwrap();
208        }
209
210        let mut env = EnvFile::new(path).unwrap();
211        env.write().unwrap();
212        let copy: &[u8] = &read(path).unwrap();
213
214        assert_eq!(copy, SAMPLE_CLEANED.as_bytes(), "Expected '{}' == '{}'", String::from_utf8_lossy(copy), SAMPLE_CLEANED);
215    }
216}