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::io::Write;
14use std::path::{Path, PathBuf};
15use std::time::{Duration, Instant};
16
17use serde_json::{Map, Value};
18
19pub const FILE: &str = "oauth_sessions.json";
20
21/// `~/.dre/oauth_sessions.json`.
22pub fn path() -> PathBuf {
23    std::env::home_dir().unwrap_or_default().join(".dre").join(FILE)
24}
25
26/// The session saved under `key`, if any.
27pub fn load(key: &str) -> Option<Value> {
28    read(&path()).remove(key)
29}
30
31/// Save (or with `None`, forget) the session under `key`.
32pub fn store(key: &str, session: Option<Value>) -> std::io::Result<()> {
33    store_at(&path(), key, session)
34}
35
36fn store_at(path: &Path, key: &str, session: Option<Value>) -> std::io::Result<()> {
37    if let Some(dir) = path.parent() {
38        std::fs::create_dir_all(dir)?;
39    }
40    let _lock = Lock::acquire(&path.with_extension("json.lock"));
41    let mut all = read(path);
42    match session {
43        Some(s) => all.insert(key.to_string(), s),
44        None => all.remove(key),
45    };
46    let bytes = serde_json::to_vec_pretty(&Value::Object(all)).unwrap_or_default();
47    let tmp = path.with_extension(format!("json.{}.tmp", std::process::id()));
48    let mut opts = OpenOptions::new();
49    opts.write(true).create(true).truncate(true);
50    #[cfg(unix)]
51    {
52        use std::os::unix::fs::OpenOptionsExt;
53        opts.mode(0o600);
54    }
55    opts.open(&tmp)?.write_all(&bytes)?;
56    std::fs::rename(&tmp, path)
57}
58
59/// Every saved session; an unreadable or missing file is empty.
60fn read(path: &Path) -> Map<String, Value> {
61    std::fs::read_to_string(path)
62        .ok()
63        .and_then(|t| serde_json::from_str(&t).ok())
64        .unwrap_or_default()
65}
66
67/// A lock file held while the sessions file is rewritten. A lock older than a few seconds was
68/// left by a process that died, so it's taken over rather than waited on forever.
69struct Lock(Option<PathBuf>);
70
71impl Lock {
72    fn acquire(path: &Path) -> Lock {
73        let started = Instant::now();
74        loop {
75            if OpenOptions::new().write(true).create_new(true).open(path).is_ok() {
76                return Lock(Some(path.to_path_buf()));
77            }
78            let stale = std::fs::metadata(path)
79                .and_then(|m| m.modified())
80                .ok()
81                .and_then(|t| t.elapsed().ok())
82                .is_some_and(|age| age > Duration::from_secs(10));
83            if stale {
84                let _ = std::fs::remove_file(path);
85                continue;
86            }
87            if started.elapsed() > Duration::from_secs(10) {
88                // Better to write unlocked than to fail a run over a stuck lock.
89                return Lock(None);
90            }
91            std::thread::sleep(Duration::from_millis(20));
92        }
93    }
94}
95
96impl Drop for Lock {
97    fn drop(&mut self) {
98        if let Some(p) = &self.0 {
99            let _ = std::fs::remove_file(p);
100        }
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107    use serde_json::json;
108
109    #[test]
110    fn concurrent_logins_keep_every_entry_and_the_file_is_private() {
111        let dir = tempfile::tempdir().unwrap();
112        let path = dir.path().join(".dre").join(FILE);
113        std::thread::scope(|s| {
114            for i in 0..8 {
115                let path = &path;
116                s.spawn(move || {
117                    store_at(path, &format!("databricks/ws-{i}/c"), Some(json!({"n": i}))).unwrap()
118                });
119            }
120        });
121        let all = read(&path);
122        assert_eq!(all.len(), 8);
123        assert_eq!(all["databricks/ws-3/c"], json!({"n": 3}));
124        store_at(&path, "databricks/ws-3/c", None).unwrap();
125        assert!(!read(&path).contains_key("databricks/ws-3/c"));
126        assert!(!path.with_extension("json.lock").exists());
127        #[cfg(unix)]
128        {
129            use std::os::unix::fs::PermissionsExt;
130            let mode = std::fs::metadata(&path).unwrap().permissions().mode();
131            assert_eq!(mode & 0o777, 0o600);
132        }
133    }
134}
dre_protocol/sessions.rs

dre_protocol/
sessions.rs
