Skip to main content

layer_client/
session_backend.rs

1//! Pluggable session storage backend.
2//!
3//! Built-in backends:
4//! * [`BinaryFileBackend`]  — compact binary file (default).
5//! * [`StringSessionBackend`] — portable base64 string (env vars, Docker, CI).
6//! * [`SqliteBackend`]      — SQLite (`sqlite-session` feature).
7//! * [`InMemoryBackend`]    — ephemeral, for tests / fresh-start bots.
8
9use std::io;
10use std::path::PathBuf;
11use crate::session::{CachedPeer, DcEntry, PersistedSession, UpdatesStateSnap};
12
13// ─── Trait ────────────────────────────────────────────────────────────────────
14
15pub trait SessionBackend: Send + Sync {
16    fn save(&self, session: &PersistedSession) -> io::Result<()>;
17    fn load(&self) -> io::Result<Option<PersistedSession>>;
18    fn delete(&self) -> io::Result<()>;
19    fn name(&self) -> &str;
20}
21
22// ─── BinaryFileBackend ────────────────────────────────────────────────────────
23
24/// Stores the session in a compact binary file (v2 format with update state + peer cache).
25pub struct BinaryFileBackend {
26    path: PathBuf,
27}
28
29impl BinaryFileBackend {
30    pub fn new(path: impl Into<PathBuf>) -> Self {
31        Self { path: path.into() }
32    }
33
34    /// Returns the path this backend writes to.
35    pub fn path(&self) -> &std::path::Path { &self.path }
36}
37
38impl SessionBackend for BinaryFileBackend {
39    fn save(&self, session: &PersistedSession) -> io::Result<()> {
40        session.save(&self.path)
41    }
42    fn load(&self) -> io::Result<Option<PersistedSession>> {
43        if !self.path.exists() { return Ok(None); }
44        PersistedSession::load(&self.path).map(Some)
45    }
46    fn delete(&self) -> io::Result<()> {
47        if self.path.exists() { std::fs::remove_file(&self.path)?; }
48        Ok(())
49    }
50    fn name(&self) -> &str { "binary-file" }
51}
52
53// ─── InMemoryBackend ─────────────────────────────────────────────────────────
54
55/// Ephemeral session — nothing persisted to disk.
56///
57/// Useful for tests or bots that always start fresh. Note that access-hash
58/// caches and update state are preserved across `save`/`load` calls *within
59/// the same process*, which is what the reconnect path needs.
60#[derive(Default)]
61pub struct InMemoryBackend {
62    data: std::sync::Mutex<Option<MemData>>,
63}
64
65#[derive(Clone)]
66struct MemData {
67    home_dc_id:    i32,
68    dcs:           Vec<DcEntry>,
69    updates_state: UpdatesStateSnap,
70    peers:         Vec<CachedPeer>,
71}
72
73impl InMemoryBackend {
74    pub fn new() -> Self { Self::default() }
75}
76
77impl SessionBackend for InMemoryBackend {
78    fn save(&self, s: &PersistedSession) -> io::Result<()> {
79        *self.data.lock().unwrap() = Some(MemData {
80            home_dc_id:    s.home_dc_id,
81            dcs:           s.dcs.clone(),
82            updates_state: s.updates_state.clone(),
83            peers:         s.peers.clone(),
84        });
85        Ok(())
86    }
87    fn load(&self) -> io::Result<Option<PersistedSession>> {
88        Ok(self.data.lock().unwrap().as_ref().map(|d| PersistedSession {
89            home_dc_id:    d.home_dc_id,
90            dcs:           d.dcs.clone(),
91            updates_state: d.updates_state.clone(),
92            peers:         d.peers.clone(),
93        }))
94    }
95    fn delete(&self) -> io::Result<()> {
96        *self.data.lock().unwrap() = None;
97        Ok(())
98    }
99    fn name(&self) -> &str { "in-memory" }
100}
101
102// ─── StringSessionBackend ─────────────────────────────────────────────────────
103
104/// Portable base64 string session backend — Pyrogram/Telethon style.
105///
106/// The session is stored as a single URL-safe base64 string with no padding.
107/// Useful for deployment pipelines, Docker, environment variables, and CI bots.
108///
109/// # Example
110/// ```rust,no_run
111/// let session_str = std::env::var("SESSION").unwrap();
112/// let backend = StringSessionBackend::new(&session_str);
113/// // ... pass to Config::session_backend
114///
115/// // Later, export the updated session:
116/// let updated = backend.current();
117/// println!("SESSION={updated}");
118/// ```
119pub struct StringSessionBackend {
120    data: std::sync::Mutex<String>,
121}
122
123impl StringSessionBackend {
124    /// Create a backend pre-loaded from an existing session string.
125    /// Pass an empty string to start a fresh session.
126    pub fn new(s: impl Into<String>) -> Self {
127        Self { data: std::sync::Mutex::new(s.into()) }
128    }
129
130    /// Return the current encoded session string.
131    /// Call after `client.save_session()` to get the latest value.
132    pub fn current(&self) -> String {
133        self.data.lock().unwrap().clone()
134    }
135}
136
137impl SessionBackend for StringSessionBackend {
138    fn save(&self, session: &PersistedSession) -> io::Result<()> {
139        *self.data.lock().unwrap() = session.to_string();
140        Ok(())
141    }
142
143    fn load(&self) -> io::Result<Option<PersistedSession>> {
144        let s = self.data.lock().unwrap().clone();
145        if s.trim().is_empty() {
146            return Ok(None);
147        }
148        PersistedSession::from_string(&s).map(Some)
149    }
150
151    fn delete(&self) -> io::Result<()> {
152        *self.data.lock().unwrap() = String::new();
153        Ok(())
154    }
155
156    fn name(&self) -> &str { "string-session" }
157}
158
159// ─── SqliteBackend ────────────────────────────────────────────────────────────
160
161#[cfg(feature = "sqlite-session")]
162pub use sqlite_backend::SqliteBackend;
163
164#[cfg(feature = "sqlite-session")]
165mod sqlite_backend {
166    use super::*;
167    use rusqlite::{Connection, params};
168
169    /// SQLite-backed session store.
170    ///
171    /// Schema (auto-created on first open):
172    /// - `meta`           — key/value string pairs (home_dc_id, pts, qts, date, seq)
173    /// - `dc_entries`     — one row per DC
174    /// - `channel_pts`    — per-channel pts counters
175    /// - `peers`          — cached access hashes
176    ///
177    /// Enable with the `sqlite-session` Cargo feature:
178    /// ```toml
179    /// layer-client = { version = "*", features = ["sqlite-session"] }
180    /// ```
181    pub struct SqliteBackend { path: PathBuf }
182
183    impl SqliteBackend {
184        pub fn new(path: impl Into<PathBuf>) -> io::Result<Self> {
185            let path = path.into();
186            let conn = Connection::open(&path)
187                .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
188            conn.execute_batch(
189                "CREATE TABLE IF NOT EXISTS meta (
190                    key   TEXT PRIMARY KEY,
191                    value TEXT NOT NULL
192                );
193                CREATE TABLE IF NOT EXISTS dc_entries (
194                    dc_id       INTEGER PRIMARY KEY,
195                    addr        TEXT    NOT NULL,
196                    auth_key    BLOB,
197                    first_salt  INTEGER NOT NULL DEFAULT 0,
198                    time_offset INTEGER NOT NULL DEFAULT 0
199                );
200                CREATE TABLE IF NOT EXISTS channel_pts (
201                    channel_id  INTEGER PRIMARY KEY,
202                    pts         INTEGER NOT NULL DEFAULT 0
203                );
204                CREATE TABLE IF NOT EXISTS peers (
205                    id          INTEGER PRIMARY KEY,
206                    access_hash INTEGER NOT NULL,
207                    is_channel  INTEGER NOT NULL DEFAULT 0
208                );",
209            ).map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
210            Ok(Self { path })
211        }
212
213        fn conn(&self) -> io::Result<Connection> {
214            Connection::open(&self.path)
215                .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
216        }
217    }
218
219    impl SessionBackend for SqliteBackend {
220        fn save(&self, s: &PersistedSession) -> io::Result<()> {
221            let conn = self.conn()?;
222            let e = |e: rusqlite::Error| io::Error::new(io::ErrorKind::Other, e);
223
224            // meta
225            for (k, v) in [
226                ("home_dc_id", s.home_dc_id.to_string()),
227                ("pts",  s.updates_state.pts.to_string()),
228                ("qts",  s.updates_state.qts.to_string()),
229                ("date", s.updates_state.date.to_string()),
230                ("seq",  s.updates_state.seq.to_string()),
231            ] {
232                conn.execute(
233                    "INSERT OR REPLACE INTO meta (key, value) VALUES (?1, ?2)",
234                    params![k, v],
235                ).map_err(e)?;
236            }
237
238            // dc_entries
239            for dc in &s.dcs {
240                conn.execute(
241                    "INSERT OR REPLACE INTO dc_entries
242                        (dc_id, addr, auth_key, first_salt, time_offset)
243                     VALUES (?1, ?2, ?3, ?4, ?5)",
244                    params![
245                        dc.dc_id,
246                        dc.addr,
247                        dc.auth_key.map(|k| k.to_vec()),
248                        dc.first_salt,
249                        dc.time_offset,
250                    ],
251                ).map_err(e)?;
252            }
253
254            // channel_pts
255            conn.execute_batch("DELETE FROM channel_pts").map_err(e)?;
256            for &(cid, cpts) in &s.updates_state.channels {
257                conn.execute(
258                    "INSERT INTO channel_pts (channel_id, pts) VALUES (?1, ?2)",
259                    params![cid, cpts],
260                ).map_err(e)?;
261            }
262
263            // peers
264            conn.execute_batch("DELETE FROM peers").map_err(e)?;
265            for p in &s.peers {
266                conn.execute(
267                    "INSERT INTO peers (id, access_hash, is_channel) VALUES (?1, ?2, ?3)",
268                    params![p.id, p.access_hash, p.is_channel as i32],
269                ).map_err(e)?;
270            }
271
272            Ok(())
273        }
274
275        fn load(&self) -> io::Result<Option<PersistedSession>> {
276            if !self.path.exists() { return Ok(None); }
277            let conn = self.conn()?;
278            let e = |err: rusqlite::Error| io::Error::new(io::ErrorKind::Other, err);
279
280            macro_rules! meta_i32 {
281                ($key:expr, $default:expr) => {
282                    conn.query_row(
283                        "SELECT value FROM meta WHERE key = ?1",
284                        params![$key],
285                        |row| row.get::<_, String>(0),
286                    ).ok()
287                    .and_then(|v| v.parse::<i32>().ok())
288                    .unwrap_or($default)
289                };
290            }
291
292            let home_dc_id = meta_i32!("home_dc_id", 0);
293            if home_dc_id == 0 { return Ok(None); }
294
295            // dc_entries
296            let mut stmt = conn.prepare(
297                "SELECT dc_id, addr, auth_key, first_salt, time_offset FROM dc_entries"
298            ).map_err(e)?;
299            let dcs: Vec<DcEntry> = stmt.query_map([], |row| {
300                let key_blob: Option<Vec<u8>> = row.get(2)?;
301                let auth_key = key_blob.and_then(|k| {
302                    if k.len() == 256 {
303                        let mut a = [0u8; 256]; a.copy_from_slice(&k); Some(a)
304                    } else { None }
305                });
306                Ok(DcEntry {
307                    dc_id:       row.get(0)?,
308                    addr:        row.get(1)?,
309                    auth_key,
310                    first_salt:  row.get(3)?,
311                    time_offset: row.get(4)?,
312                })
313            }).map_err(e)?.filter_map(|r| r.ok()).collect();
314
315            // update state
316            let pts  = meta_i32!("pts",  0);
317            let qts  = meta_i32!("qts",  0);
318            let date = meta_i32!("date", 0);
319            let seq  = meta_i32!("seq",  0);
320
321            let mut ch_stmt = conn.prepare(
322                "SELECT channel_id, pts FROM channel_pts"
323            ).map_err(e)?;
324            let channels: Vec<(i64, i32)> = ch_stmt
325                .query_map([], |row| Ok((row.get::<_, i64>(0)?, row.get::<_, i32>(1)?)))
326                .map_err(e)?.filter_map(|r| r.ok()).collect();
327
328            // peers
329            let mut peer_stmt = conn.prepare(
330                "SELECT id, access_hash, is_channel FROM peers"
331            ).map_err(e)?;
332            let peers: Vec<CachedPeer> = peer_stmt
333                .query_map([], |row| Ok(CachedPeer {
334                    id:          row.get(0)?,
335                    access_hash: row.get(1)?,
336                    is_channel:  row.get::<_, i32>(2)? != 0,
337                }))
338                .map_err(e)?.filter_map(|r| r.ok()).collect();
339
340            Ok(Some(PersistedSession {
341                home_dc_id,
342                dcs,
343                updates_state: UpdatesStateSnap { pts, qts, date, seq, channels },
344                peers,
345            }))
346        }
347
348        fn delete(&self) -> io::Result<()> {
349            if self.path.exists() { std::fs::remove_file(&self.path)?; }
350            Ok(())
351        }
352
353        fn name(&self) -> &str { "sqlite" }
354    }
355}