Skip to main content

dre_protocol/
sessions.rs

1//! `~/.dre/oauth_sessions.json`: OAuth sessions plugins keep between runs.
2//!
3//! One JSON object with an entry per login, keyed by the plugin as `<platform>/<host>/<client>`,
4//! so each workspace or account has its own entry and a source and destination on the same one
5//! share it. Plugins read their entry when they connect and write it after signing in or
6//! renewing a token. Only sessions go here: long-lived secrets stay in the environment.
7//!
8//! The file is only readable by its owner: mode 0600 on Unix, and on Windows the home folder's
9//! access list already limits it to the user. Writes lock the file, re-read it and replace it
10//! atomically, so two plugins saving different logins at once don't lose either.
11
12use std::fs::OpenOptions;
13use std::path::{Path, PathBuf};
14use std::time::{Duration, Instant};
15
16use serde_json::{Map, Value};
17
18pub const FILE: &str = "oauth_sessions.json";
19
20/// `~/.dre/oauth_sessions.json`.
21pub fn path() -> PathBuf {
22    std::env::home_dir().unwrap_or_default().join(".dre").join(FILE)
23}
24
25/// The session saved under `key`, if any.
26pub fn load(key: &str) -> Option<Value> {
27    read(&path()).remove(key)
28}
29
30/// Save (or with `None`, forget) the session under `key`.
31pub fn store(key: &str, session: Option<Value>) -> std::io::Result<()> {
32    store_at(&path(), key, session)
33}
34
35fn store_at(path: &Path, key: &str, session: Option<Value>) -> std::io::Result<()> {
36    if let Some(dir) = path.parent() {
37        std::fs::create_dir_all(dir)?;
38    }
39    let _lock = Lock::acquire(&path.with_extension("json.lock"));
40    let mut all = read(path);
41    match session {
42        Some(s) => all.insert(key.to_string(), s),
43        None => all.remove(key),
44    };
45    let bytes = serde_json::to_vec_pretty(&Value::Object(all)).unwrap_or_default();
46    crate::util::write_atomic_mode(path, &bytes, Some(0o600))
47}
48
49/// Every saved session; an unreadable or missing file is empty.
50fn read(path: &Path) -> Map<String, Value> {
51    std::fs::read_to_string(path)
52        .ok()
53        .and_then(|t| serde_json::from_str(&t).ok())
54        .unwrap_or_default()
55}
56
57/// A lock file held while the sessions file is rewritten. A lock older than a few seconds was
58/// left by a process that died, so it's taken over rather than waited on forever.
59struct Lock(Option<PathBuf>);
60
61impl Lock {
62    fn acquire(path: &Path) -> Lock {
63        let started = Instant::now();
64        loop {
65            if OpenOptions::new().write(true).create_new(true).open(path).is_ok() {
66                return Lock(Some(path.to_path_buf()));
67            }
68            let stale = std::fs::metadata(path)
69                .and_then(|m| m.modified())
70                .ok()
71                .and_then(|t| t.elapsed().ok())
72                .is_some_and(|age| age > Duration::from_secs(10));
73            if stale {
74                let _ = std::fs::remove_file(path);
75                continue;
76            }
77            if started.elapsed() > Duration::from_secs(10) {
78                // Better to write unlocked than to fail a run over a stuck lock.
79                return Lock(None);
80            }
81            std::thread::sleep(Duration::from_millis(20));
82        }
83    }
84}
85
86impl Drop for Lock {
87    fn drop(&mut self) {
88        if let Some(p) = &self.0 {
89            let _ = std::fs::remove_file(p);
90        }
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97    use serde_json::json;
98
99    #[test]
100    fn concurrent_logins_keep_every_entry_and_the_file_is_private() {
101        let dir = tempfile::tempdir().unwrap();
102        let path = dir.path().join(".dre").join(FILE);
103        std::thread::scope(|s| {
104            for i in 0..8 {
105                let path = &path;
106                s.spawn(move || {
107                    store_at(path, &format!("databricks/ws-{i}/c"), Some(json!({"n": i}))).unwrap()
108                });
109            }
110        });
111        let all = read(&path);
112        assert_eq!(all.len(), 8);
113        assert_eq!(all["databricks/ws-3/c"], json!({"n": 3}));
114        store_at(&path, "databricks/ws-3/c", None).unwrap();
115        assert!(!read(&path).contains_key("databricks/ws-3/c"));
116        assert!(!path.with_extension("json.lock").exists());
117        #[cfg(unix)]
118        {
119            use std::os::unix::fs::PermissionsExt;
120            let mode = std::fs::metadata(&path).unwrap().permissions().mode();
121            assert_eq!(mode & 0o777, 0o600);
122        }
123    }
124}