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.
59#[must_use = "dropping the lock releases it at once"]
60struct Lock(Option<PathBuf>);
61
62impl Lock {
63    fn acquire(path: &Path) -> Lock {
64        let started = Instant::now();
65        loop {
66            if OpenOptions::new().write(true).create_new(true).open(path).is_ok() {
67                return Lock(Some(path.to_path_buf()));
68            }
69            let stale = std::fs::metadata(path)
70                .and_then(|m| m.modified())
71                .ok()
72                .and_then(|t| t.elapsed().ok())
73                .is_some_and(|age| age > Duration::from_secs(10));
74            if stale {
75                let _ = std::fs::remove_file(path);
76                continue;
77            }
78            if started.elapsed() > Duration::from_secs(10) {
79                // Better to write unlocked than to fail a run over a stuck lock.
80                return Lock(None);
81            }
82            std::thread::sleep(Duration::from_millis(20));
83        }
84    }
85}
86
87impl Drop for Lock {
88    fn drop(&mut self) {
89        if let Some(p) = &self.0 {
90            let _ = std::fs::remove_file(p);
91        }
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use serde_json::json;
99
100    #[test]
101    fn concurrent_logins_keep_every_entry_and_the_file_is_private() {
102        let dir = tempfile::tempdir().unwrap();
103        let path = dir.path().join(".dre").join(FILE);
104        std::thread::scope(|s| {
105            for i in 0..8 {
106                let path = &path;
107                s.spawn(move || {
108                    store_at(path, &format!("databricks/ws-{i}/c"), Some(json!({"n": i}))).unwrap()
109                });
110            }
111        });
112        let all = read(&path);
113        assert_eq!(all.len(), 8);
114        assert_eq!(all["databricks/ws-3/c"], json!({"n": 3}));
115        store_at(&path, "databricks/ws-3/c", None).unwrap();
116        assert!(!read(&path).contains_key("databricks/ws-3/c"));
117        assert!(!path.with_extension("json.lock").exists());
118        #[cfg(unix)]
119        {
120            use std::os::unix::fs::PermissionsExt;
121            let mode = std::fs::metadata(&path).unwrap().permissions().mode();
122            assert_eq!(mode & 0o777, 0o600);
123        }
124    }
125}
dre_protocol/sessions.rs

dre_protocol/
sessions.rs
