1
//! Filesystem + JSON implementation of StateMgr.
2

            
3
use crate::{load_error, store_error};
4
use crate::{Error, LockStatus, Result, StateMgr};
5
use serde::{de::DeserializeOwned, Serialize};
6
use std::path::{Path, PathBuf};
7
use std::sync::{Arc, Mutex};
8

            
9
#[cfg(target_family = "unix")]
10
use std::os::unix::fs::DirBuilderExt;
11

            
12
/// Implementation of StateMgr that stores state as JSON files on disk.
13
///
14
/// # Locking
15
///
16
/// This manager uses a lock file to determine whether it's allowed to
17
/// write to the disk.  Only one process should write to the disk at
18
/// a time, though any number may read from the disk.
19
///
20
/// By default, every `FsStateMgr` starts out unlocked, and only able
21
/// to read.  Use [`FsStateMgr::try_lock()`] to lock it.
22
///
23
/// # Limitations
24
///
25
/// 1) This manager only accepts objects that can be serialized as
26
/// JSON documents.  Some types (like maps with non-string keys) can't
27
/// be serialized as JSON.
28
///
29
/// 2) This manager normalizes keys to an fs-safe format before saving
30
/// data with them.  This keeps you from accidentally creating or
31
/// reading files elsewhere in the filesystem, but it doesn't prevent
32
/// collisions when two keys collapse to the same fs-safe filename.
33
/// Therefore, you should probably only use ascii keys that are
34
/// fs-safe on all systems.
35
///
36
/// NEVER use user-controlled or remote-controlled data for your keys.
37
6
#[derive(Clone, Debug)]
38
pub struct FsStateMgr {
39
    /// Inner reference-counted object.
40
    inner: Arc<FsStateMgrInner>,
41
}
42

            
43
/// Inner reference-counted object, used by `FsStateMgr`.
44
#[derive(Debug)]
45
struct FsStateMgrInner {
46
    /// Directory in which we store state files.
47
    statepath: PathBuf,
48
    /// Lockfile to achieve exclusive access to state files.
49
    lockfile: Mutex<fslock::LockFile>,
50
}
51

            
52
impl FsStateMgr {
53
    /// Construct a new `FsStateMgr` to store data in `path`.
54
    ///
55
    /// This function will try to create `path` if it does not already
56
    /// exist.
57
4
    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<Self> {
58
4
        let path = path.as_ref();
59
4
        let statepath = path.join("state");
60
4
        let lockpath = path.join("state.lock");
61
4

            
62
4
        {
63
4
            let mut builder = std::fs::DirBuilder::new();
64
4
            #[cfg(target_family = "unix")]
65
4
            builder.mode(0o700);
66
4
            builder.recursive(true).create(&statepath)?;
67
        }
68

            
69
4
        let lockfile = Mutex::new(fslock::LockFile::open(&lockpath)?);
70

            
71
4
        Ok(FsStateMgr {
72
4
            inner: Arc::new(FsStateMgrInner {
73
4
                statepath,
74
4
                lockfile,
75
4
            }),
76
4
        })
77
4
    }
78
    /// Return a filename to use for storing data with `key`.
79
    ///
80
    /// See "Limitations" section on [`FsStateMgr`] for caveats.
81
134
    fn filename(&self, key: &str) -> PathBuf {
82
134
        self.inner
83
134
            .statepath
84
134
            .join(sanitize_filename::sanitize(key) + ".json")
85
134
    }
86
    /// Return the top-level directory for this storage manager.
87
    ///
88
    /// (This is the same directory passed to [`FsStateMgr::from_path`].)
89
1
    pub fn path(&self) -> &Path {
90
1
        self.inner
91
1
            .statepath
92
1
            .parent()
93
1
            .expect("No parent directory even after path.join?")
94
1
    }
95
}
96

            
97
impl StateMgr for FsStateMgr {
98
99
    fn can_store(&self) -> bool {
99
99
        let lockfile = self
100
99
            .inner
101
99
            .lockfile
102
99
            .lock()
103
99
            .expect("Poisoned lock on state lockfile");
104
99
        lockfile.owns_lock()
105
99
    }
106
34
    fn try_lock(&self) -> Result<LockStatus> {
107
34
        let mut lockfile = self
108
34
            .inner
109
34
            .lockfile
110
34
            .lock()
111
34
            .expect("Poisoned lock on state lockfile");
112
34
        if lockfile.owns_lock() {
113
            Ok(LockStatus::AlreadyHeld)
114
34
        } else if lockfile.try_lock()? {
115
34
            Ok(LockStatus::NewlyAcquired)
116
        } else {
117
            Ok(LockStatus::NoLock)
118
        }
119
34
    }
120
    fn unlock(&self) -> Result<()> {
121
        let mut lockfile = self
122
            .inner
123
            .lockfile
124
            .lock()
125
            .expect("Poisoned lock on state lockfile");
126
        if lockfile.owns_lock() {
127
            lockfile.unlock()?;
128
        }
129
        Ok(())
130
    }
131
12
    fn load<D>(&self, key: &str) -> Result<Option<D>>
132
12
    where
133
12
        D: DeserializeOwned,
134
12
    {
135
12
        let fname = self.filename(key);
136

            
137
12
        let string = match std::fs::read_to_string(fname) {
138
3
            Ok(s) => s,
139
9
            Err(e) => {
140
9
                if e.kind() == std::io::ErrorKind::NotFound {
141
9
                    return Ok(None);
142
                } else {
143
                    return Err(e.into());
144
                }
145
            }
146
        };
147

            
148
3
        Ok(Some(serde_json::from_str(&string).map_err(load_error)?))
149
12
    }
150

            
151
3
    fn store<S>(&self, key: &str, val: &S) -> Result<()>
152
3
    where
153
3
        S: Serialize,
154
3
    {
155
3
        if !self.can_store() {
156
1
            return Err(Error::NoLock);
157
2
        }
158
2

            
159
2
        let fname = self.filename(key);
160

            
161
2
        let output = serde_json::to_string_pretty(val).map_err(store_error)?;
162

            
163
2
        let fname_tmp = fname.with_extension("tmp");
164
2
        std::fs::write(&fname_tmp, &output)?;
165
2
        std::fs::rename(fname_tmp, fname)?;
166

            
167
2
        Ok(())
168
3
    }
169
}
170

            
171
#[cfg(test)]
172
mod test {
173
    #![allow(clippy::unwrap_used)]
174
    use super::*;
175
    use std::collections::HashMap;
176

            
177
    #[test]
178
    fn simple() -> Result<()> {
179
        let dir = tempfile::TempDir::new().unwrap();
180
        let store = FsStateMgr::from_path(dir.path())?;
181

            
182
        assert_eq!(store.try_lock()?, LockStatus::NewlyAcquired);
183
        let stuff: HashMap<_, _> = vec![("hello".to_string(), "world".to_string())]
184
            .into_iter()
185
            .collect();
186
        store.store("xyz", &stuff)?;
187

            
188
        let stuff2: Option<HashMap<String, String>> = store.load("xyz")?;
189
        let nothing: Option<HashMap<String, String>> = store.load("abc")?;
190

            
191
        assert_eq!(Some(stuff), stuff2);
192
        assert!(nothing.is_none());
193

            
194
        assert_eq!(dir.path(), store.path());
195

            
196
        drop(store); // Do this to release the fs lock.
197
        let store = FsStateMgr::from_path(dir.path())?;
198
        let stuff3: Option<HashMap<String, String>> = store.load("xyz")?;
199
        assert_eq!(stuff2, stuff3);
200

            
201
        let stuff4: HashMap<_, _> = vec![("greetings".to_string(), "humans".to_string())]
202
            .into_iter()
203
            .collect();
204

            
205
        assert!(matches!(store.store("xyz", &stuff4), Err(Error::NoLock)));
206

            
207
        assert_eq!(store.try_lock()?, LockStatus::NewlyAcquired);
208
        store.store("xyz", &stuff4)?;
209

            
210
        let stuff5: Option<HashMap<String, String>> = store.load("xyz")?;
211
        assert_eq!(Some(stuff4), stuff5);
212

            
213
        Ok(())
214
    }
215
}