1#![deny(unsafe_code)]
14#![cfg_attr(docsrs, feature(doc_cfg))]
15#![doc(html_root_url = "https://docs.rs/ferogram-session/0.6.4")]
16use std::collections::HashMap;
79use std::io::{self, ErrorKind};
80use std::path::Path;
81
82#[cfg(feature = "serde")]
83mod auth_key_serde {
84 use serde::{Deserialize, Deserializer, Serializer};
85
86 pub fn serialize<S>(value: &Option<[u8; 256]>, s: S) -> Result<S::Ok, S::Error>
89 where
90 S: Serializer,
91 {
92 match value {
93 Some(k) => s.serialize_some(k.as_slice()),
94 None => s.serialize_none(),
95 }
96 }
97
98 pub fn deserialize<'de, D>(d: D) -> Result<Option<[u8; 256]>, D::Error>
101 where
102 D: Deserializer<'de>,
103 {
104 let opt: Option<Vec<u8>> = Option::deserialize(d)?;
105 match opt {
106 None => Ok(None),
107 Some(v) => {
108 let arr: [u8; 256] = v
109 .try_into()
110 .map_err(|_| serde::de::Error::custom("auth_key must be exactly 256 bytes"))?;
111 Ok(Some(arr))
112 }
113 }
114 }
115}
116
117#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
121#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
122pub struct DcFlags(pub u8);
123
124impl DcFlags {
125 pub const NONE: DcFlags = DcFlags(0);
126 pub const IPV6: DcFlags = DcFlags(1 << 0);
127 pub const MEDIA_ONLY: DcFlags = DcFlags(1 << 1);
128 pub const TCPO_ONLY: DcFlags = DcFlags(1 << 2);
129 pub const CDN: DcFlags = DcFlags(1 << 3);
130 pub const STATIC: DcFlags = DcFlags(1 << 4);
131
132 pub fn contains(self, other: DcFlags) -> bool {
134 self.0 & other.0 == other.0
135 }
136
137 pub fn set(&mut self, flag: DcFlags) {
139 self.0 |= flag.0;
140 }
141}
142
143impl std::ops::BitOr for DcFlags {
144 type Output = DcFlags;
145 fn bitor(self, rhs: DcFlags) -> DcFlags {
146 DcFlags(self.0 | rhs.0)
147 }
148}
149
150#[derive(Clone, Debug)]
152#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
153pub struct DcEntry {
154 pub dc_id: i32,
155 pub addr: String,
156 #[cfg_attr(feature = "serde", serde(with = "auth_key_serde"))]
157 pub auth_key: Option<[u8; 256]>,
158 pub first_salt: i64,
159 pub time_offset: i32,
160 pub flags: DcFlags,
162}
163
164impl DcEntry {
165 #[inline]
167 pub fn is_ipv6(&self) -> bool {
168 self.flags.contains(DcFlags::IPV6)
169 }
170
171 pub fn socket_addr(&self) -> io::Result<std::net::SocketAddr> {
178 self.addr.parse::<std::net::SocketAddr>().map_err(|_| {
179 io::Error::new(
180 io::ErrorKind::InvalidData,
181 format!("invalid DC address: {:?}", self.addr),
182 )
183 })
184 }
185
186 pub fn from_parts(dc_id: i32, ip: &str, port: u16, flags: DcFlags) -> Self {
199 let addr = if ip.contains(':') {
201 format!("[{ip}]:{port}")
202 } else {
203 format!("{ip}:{port}")
204 };
205 Self {
206 dc_id,
207 addr,
208 auth_key: None,
209 first_salt: 0,
210 time_offset: 0,
211 flags,
212 }
213 }
214}
215
216#[derive(Clone, Debug, Default)]
219#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
220pub struct UpdatesStateSnap {
221 pub pts: i32,
223 pub qts: i32,
225 pub date: i32,
227 pub seq: i32,
229 pub channels: Vec<(i64, i32)>,
231}
232
233impl UpdatesStateSnap {
234 #[inline]
236 pub fn is_initialised(&self) -> bool {
237 self.pts > 0
238 }
239
240 pub fn set_channel_pts(&mut self, channel_id: i64, pts: i32) {
242 if let Some(entry) = self.channels.iter_mut().find(|c| c.0 == channel_id) {
243 entry.1 = pts;
244 } else {
245 self.channels.push((channel_id, pts));
246 }
247 }
248
249 pub fn channel_pts(&self, channel_id: i64) -> i32 {
251 self.channels
252 .iter()
253 .find(|c| c.0 == channel_id)
254 .map(|c| c.1)
255 .unwrap_or(0)
256 }
257}
258
259#[derive(Clone, Copy, Debug, PartialEq, Eq)]
262#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
263#[repr(u8)]
264pub enum ChannelKind {
265 Broadcast = 0,
267 Megagroup = 1,
269 Gigagroup = 2,
271}
272
273impl ChannelKind {
274 pub fn from_byte(b: u8) -> Option<Self> {
277 match b {
278 0 => Some(Self::Broadcast),
279 1 => Some(Self::Megagroup),
280 2 => Some(Self::Gigagroup),
281 _ => None,
282 }
283 }
284
285 pub fn to_byte(self) -> u8 {
287 self as u8
288 }
289}
290
291#[derive(Clone, Debug)]
294#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
295pub struct CachedPeer {
296 pub id: i64,
298 pub access_hash: i64,
301 pub is_channel: bool,
303 pub is_chat: bool,
306 pub channel_kind: Option<ChannelKind>,
309}
310
311#[derive(Clone, Debug)]
315#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
316pub struct CachedMinPeer {
317 pub user_id: i64,
319 pub peer_id: i64,
321 pub msg_id: i32,
323}
324
325#[derive(Clone, Debug, Default)]
327#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
328pub struct PersistedSession {
329 pub home_dc_id: i32,
330 pub dcs: Vec<DcEntry>,
331 pub updates_state: UpdatesStateSnap,
333 pub peers: Vec<CachedPeer>,
336 pub min_peers: Vec<CachedMinPeer>,
339 pub future_auth_token: Option<Vec<u8>>,
346}
347
348impl PersistedSession {
349 pub fn to_bytes(&self) -> Vec<u8> {
351 let mut b = Vec::with_capacity(512);
352
353 b.push(0x07u8); b.extend_from_slice(&self.home_dc_id.to_le_bytes());
356
357 b.push(self.dcs.len() as u8);
358 for d in &self.dcs {
359 b.extend_from_slice(&d.dc_id.to_le_bytes());
360 match &d.auth_key {
361 Some(k) => {
362 b.push(1);
363 b.extend_from_slice(k);
364 }
365 None => {
366 b.push(0);
367 }
368 }
369 b.extend_from_slice(&d.first_salt.to_le_bytes());
370 b.extend_from_slice(&d.time_offset.to_le_bytes());
371 let ab = d.addr.as_bytes();
372 b.push(ab.len() as u8);
373 b.extend_from_slice(ab);
374 b.push(d.flags.0);
375 }
376
377 b.extend_from_slice(&self.updates_state.pts.to_le_bytes());
378 b.extend_from_slice(&self.updates_state.qts.to_le_bytes());
379 b.extend_from_slice(&self.updates_state.date.to_le_bytes());
380 b.extend_from_slice(&self.updates_state.seq.to_le_bytes());
381 let ch = &self.updates_state.channels;
382 b.extend_from_slice(&(ch.len() as u16).to_le_bytes());
383 for &(cid, cpts) in ch {
384 b.extend_from_slice(&cid.to_le_bytes());
385 b.extend_from_slice(&cpts.to_le_bytes());
386 }
387
388 b.extend_from_slice(&(self.peers.len() as u16).to_le_bytes());
390 for p in &self.peers {
391 b.extend_from_slice(&p.id.to_le_bytes());
392 b.extend_from_slice(&p.access_hash.to_le_bytes());
393 let peer_type: u8 = if p.is_chat {
394 2
395 } else if p.is_channel {
396 1
397 } else {
398 0
399 };
400 b.push(peer_type);
401 b.push(p.channel_kind.map(|k| k.to_byte()).unwrap_or(0xFF));
403 }
404
405 b.extend_from_slice(&(self.min_peers.len() as u16).to_le_bytes());
406 for m in &self.min_peers {
407 b.extend_from_slice(&m.user_id.to_le_bytes());
408 b.extend_from_slice(&m.peer_id.to_le_bytes());
409 b.extend_from_slice(&m.msg_id.to_le_bytes());
410 }
411
412 match &self.future_auth_token {
414 Some(t) => {
415 b.push(1);
416 b.extend_from_slice(&(t.len() as u16).to_le_bytes());
417 b.extend_from_slice(t);
418 }
419 None => b.push(0),
420 }
421
422 b
423 }
424
425 pub fn save(&self, path: &Path) -> io::Result<()> {
432 use std::sync::atomic::{AtomicU64, Ordering};
433 static SEQ: AtomicU64 = AtomicU64::new(0);
434 let n = SEQ.fetch_add(1, Ordering::Relaxed);
435 let tmp = path.with_extension(format!("{n}.tmp"));
436 std::fs::write(&tmp, self.to_bytes())?;
437 std::fs::rename(&tmp, path).inspect_err(|_e| {
438 let _ = std::fs::remove_file(&tmp);
439 })
440 }
441
442 pub fn from_bytes(buf: &[u8]) -> io::Result<Self> {
444 if buf.is_empty() {
445 return Err(io::Error::new(ErrorKind::InvalidData, "empty session data"));
446 }
447
448 let mut p = 0usize;
449
450 macro_rules! r {
451 ($n:expr) => {{
452 if p + $n > buf.len() {
453 return Err(io::Error::new(ErrorKind::InvalidData, "truncated session"));
454 }
455 let s = &buf[p..p + $n];
456 p += $n;
457 s
458 }};
459 }
460 macro_rules! r_i32 {
461 () => {
462 i32::from_le_bytes(r!(4).try_into().unwrap())
463 };
464 }
465 macro_rules! r_i64 {
466 () => {
467 i64::from_le_bytes(r!(8).try_into().unwrap())
468 };
469 }
470 macro_rules! r_u8 {
471 () => {
472 r!(1)[0]
473 };
474 }
475 macro_rules! r_u16 {
476 () => {
477 u16::from_le_bytes(r!(2).try_into().unwrap())
478 };
479 }
480
481 let first_byte = r_u8!();
482
483 let (home_dc_id, version) = if first_byte == 0x07 {
484 (r_i32!(), 7u8)
485 } else if first_byte == 0x06 {
486 (r_i32!(), 6u8)
487 } else if first_byte == 0x05 {
488 (r_i32!(), 5u8)
489 } else if first_byte == 0x04 {
490 (r_i32!(), 4u8)
491 } else if first_byte == 0x03 {
492 (r_i32!(), 3u8)
493 } else if first_byte == 0x02 {
494 (r_i32!(), 2u8)
495 } else {
496 let rest = r!(3);
497 let mut bytes = [0u8; 4];
498 bytes[0] = first_byte;
499 bytes[1..4].copy_from_slice(rest);
500 (i32::from_le_bytes(bytes), 1u8)
501 };
502
503 let dc_count = r_u8!() as usize;
504 let mut dcs = Vec::with_capacity(dc_count);
505 for _ in 0..dc_count {
506 let dc_id = r_i32!();
507 let has_key = r_u8!();
508 let auth_key = if has_key == 1 {
509 let mut k = [0u8; 256];
510 k.copy_from_slice(r!(256));
511 Some(k)
512 } else {
513 None
514 };
515 let first_salt = r_i64!();
516 let time_offset = r_i32!();
517 let al = r_u8!() as usize;
518 let addr = String::from_utf8_lossy(r!(al)).into_owned();
519 let flags = if version >= 3 {
520 DcFlags(r_u8!())
521 } else {
522 DcFlags::NONE
523 };
524 dcs.push(DcEntry {
525 dc_id,
526 addr,
527 auth_key,
528 first_salt,
529 time_offset,
530 flags,
531 });
532 }
533
534 if version < 2 {
535 return Ok(Self {
536 home_dc_id,
537 dcs,
538 updates_state: UpdatesStateSnap::default(),
539 peers: Vec::new(),
540 min_peers: Vec::new(),
541 future_auth_token: None,
542 });
543 }
544
545 let pts = r_i32!();
546 let qts = r_i32!();
547 let date = r_i32!();
548 let seq = r_i32!();
549 let ch_count = r_u16!() as usize;
550 let mut channels = Vec::with_capacity(ch_count);
551 for _ in 0..ch_count {
552 let cid = r_i64!();
553 let cpts = r_i32!();
554 channels.push((cid, cpts));
555 }
556
557 let peer_count = r_u16!() as usize;
558 let mut peers = Vec::with_capacity(peer_count);
559 for _ in 0..peer_count {
560 let id = r_i64!();
561 let access_hash = r_i64!();
562 let peer_type = r_u8!();
564 let is_channel = peer_type == 1;
565 let is_chat = peer_type == 2;
566 let channel_kind = if version >= 6 {
568 let kb = r_u8!();
569 if kb == 0xFF {
570 None
571 } else {
572 ChannelKind::from_byte(kb)
573 }
574 } else {
575 None
576 };
577 peers.push(CachedPeer {
578 id,
579 access_hash,
580 is_channel,
581 is_chat,
582 channel_kind,
583 });
584 }
585
586 let min_peers = if version >= 4 {
588 let count = r_u16!() as usize;
589 let mut v = Vec::with_capacity(count);
590 for _ in 0..count {
591 let user_id = r_i64!();
592 let peer_id = r_i64!();
593 let msg_id = r_i32!();
594 v.push(CachedMinPeer {
595 user_id,
596 peer_id,
597 msg_id,
598 });
599 }
600 v
601 } else {
602 Vec::new()
603 };
604
605 let future_auth_token = if version >= 7 {
607 let has_token = r_u8!();
608 if has_token == 1 {
609 let tl = r_u16!() as usize;
610 Some(r!(tl).to_vec())
611 } else {
612 None
613 }
614 } else {
615 None
616 };
617 let _ = p;
621
622 Ok(Self {
623 home_dc_id,
624 dcs,
625 updates_state: UpdatesStateSnap {
626 pts,
627 qts,
628 date,
629 seq,
630 channels,
631 },
632 peers,
633 min_peers,
634 future_auth_token,
635 })
636 }
637
638 pub fn from_string(s: &str) -> io::Result<Self> {
640 use base64::Engine as _;
641 let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
642 .decode(s.trim())
643 .map_err(|e| io::Error::new(ErrorKind::InvalidData, e))?;
644 Self::from_bytes(&bytes)
645 }
646
647 pub fn load(path: &Path) -> io::Result<Self> {
650 let buf = std::fs::read(path)?;
651 Self::from_bytes(&buf)
652 }
653
654 pub fn dc_for(&self, dc_id: i32, prefer_ipv6: bool) -> Option<&DcEntry> {
666 let mut candidates = self.dcs.iter().filter(|d| d.dc_id == dc_id).peekable();
667 candidates.peek()?;
668 let cands: Vec<&DcEntry> = self.dcs.iter().filter(|d| d.dc_id == dc_id).collect();
670 cands
672 .iter()
673 .copied()
674 .find(|d| d.is_ipv6() == prefer_ipv6)
675 .or_else(|| cands.first().copied())
676 }
677
678 pub fn all_dcs_for(&self, dc_id: i32) -> impl Iterator<Item = &DcEntry> {
683 self.dcs.iter().filter(move |d| d.dc_id == dc_id)
684 }
685
686 pub fn stats(&self) -> SessionStats {
695 let users = self
696 .peers
697 .iter()
698 .filter(|p| !p.is_channel && !p.is_chat)
699 .count();
700 let channels = self.peers.iter().filter(|p| p.is_channel).count();
701 let chats = self.peers.iter().filter(|p| p.is_chat).count();
702
703 SessionStats {
704 home_dc_id: self.home_dc_id,
705 dcs_total: self.dcs.len(),
706 dcs_with_auth_key: self.dcs.iter().filter(|d| d.auth_key.is_some()).count(),
707 peers_total: self.peers.len(),
708 peers_users: users,
709 peers_channels: channels,
710 peers_chats: chats,
711 min_peers: self.min_peers.len(),
712 channel_pts_entries: self.updates_state.channels.len(),
713 updates_initialised: self.updates_state.is_initialised(),
714 approx_bytes: self.to_bytes().len(),
715 }
716 }
717}
718
719#[derive(Clone, Debug)]
723pub struct SessionStats {
724 pub home_dc_id: i32,
726 pub dcs_total: usize,
728 pub dcs_with_auth_key: usize,
730 pub peers_total: usize,
732 pub peers_users: usize,
734 pub peers_channels: usize,
736 pub peers_chats: usize,
738 pub min_peers: usize,
741 pub channel_pts_entries: usize,
743 pub updates_initialised: bool,
745 pub approx_bytes: usize,
747}
748
749impl std::fmt::Display for PersistedSession {
750 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
751 use base64::Engine as _;
752 f.write_str(&base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(self.to_bytes()))
753 }
754}
755
756pub fn default_dc_addresses() -> HashMap<i32, String> {
758 [
759 (1, "149.154.175.53:443"),
760 (2, "149.154.167.51:443"),
761 (3, "149.154.175.100:443"),
762 (4, "149.154.167.91:443"),
763 (5, "91.108.56.130:443"),
764 ]
765 .into_iter()
766 .map(|(id, addr)| (id, addr.to_string()))
767 .collect()
768}
769
770use std::path::PathBuf;
775
776pub trait SessionBackend: Send + Sync {
784 fn save(&self, session: &PersistedSession) -> io::Result<()>;
785 fn load(&self) -> io::Result<Option<PersistedSession>>;
786 fn delete(&self) -> io::Result<()>;
787
788 fn name(&self) -> &str;
790
791 fn update_dc(&self, entry: &DcEntry) -> io::Result<()> {
804 let mut s = self.load()?.unwrap_or_default();
805 if let Some(existing) = s
807 .dcs
808 .iter_mut()
809 .find(|d| d.dc_id == entry.dc_id && d.is_ipv6() == entry.is_ipv6())
810 {
811 *existing = entry.clone();
812 } else {
813 s.dcs.push(entry.clone());
814 }
815 self.save(&s)
816 }
817
818 fn set_home_dc(&self, dc_id: i32) -> io::Result<()> {
824 let mut s = self.load()?.unwrap_or_default();
825 s.home_dc_id = dc_id;
826 self.save(&s)
827 }
828
829 fn apply_update_state(&self, update: UpdateStateChange) -> io::Result<()> {
834 let mut s = self.load()?.unwrap_or_default();
835 update.apply_to(&mut s.updates_state);
836 self.save(&s)
837 }
838
839 fn cache_peer(&self, peer: &CachedPeer) -> io::Result<()> {
845 let mut s = self.load()?.unwrap_or_default();
846 if let Some(existing) = s.peers.iter_mut().find(|p| p.id == peer.id) {
847 *existing = peer.clone();
848 } else {
849 s.peers.push(peer.clone());
850 }
851 self.save(&s)
852 }
853}
854
855#[derive(Debug, Clone)]
870pub enum UpdateStateChange {
871 All(UpdatesStateSnap),
873 Primary { pts: i32, date: i32, seq: i32 },
875 Secondary { qts: i32 },
877 Channel { id: i64, pts: i32 },
879}
880
881impl UpdateStateChange {
882 pub fn apply_to(&self, snap: &mut UpdatesStateSnap) {
884 match self {
885 Self::All(new_snap) => *snap = new_snap.clone(),
886 Self::Primary { pts, date, seq } => {
887 snap.pts = *pts;
888 snap.date = *date;
889 snap.seq = *seq;
890 }
891 Self::Secondary { qts } => {
892 snap.qts = *qts;
893 }
894 Self::Channel { id, pts } => {
895 if let Some(existing) = snap.channels.iter_mut().find(|c| c.0 == *id) {
897 existing.1 = *pts;
898 } else {
899 snap.channels.push((*id, *pts));
900 }
901 }
902 }
903 }
904}
905
906pub struct BinaryFileBackend {
910 path: PathBuf,
911 write_lock: std::sync::Mutex<()>,
916}
917
918impl BinaryFileBackend {
919 pub fn new(path: impl Into<PathBuf>) -> Self {
922 Self {
923 path: path.into(),
924 write_lock: std::sync::Mutex::new(()),
925 }
926 }
927
928 pub fn path(&self) -> &std::path::Path {
930 &self.path
931 }
932}
933
934impl SessionBackend for BinaryFileBackend {
935 fn save(&self, session: &PersistedSession) -> io::Result<()> {
936 let _guard = self.write_lock.lock().unwrap();
937 session.save(&self.path)
938 }
939
940 fn load(&self) -> io::Result<Option<PersistedSession>> {
941 if !self.path.exists() {
942 return Ok(None);
943 }
944 match PersistedSession::load(&self.path) {
945 Ok(s) => Ok(Some(s)),
946 Err(e) => {
947 let bak = self.path.with_extension("bak");
948 tracing::warn!(
949 "[ferogram::session] session file {:?} could not be read ({e}); \
950 backing it up to {:?} and starting a new session",
951 self.path,
952 bak
953 );
954 let _ = std::fs::rename(&self.path, &bak);
955 Ok(None)
956 }
957 }
958 }
959
960 fn delete(&self) -> io::Result<()> {
961 if self.path.exists() {
962 std::fs::remove_file(&self.path)?;
963 }
964 Ok(())
965 }
966
967 fn name(&self) -> &str {
968 "binary-file"
969 }
970
971 }
974
975#[derive(Default)]
983pub struct InMemoryBackend {
984 data: std::sync::Mutex<Option<PersistedSession>>,
985}
986
987impl InMemoryBackend {
988 pub fn new() -> Self {
990 Self::default()
991 }
992
993 pub fn snapshot(&self) -> Option<PersistedSession> {
995 self.data.lock().unwrap().clone()
996 }
997}
998
999impl SessionBackend for InMemoryBackend {
1000 fn save(&self, s: &PersistedSession) -> io::Result<()> {
1001 *self.data.lock().unwrap() = Some(s.clone());
1002 Ok(())
1003 }
1004
1005 fn load(&self) -> io::Result<Option<PersistedSession>> {
1006 Ok(self.data.lock().unwrap().clone())
1007 }
1008
1009 fn delete(&self) -> io::Result<()> {
1010 *self.data.lock().unwrap() = None;
1011 Ok(())
1012 }
1013
1014 fn name(&self) -> &str {
1015 "in-memory"
1016 }
1017
1018 fn update_dc(&self, entry: &DcEntry) -> io::Result<()> {
1021 let mut guard = self.data.lock().unwrap();
1022 let s = guard.get_or_insert_with(PersistedSession::default);
1023 if let Some(existing) = s
1024 .dcs
1025 .iter_mut()
1026 .find(|d| d.dc_id == entry.dc_id && d.is_ipv6() == entry.is_ipv6())
1027 {
1028 *existing = entry.clone();
1029 } else {
1030 s.dcs.push(entry.clone());
1031 }
1032 Ok(())
1033 }
1034
1035 fn set_home_dc(&self, dc_id: i32) -> io::Result<()> {
1036 let mut guard = self.data.lock().unwrap();
1037 let s = guard.get_or_insert_with(PersistedSession::default);
1038 s.home_dc_id = dc_id;
1039 Ok(())
1040 }
1041
1042 fn apply_update_state(&self, update: UpdateStateChange) -> io::Result<()> {
1043 let mut guard = self.data.lock().unwrap();
1044 let s = guard.get_or_insert_with(PersistedSession::default);
1045 update.apply_to(&mut s.updates_state);
1046 Ok(())
1047 }
1048
1049 fn cache_peer(&self, peer: &CachedPeer) -> io::Result<()> {
1050 let mut guard = self.data.lock().unwrap();
1051 let s = guard.get_or_insert_with(PersistedSession::default);
1052 if let Some(existing) = s.peers.iter_mut().find(|p| p.id == peer.id) {
1053 *existing = peer.clone();
1054 } else {
1055 s.peers.push(peer.clone());
1056 }
1057 Ok(())
1058 }
1059}
1060
1061pub struct StringSessionBackend {
1065 data: std::sync::Mutex<String>,
1066}
1067
1068impl StringSessionBackend {
1069 pub fn new(s: impl Into<String>) -> Self {
1072 Self {
1073 data: std::sync::Mutex::new(s.into()),
1074 }
1075 }
1076
1077 pub fn current(&self) -> String {
1079 self.data.lock().unwrap().clone()
1080 }
1081}
1082
1083impl SessionBackend for StringSessionBackend {
1084 fn save(&self, session: &PersistedSession) -> io::Result<()> {
1085 *self.data.lock().unwrap() = session.to_string();
1086 Ok(())
1087 }
1088
1089 fn load(&self) -> io::Result<Option<PersistedSession>> {
1090 let s = self.data.lock().unwrap().clone();
1091 if s.trim().is_empty() {
1092 return Ok(None);
1093 }
1094 PersistedSession::from_string(&s).map(Some)
1095 }
1096
1097 fn delete(&self) -> io::Result<()> {
1098 *self.data.lock().unwrap() = String::new();
1099 Ok(())
1100 }
1101
1102 fn name(&self) -> &str {
1103 "string-session"
1104 }
1105}
1106
1107pub mod string_session;
1115pub use string_session::{FullSession, Session, StringSession, StringSessionError};
1116
1117#[cfg(test)]
1118mod tests {
1119 use super::*;
1120
1121 fn make_dc(id: i32) -> DcEntry {
1122 DcEntry {
1123 dc_id: id,
1124 addr: format!("1.2.3.{id}:443"),
1125 auth_key: None,
1126 first_salt: 0,
1127 time_offset: 0,
1128 flags: DcFlags::NONE,
1129 }
1130 }
1131
1132 fn make_peer(id: i64, hash: i64) -> CachedPeer {
1133 CachedPeer {
1134 id,
1135 access_hash: hash,
1136 is_channel: false,
1137 is_chat: false,
1138 channel_kind: None,
1139 }
1140 }
1141
1142 #[test]
1145 fn inmemory_load_returns_none_when_empty() {
1146 let b = InMemoryBackend::new();
1147 assert!(b.load().unwrap().is_none());
1148 }
1149
1150 #[test]
1151 fn inmemory_save_then_load_round_trips() {
1152 let b = InMemoryBackend::new();
1153 let mut s = PersistedSession::default();
1154 s.home_dc_id = 3;
1155 s.dcs.push(make_dc(3));
1156 b.save(&s).unwrap();
1157
1158 let loaded = b.load().unwrap().unwrap();
1159 assert_eq!(loaded.home_dc_id, 3);
1160 assert_eq!(loaded.dcs.len(), 1);
1161 }
1162
1163 #[test]
1164 fn inmemory_delete_clears_state() {
1165 let b = InMemoryBackend::new();
1166 let mut s = PersistedSession::default();
1167 s.home_dc_id = 2;
1168 b.save(&s).unwrap();
1169 b.delete().unwrap();
1170 assert!(b.load().unwrap().is_none());
1171 }
1172
1173 #[test]
1176 fn inmemory_update_dc_inserts_new() {
1177 let b = InMemoryBackend::new();
1178 b.update_dc(&make_dc(4)).unwrap();
1179 let s = b.snapshot().unwrap();
1180 assert_eq!(s.dcs.len(), 1);
1181 assert_eq!(s.dcs[0].dc_id, 4);
1182 }
1183
1184 #[test]
1185 fn inmemory_update_dc_replaces_existing() {
1186 let b = InMemoryBackend::new();
1187 b.update_dc(&make_dc(2)).unwrap();
1188 let mut updated = make_dc(2);
1189 updated.addr = "9.9.9.9:443".to_string();
1190 b.update_dc(&updated).unwrap();
1191
1192 let s = b.snapshot().unwrap();
1193 assert_eq!(s.dcs.len(), 1);
1194 assert_eq!(s.dcs[0].addr, "9.9.9.9:443");
1195 }
1196
1197 #[test]
1198 fn inmemory_set_home_dc() {
1199 let b = InMemoryBackend::new();
1200 b.set_home_dc(5).unwrap();
1201 assert_eq!(b.snapshot().unwrap().home_dc_id, 5);
1202 }
1203
1204 #[test]
1205 fn inmemory_cache_peer_inserts() {
1206 let b = InMemoryBackend::new();
1207 b.cache_peer(&make_peer(100, 0xdeadbeef)).unwrap();
1208 let s = b.snapshot().unwrap();
1209 assert_eq!(s.peers.len(), 1);
1210 assert_eq!(s.peers[0].id, 100);
1211 }
1212
1213 #[test]
1214 fn inmemory_cache_peer_updates_existing() {
1215 let b = InMemoryBackend::new();
1216 b.cache_peer(&make_peer(100, 111)).unwrap();
1217 b.cache_peer(&make_peer(100, 222)).unwrap();
1218 let s = b.snapshot().unwrap();
1219 assert_eq!(s.peers.len(), 1);
1220 assert_eq!(s.peers[0].access_hash, 222);
1221 }
1222
1223 #[test]
1226 fn update_state_primary() {
1227 let mut snap = UpdatesStateSnap {
1228 pts: 0,
1229 qts: 0,
1230 date: 0,
1231 seq: 0,
1232 channels: vec![],
1233 };
1234 UpdateStateChange::Primary {
1235 pts: 10,
1236 date: 20,
1237 seq: 30,
1238 }
1239 .apply_to(&mut snap);
1240 assert_eq!(snap.pts, 10);
1241 assert_eq!(snap.date, 20);
1242 assert_eq!(snap.seq, 30);
1243 assert_eq!(snap.qts, 0); }
1245
1246 #[test]
1247 fn update_state_secondary() {
1248 let mut snap = UpdatesStateSnap {
1249 pts: 5,
1250 qts: 0,
1251 date: 0,
1252 seq: 0,
1253 channels: vec![],
1254 };
1255 UpdateStateChange::Secondary { qts: 99 }.apply_to(&mut snap);
1256 assert_eq!(snap.qts, 99);
1257 assert_eq!(snap.pts, 5); }
1259
1260 #[test]
1261 fn update_state_channel_inserts() {
1262 let mut snap = UpdatesStateSnap {
1263 pts: 0,
1264 qts: 0,
1265 date: 0,
1266 seq: 0,
1267 channels: vec![],
1268 };
1269 UpdateStateChange::Channel { id: 12345, pts: 42 }.apply_to(&mut snap);
1270 assert_eq!(snap.channels, vec![(12345, 42)]);
1271 }
1272
1273 #[test]
1274 fn update_state_channel_updates_existing() {
1275 let mut snap = UpdatesStateSnap {
1276 pts: 0,
1277 qts: 0,
1278 date: 0,
1279 seq: 0,
1280 channels: vec![(12345, 10), (67890, 5)],
1281 };
1282 UpdateStateChange::Channel { id: 12345, pts: 99 }.apply_to(&mut snap);
1283 assert_eq!(snap.channels[0], (12345, 99));
1285 assert_eq!(snap.channels[1], (67890, 5));
1286 }
1287
1288 #[test]
1289 fn apply_update_state_via_backend() {
1290 let b = InMemoryBackend::new();
1291 b.apply_update_state(UpdateStateChange::Primary {
1292 pts: 7,
1293 date: 8,
1294 seq: 9,
1295 })
1296 .unwrap();
1297 let s = b.snapshot().unwrap();
1298 assert_eq!(s.updates_state.pts, 7);
1299 }
1300
1301 #[test]
1304 fn default_update_dc_via_trait_object() {
1305 let b: Box<dyn SessionBackend> = Box::new(InMemoryBackend::new());
1306 b.update_dc(&make_dc(1)).unwrap();
1307 b.update_dc(&make_dc(2)).unwrap();
1308 let loaded = b.load().unwrap().unwrap();
1310 assert_eq!(loaded.dcs.len(), 2);
1311 }
1312
1313 fn make_dc_v6(id: i32) -> DcEntry {
1316 DcEntry {
1317 dc_id: id,
1318 addr: format!("[2001:b28:f23d:f00{}::a]:443", id),
1319 auth_key: None,
1320 first_salt: 0,
1321 time_offset: 0,
1322 flags: DcFlags::IPV6,
1323 }
1324 }
1325
1326 #[test]
1327 fn dc_entry_from_parts_ipv4() {
1328 let dc = DcEntry::from_parts(1, "149.154.175.53", 443, DcFlags::NONE);
1329 assert_eq!(dc.addr, "149.154.175.53:443");
1330 assert!(!dc.is_ipv6());
1331 let sa = dc.socket_addr().unwrap();
1332 assert_eq!(sa.port(), 443);
1333 }
1334
1335 #[test]
1336 fn dc_entry_from_parts_ipv6() {
1337 let dc = DcEntry::from_parts(2, "2001:b28:f23d:f001::a", 443, DcFlags::IPV6);
1338 assert_eq!(dc.addr, "[2001:b28:f23d:f001::a]:443");
1339 assert!(dc.is_ipv6());
1340 let sa = dc.socket_addr().unwrap();
1341 assert_eq!(sa.port(), 443);
1342 }
1343
1344 #[test]
1345 fn persisted_session_dc_for_prefers_ipv6() {
1346 let mut s = PersistedSession::default();
1347 s.dcs.push(make_dc(2)); s.dcs.push(make_dc_v6(2)); let v6 = s.dc_for(2, true).unwrap();
1351 assert!(v6.is_ipv6());
1352
1353 let v4 = s.dc_for(2, false).unwrap();
1354 assert!(!v4.is_ipv6());
1355 }
1356
1357 #[test]
1358 fn persisted_session_dc_for_falls_back_when_only_ipv4() {
1359 let mut s = PersistedSession::default();
1360 s.dcs.push(make_dc(3)); let dc = s.dc_for(3, true).unwrap();
1364 assert!(!dc.is_ipv6());
1365 }
1366
1367 #[test]
1368 fn persisted_session_all_dcs_for_returns_both() {
1369 let mut s = PersistedSession::default();
1370 s.dcs.push(make_dc(1));
1371 s.dcs.push(make_dc_v6(1));
1372 s.dcs.push(make_dc(2));
1373
1374 assert_eq!(s.all_dcs_for(1).count(), 2);
1375 assert_eq!(s.all_dcs_for(2).count(), 1);
1376 assert_eq!(s.all_dcs_for(5).count(), 0);
1377 }
1378
1379 #[test]
1380 fn inmemory_ipv4_and_ipv6_coexist() {
1381 let b = InMemoryBackend::new();
1382 b.update_dc(&make_dc(2)).unwrap(); b.update_dc(&make_dc_v6(2)).unwrap(); let s = b.snapshot().unwrap();
1386 assert_eq!(s.dcs.iter().filter(|d| d.dc_id == 2).count(), 2);
1388 }
1389
1390 #[test]
1391 fn binary_roundtrip_ipv4_and_ipv6() {
1392 let mut s = PersistedSession::default();
1393 s.home_dc_id = 2;
1394 s.dcs.push(make_dc(2));
1395 s.dcs.push(make_dc_v6(2));
1396
1397 let bytes = s.to_bytes();
1398 let loaded = PersistedSession::from_bytes(&bytes).unwrap();
1399 assert_eq!(loaded.dcs.len(), 2);
1400 assert_eq!(loaded.dcs.iter().filter(|d| d.is_ipv6()).count(), 1);
1401 assert_eq!(loaded.dcs.iter().filter(|d| !d.is_ipv6()).count(), 1);
1402 }
1403
1404 #[test]
1405 fn v6_channel_kind_roundtrip_all_variants() {
1406 let mut s = PersistedSession::default();
1407 s.home_dc_id = 1;
1408 s.peers.push(CachedPeer {
1409 id: 1001,
1410 access_hash: 0xaaaa,
1411 is_channel: true,
1412 is_chat: false,
1413 channel_kind: Some(ChannelKind::Broadcast),
1414 });
1415 s.peers.push(CachedPeer {
1416 id: 1002,
1417 access_hash: 0xbbbb,
1418 is_channel: true,
1419 is_chat: false,
1420 channel_kind: Some(ChannelKind::Megagroup),
1421 });
1422 s.peers.push(CachedPeer {
1423 id: 1003,
1424 access_hash: 0xcccc,
1425 is_channel: true,
1426 is_chat: false,
1427 channel_kind: Some(ChannelKind::Gigagroup),
1428 });
1429 s.peers.push(CachedPeer {
1430 id: 1004,
1431 access_hash: 0xdddd,
1432 is_channel: false,
1433 is_chat: false,
1434 channel_kind: None,
1435 });
1436
1437 let bytes = s.to_bytes();
1438 assert_eq!(bytes[0], 0x07, "version byte must be 7");
1439
1440 let loaded = PersistedSession::from_bytes(&bytes).unwrap();
1441 assert_eq!(loaded.peers.len(), 4);
1442
1443 let p = &loaded.peers[0];
1444 assert_eq!(p.id, 1001);
1445 assert_eq!(p.channel_kind, Some(ChannelKind::Broadcast));
1446
1447 let p = &loaded.peers[1];
1448 assert_eq!(p.id, 1002);
1449 assert_eq!(p.channel_kind, Some(ChannelKind::Megagroup));
1450
1451 let p = &loaded.peers[2];
1452 assert_eq!(p.id, 1003);
1453 assert_eq!(p.channel_kind, Some(ChannelKind::Gigagroup));
1454
1455 let p = &loaded.peers[3];
1456 assert_eq!(p.id, 1004);
1457 assert_eq!(p.channel_kind, None);
1458 }
1459
1460 #[test]
1461 fn v7_future_auth_token_roundtrip() {
1462 let mut s = PersistedSession::default();
1463 s.home_dc_id = 2;
1464 s.future_auth_token = Some(vec![1, 2, 3, 4, 5]);
1465
1466 let bytes = s.to_bytes();
1467 assert_eq!(bytes[0], 0x07, "version byte must be 7");
1468
1469 let loaded = PersistedSession::from_bytes(&bytes).unwrap();
1470 assert_eq!(loaded.future_auth_token, Some(vec![1, 2, 3, 4, 5]));
1471 }
1472
1473 #[test]
1474 fn v7_future_auth_token_absent() {
1475 let s = PersistedSession {
1476 home_dc_id: 2,
1477 ..Default::default()
1478 };
1479
1480 let bytes = s.to_bytes();
1481 let loaded = PersistedSession::from_bytes(&bytes).unwrap();
1482 assert_eq!(loaded.future_auth_token, None);
1483 }
1484
1485 #[test]
1486 fn v6_file_without_token_loads_as_none() {
1487 let mut s = PersistedSession::default();
1490 s.home_dc_id = 3;
1491 let mut bytes = s.to_bytes();
1492 bytes[0] = 0x06; bytes.truncate(bytes.len() - 1); let loaded = PersistedSession::from_bytes(&bytes).unwrap();
1496 assert_eq!(loaded.home_dc_id, 3);
1497 assert_eq!(loaded.future_auth_token, None);
1498 }
1499
1500 #[test]
1501 fn v6_channel_kind_absent_sentinel_roundtrip() {
1502 let mut s = PersistedSession::default();
1504 s.home_dc_id = 1;
1505 s.peers.push(CachedPeer {
1506 id: 555,
1507 access_hash: 0x1234,
1508 is_channel: false,
1509 is_chat: false,
1510 channel_kind: None,
1511 });
1512 let loaded = PersistedSession::from_bytes(&s.to_bytes()).unwrap();
1513 assert_eq!(loaded.peers[0].channel_kind, None);
1514 }
1515}
1516
1517#[cfg(feature = "sqlite-session")]
1543pub struct SqliteBackend {
1544 conn: std::sync::Mutex<rusqlite::Connection>,
1545 label: String,
1546}
1547
1548#[cfg(feature = "sqlite-session")]
1549impl SqliteBackend {
1550 const SCHEMA: &'static str = "
1551 PRAGMA journal_mode = WAL;
1552 PRAGMA synchronous = NORMAL;
1553
1554 CREATE TABLE IF NOT EXISTS meta (
1555 key TEXT PRIMARY KEY,
1556 value INTEGER NOT NULL DEFAULT 0
1557 );
1558
1559 CREATE TABLE IF NOT EXISTS dcs (
1560 dc_id INTEGER NOT NULL,
1561 flags INTEGER NOT NULL DEFAULT 0,
1562 addr TEXT NOT NULL,
1563 auth_key BLOB,
1564 first_salt INTEGER NOT NULL DEFAULT 0,
1565 time_offset INTEGER NOT NULL DEFAULT 0,
1566 PRIMARY KEY (dc_id, flags)
1567 );
1568
1569 CREATE TABLE IF NOT EXISTS update_state (
1570 id INTEGER PRIMARY KEY CHECK (id = 1),
1571 pts INTEGER NOT NULL DEFAULT 0,
1572 qts INTEGER NOT NULL DEFAULT 0,
1573 date INTEGER NOT NULL DEFAULT 0,
1574 seq INTEGER NOT NULL DEFAULT 0
1575 );
1576
1577 CREATE TABLE IF NOT EXISTS channel_pts (
1578 channel_id INTEGER PRIMARY KEY,
1579 pts INTEGER NOT NULL
1580 );
1581
1582 CREATE TABLE IF NOT EXISTS peers (
1583 id INTEGER PRIMARY KEY,
1584 access_hash INTEGER NOT NULL,
1585 is_channel INTEGER NOT NULL DEFAULT 0,
1586 is_chat INTEGER NOT NULL DEFAULT 0,
1587 channel_kind INTEGER
1588 );
1589
1590 CREATE TABLE IF NOT EXISTS min_peers (
1591 user_id INTEGER PRIMARY KEY,
1592 peer_id INTEGER NOT NULL,
1593 msg_id INTEGER NOT NULL
1594 );
1595 ";
1596
1597 pub fn open(path: impl Into<PathBuf>) -> io::Result<Self> {
1599 let path = path.into();
1600 let label = path.display().to_string();
1601 let conn = rusqlite::Connection::open(&path).map_err(io::Error::other)?;
1602 conn.execute_batch(Self::SCHEMA).map_err(io::Error::other)?;
1603 Self::migrate_legacy_sqlite_schema(&conn)?;
1604 Ok(Self {
1605 conn: std::sync::Mutex::new(conn),
1606 label,
1607 })
1608 }
1609
1610 pub fn in_memory() -> io::Result<Self> {
1612 let conn = rusqlite::Connection::open_in_memory().map_err(io::Error::other)?;
1613 conn.execute_batch(Self::SCHEMA).map_err(io::Error::other)?;
1614 Self::migrate_legacy_sqlite_schema(&conn)?;
1615 Ok(Self {
1616 conn: std::sync::Mutex::new(conn),
1617 label: ":memory:".into(),
1618 })
1619 }
1620
1621 fn map_err(e: rusqlite::Error) -> io::Error {
1622 io::Error::other(e)
1623 }
1624
1625 fn migrate_legacy_sqlite_schema(conn: &rusqlite::Connection) -> io::Result<()> {
1629 let mut has_is_chat = false;
1630 let mut stmt = conn
1631 .prepare("PRAGMA table_info(peers)")
1632 .map_err(Self::map_err)?;
1633 let cols = stmt
1634 .query_map([], |row| row.get::<_, String>(1))
1635 .map_err(Self::map_err)?;
1636 for col in cols.filter_map(|r| r.ok()) {
1637 if col == "is_chat" {
1638 has_is_chat = true;
1639 break;
1640 }
1641 }
1642 if !has_is_chat {
1643 conn.execute_batch("ALTER TABLE peers ADD COLUMN is_chat INTEGER NOT NULL DEFAULT 0;")
1644 .map_err(Self::map_err)?;
1645 }
1646 let mut has_channel_kind = false;
1648 let mut stmt2 = conn
1649 .prepare("PRAGMA table_info(peers)")
1650 .map_err(Self::map_err)?;
1651 let cols2 = stmt2
1652 .query_map([], |row| row.get::<_, String>(1))
1653 .map_err(Self::map_err)?;
1654 for col in cols2.filter_map(|r| r.ok()) {
1655 if col == "channel_kind" {
1656 has_channel_kind = true;
1657 break;
1658 }
1659 }
1660 if !has_channel_kind {
1661 conn.execute_batch("ALTER TABLE peers ADD COLUMN channel_kind INTEGER;")
1662 .map_err(Self::map_err)?;
1663 }
1664 conn.execute_batch(
1665 "CREATE TABLE IF NOT EXISTS min_peers (
1666 user_id INTEGER PRIMARY KEY,
1667 peer_id INTEGER NOT NULL,
1668 msg_id INTEGER NOT NULL
1669 );",
1670 )
1671 .map_err(Self::map_err)?;
1672 Ok(())
1673 }
1674
1675 fn read_session(conn: &rusqlite::Connection) -> io::Result<PersistedSession> {
1677 let home_dc_id: i32 = conn
1679 .query_row("SELECT value FROM meta WHERE key = 'home_dc_id'", [], |r| {
1680 r.get(0)
1681 })
1682 .unwrap_or(0);
1683
1684 let future_auth_token: Option<Vec<u8>> = conn
1686 .query_row(
1687 "SELECT value FROM meta WHERE key = 'future_auth_token'",
1688 [],
1689 |r| r.get::<_, Vec<u8>>(0),
1690 )
1691 .ok();
1692
1693 let mut stmt = conn
1695 .prepare("SELECT dc_id, flags, addr, auth_key, first_salt, time_offset FROM dcs")
1696 .map_err(Self::map_err)?;
1697 let dcs = stmt
1698 .query_map([], |row| {
1699 let dc_id: i32 = row.get(0)?;
1700 let flags_raw: u8 = row.get(1)?;
1701 let addr: String = row.get(2)?;
1702 let key_blob: Option<Vec<u8>> = row.get(3)?;
1703 let first_salt: i64 = row.get(4)?;
1704 let time_offset: i32 = row.get(5)?;
1705 Ok((dc_id, addr, key_blob, first_salt, time_offset, flags_raw))
1706 })
1707 .map_err(Self::map_err)?
1708 .filter_map(|r| r.ok())
1709 .map(
1710 |(dc_id, addr, key_blob, first_salt, time_offset, flags_raw)| {
1711 let auth_key = key_blob.and_then(|b| {
1712 if b.len() == 256 {
1713 let mut k = [0u8; 256];
1714 k.copy_from_slice(&b);
1715 Some(k)
1716 } else {
1717 None
1718 }
1719 });
1720 DcEntry {
1721 dc_id,
1722 addr,
1723 auth_key,
1724 first_salt,
1725 time_offset,
1726 flags: DcFlags(flags_raw),
1727 }
1728 },
1729 )
1730 .collect();
1731
1732 let updates_state = conn
1734 .query_row(
1735 "SELECT pts, qts, date, seq FROM update_state WHERE id = 1",
1736 [],
1737 |r| {
1738 Ok(UpdatesStateSnap {
1739 pts: r.get(0)?,
1740 qts: r.get(1)?,
1741 date: r.get(2)?,
1742 seq: r.get(3)?,
1743 channels: vec![],
1744 })
1745 },
1746 )
1747 .unwrap_or_default();
1748
1749 let mut ch_stmt = conn
1751 .prepare("SELECT channel_id, pts FROM channel_pts")
1752 .map_err(Self::map_err)?;
1753 let channels: Vec<(i64, i32)> = ch_stmt
1754 .query_map([], |r| Ok((r.get::<_, i64>(0)?, r.get::<_, i32>(1)?)))
1755 .map_err(Self::map_err)?
1756 .filter_map(|r| r.ok())
1757 .collect();
1758
1759 let mut peer_stmt = conn
1761 .prepare("SELECT id, access_hash, is_channel, is_chat, channel_kind FROM peers")
1762 .map_err(Self::map_err)?;
1763 let peers: Vec<CachedPeer> = peer_stmt
1764 .query_map([], |r| {
1765 let kind_raw: Option<i32> = r.get(4)?;
1766 Ok(CachedPeer {
1767 id: r.get(0)?,
1768 access_hash: r.get(1)?,
1769 is_channel: r.get::<_, i32>(2)? != 0,
1770 is_chat: r.get::<_, i32>(3)? != 0,
1771 channel_kind: kind_raw.and_then(|k| ChannelKind::from_byte(k as u8)),
1772 })
1773 })
1774 .map_err(Self::map_err)?
1775 .filter_map(|r| r.ok())
1776 .collect();
1777
1778 let mut min_stmt = conn
1780 .prepare("SELECT user_id, peer_id, msg_id FROM min_peers")
1781 .map_err(Self::map_err)?;
1782 let min_peers: Vec<CachedMinPeer> = min_stmt
1783 .query_map([], |r| {
1784 Ok(CachedMinPeer {
1785 user_id: r.get(0)?,
1786 peer_id: r.get(1)?,
1787 msg_id: r.get(2)?,
1788 })
1789 })
1790 .map_err(Self::map_err)?
1791 .filter_map(|r| r.ok())
1792 .collect();
1793
1794 Ok(PersistedSession {
1795 home_dc_id,
1796 dcs,
1797 updates_state: UpdatesStateSnap {
1798 channels,
1799 ..updates_state
1800 },
1801 peers,
1802 min_peers,
1803 future_auth_token,
1804 })
1805 }
1806
1807 fn write_session(conn: &rusqlite::Connection, s: &PersistedSession) -> io::Result<()> {
1809 conn.execute_batch("BEGIN IMMEDIATE")
1810 .map_err(Self::map_err)?;
1811
1812 conn.execute(
1813 "INSERT INTO meta (key, value) VALUES ('home_dc_id', ?1)
1814 ON CONFLICT(key) DO UPDATE SET value = excluded.value",
1815 rusqlite::params![s.home_dc_id],
1816 )
1817 .map_err(Self::map_err)?;
1818
1819 match &s.future_auth_token {
1820 Some(token) => conn.execute(
1821 "INSERT INTO meta (key, value) VALUES ('future_auth_token', ?1)
1822 ON CONFLICT(key) DO UPDATE SET value = excluded.value",
1823 rusqlite::params![token],
1824 ),
1825 None => conn.execute("DELETE FROM meta WHERE key = 'future_auth_token'", []),
1826 }
1827 .map_err(Self::map_err)?;
1828
1829 conn.execute("DELETE FROM dcs", []).map_err(Self::map_err)?;
1831 for d in &s.dcs {
1832 conn.execute(
1833 "INSERT INTO dcs (dc_id, flags, addr, auth_key, first_salt, time_offset)
1834 VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
1835 rusqlite::params![
1836 d.dc_id,
1837 d.flags.0,
1838 d.addr,
1839 d.auth_key.as_ref().map(|k| k.as_ref()),
1840 d.first_salt,
1841 d.time_offset,
1842 ],
1843 )
1844 .map_err(Self::map_err)?;
1845 }
1846
1847 let us = &s.updates_state;
1851 conn.execute(
1852 "INSERT INTO update_state (id, pts, qts, date, seq) VALUES (1, ?1, ?2, ?3, ?4)
1853 ON CONFLICT(id) DO UPDATE SET
1854 pts = MAX(excluded.pts, update_state.pts),
1855 qts = MAX(excluded.qts, update_state.qts),
1856 date = excluded.date,
1857 seq = excluded.seq",
1858 rusqlite::params![us.pts, us.qts, us.date, us.seq],
1859 )
1860 .map_err(Self::map_err)?;
1861
1862 conn.execute("DELETE FROM channel_pts", [])
1863 .map_err(Self::map_err)?;
1864 for &(cid, cpts) in &us.channels {
1865 conn.execute(
1866 "INSERT INTO channel_pts (channel_id, pts) VALUES (?1, ?2)",
1867 rusqlite::params![cid, cpts],
1868 )
1869 .map_err(Self::map_err)?;
1870 }
1871
1872 conn.execute("DELETE FROM peers", [])
1874 .map_err(Self::map_err)?;
1875 for p in &s.peers {
1876 conn.execute(
1877 "INSERT INTO peers (id, access_hash, is_channel, is_chat, channel_kind) VALUES (?1, ?2, ?3, ?4, ?5)",
1878 rusqlite::params![
1879 p.id,
1880 p.access_hash,
1881 p.is_channel as i32,
1882 p.is_chat as i32,
1883 p.channel_kind.map(|k| k.to_byte() as i32),
1884 ],
1885 )
1886 .map_err(Self::map_err)?;
1887 }
1888
1889 conn.execute("DELETE FROM min_peers", [])
1891 .map_err(Self::map_err)?;
1892 for m in &s.min_peers {
1893 conn.execute(
1894 "INSERT INTO min_peers (user_id, peer_id, msg_id) VALUES (?1, ?2, ?3)",
1895 rusqlite::params![m.user_id, m.peer_id, m.msg_id],
1896 )
1897 .map_err(Self::map_err)?;
1898 }
1899
1900 conn.execute_batch("COMMIT").map_err(Self::map_err)
1901 }
1902}
1903
1904#[cfg(feature = "sqlite-session")]
1905impl SessionBackend for SqliteBackend {
1906 fn save(&self, session: &PersistedSession) -> io::Result<()> {
1907 let conn = self.conn.lock().unwrap();
1908 Self::write_session(&conn, session)
1909 }
1910
1911 fn load(&self) -> io::Result<Option<PersistedSession>> {
1912 let conn = self.conn.lock().unwrap();
1913 let count: i64 = conn
1915 .query_row("SELECT COUNT(*) FROM meta", [], |r| r.get(0))
1916 .map_err(Self::map_err)?;
1917 if count == 0 {
1918 return Ok(None);
1919 }
1920 Self::read_session(&conn).map(Some)
1921 }
1922
1923 fn delete(&self) -> io::Result<()> {
1924 let conn = self.conn.lock().unwrap();
1925 conn.execute_batch(
1926 "BEGIN IMMEDIATE;
1927 DELETE FROM meta;
1928 DELETE FROM dcs;
1929 DELETE FROM update_state;
1930 DELETE FROM channel_pts;
1931 DELETE FROM peers;
1932 DELETE FROM min_peers;
1933 COMMIT;",
1934 )
1935 .map_err(Self::map_err)
1936 }
1937
1938 fn name(&self) -> &str {
1939 &self.label
1940 }
1941
1942 fn update_dc(&self, entry: &DcEntry) -> io::Result<()> {
1945 let conn = self.conn.lock().unwrap();
1946 conn.execute(
1947 "INSERT INTO dcs (dc_id, flags, addr, auth_key, first_salt, time_offset)
1948 VALUES (?1, ?6, ?2, ?3, ?4, ?5)
1949 ON CONFLICT(dc_id, flags) DO UPDATE SET
1950 addr = excluded.addr,
1951 auth_key = excluded.auth_key,
1952 first_salt = excluded.first_salt,
1953 time_offset = excluded.time_offset",
1954 rusqlite::params![
1955 entry.dc_id,
1956 entry.addr,
1957 entry.auth_key.as_ref().map(|k| k.as_ref()),
1958 entry.first_salt,
1959 entry.time_offset,
1960 entry.flags.0,
1961 ],
1962 )
1963 .map(|_| ())
1964 .map_err(Self::map_err)
1965 }
1966
1967 fn set_home_dc(&self, dc_id: i32) -> io::Result<()> {
1968 let conn = self.conn.lock().unwrap();
1969 conn.execute(
1970 "INSERT INTO meta (key, value) VALUES ('home_dc_id', ?1)
1971 ON CONFLICT(key) DO UPDATE SET value = excluded.value",
1972 rusqlite::params![dc_id],
1973 )
1974 .map(|_| ())
1975 .map_err(Self::map_err)
1976 }
1977
1978 fn apply_update_state(&self, update: UpdateStateChange) -> io::Result<()> {
1979 let conn = self.conn.lock().unwrap();
1980 match update {
1981 UpdateStateChange::All(snap) => {
1982 conn.execute(
1983 "INSERT INTO update_state (id, pts, qts, date, seq) VALUES (1,?1,?2,?3,?4)
1984 ON CONFLICT(id) DO UPDATE SET
1985 pts=excluded.pts, qts=excluded.qts,
1986 date=excluded.date, seq=excluded.seq",
1987 rusqlite::params![snap.pts, snap.qts, snap.date, snap.seq],
1988 )
1989 .map_err(Self::map_err)?;
1990 conn.execute("DELETE FROM channel_pts", [])
1991 .map_err(Self::map_err)?;
1992 for &(cid, cpts) in &snap.channels {
1993 conn.execute(
1994 "INSERT INTO channel_pts (channel_id, pts) VALUES (?1, ?2)",
1995 rusqlite::params![cid, cpts],
1996 )
1997 .map_err(Self::map_err)?;
1998 }
1999 Ok(())
2000 }
2001 UpdateStateChange::Primary { pts, date, seq } => conn
2002 .execute(
2003 "INSERT INTO update_state (id, pts, qts, date, seq) VALUES (1,?1,0,?2,?3)
2004 ON CONFLICT(id) DO UPDATE SET pts=excluded.pts, date=excluded.date,
2005 seq=excluded.seq",
2006 rusqlite::params![pts, date, seq],
2007 )
2008 .map(|_| ())
2009 .map_err(Self::map_err),
2010 UpdateStateChange::Secondary { qts } => conn
2011 .execute(
2012 "INSERT INTO update_state (id, pts, qts, date, seq) VALUES (1,0,?1,0,0)
2013 ON CONFLICT(id) DO UPDATE SET qts = excluded.qts",
2014 rusqlite::params![qts],
2015 )
2016 .map(|_| ())
2017 .map_err(Self::map_err),
2018 UpdateStateChange::Channel { id, pts } => conn
2019 .execute(
2020 "INSERT INTO channel_pts (channel_id, pts) VALUES (?1, ?2)
2021 ON CONFLICT(channel_id) DO UPDATE SET pts = excluded.pts",
2022 rusqlite::params![id, pts],
2023 )
2024 .map(|_| ())
2025 .map_err(Self::map_err),
2026 }
2027 }
2028
2029 fn cache_peer(&self, peer: &CachedPeer) -> io::Result<()> {
2030 let conn = self.conn.lock().unwrap();
2031 conn.execute(
2032 "INSERT INTO peers (id, access_hash, is_channel, is_chat, channel_kind) VALUES (?1, ?2, ?3, ?4, ?5)
2033 ON CONFLICT(id) DO UPDATE SET
2034 access_hash = excluded.access_hash,
2035 is_channel = excluded.is_channel,
2036 is_chat = excluded.is_chat,
2037 channel_kind = excluded.channel_kind",
2038 rusqlite::params![
2039 peer.id,
2040 peer.access_hash,
2041 peer.is_channel as i32,
2042 peer.is_chat as i32,
2043 peer.channel_kind.map(|k| k.to_byte() as i32),
2044 ],
2045 )
2046 .map(|_| ())
2047 .map_err(Self::map_err)
2048 }
2049}
2050
2051#[cfg(all(test, feature = "sqlite-session"))]
2052mod sqlite_backend_tests {
2053 use super::*;
2054
2055 #[test]
2056 fn future_auth_token_roundtrip() {
2057 let b = SqliteBackend::in_memory().unwrap();
2058 let mut s = PersistedSession::default();
2059 s.home_dc_id = 2;
2060 s.future_auth_token = Some(vec![1, 2, 3, 4, 5]);
2061 b.save(&s).unwrap();
2062
2063 let loaded = b.load().unwrap().unwrap();
2064 assert_eq!(loaded.future_auth_token, Some(vec![1, 2, 3, 4, 5]));
2065 }
2066
2067 #[test]
2068 fn future_auth_token_cleared_on_none() {
2069 let b = SqliteBackend::in_memory().unwrap();
2070 let mut s = PersistedSession::default();
2071 s.home_dc_id = 2;
2072 s.future_auth_token = Some(vec![9, 9, 9]);
2073 b.save(&s).unwrap();
2074 assert_eq!(
2075 b.load().unwrap().unwrap().future_auth_token,
2076 Some(vec![9, 9, 9])
2077 );
2078
2079 s.future_auth_token = None;
2080 b.save(&s).unwrap();
2081 assert_eq!(b.load().unwrap().unwrap().future_auth_token, None);
2082 }
2083}
2084
2085#[cfg(feature = "libsql-session")]
2105pub struct LibSqlBackend {
2106 conn: libsql::Connection,
2107 label: String,
2108}
2109
2110#[cfg(feature = "libsql-session")]
2111impl LibSqlBackend {
2112 const SCHEMA: &'static str = "
2113 CREATE TABLE IF NOT EXISTS meta (
2114 key TEXT PRIMARY KEY,
2115 value INTEGER NOT NULL DEFAULT 0
2116 );
2117 CREATE TABLE IF NOT EXISTS dcs (
2118 dc_id INTEGER NOT NULL,
2119 flags INTEGER NOT NULL DEFAULT 0,
2120 addr TEXT NOT NULL,
2121 auth_key BLOB,
2122 first_salt INTEGER NOT NULL DEFAULT 0,
2123 time_offset INTEGER NOT NULL DEFAULT 0,
2124 PRIMARY KEY (dc_id, flags)
2125 );
2126 CREATE TABLE IF NOT EXISTS update_state (
2127 id INTEGER PRIMARY KEY CHECK (id = 1),
2128 pts INTEGER NOT NULL DEFAULT 0,
2129 qts INTEGER NOT NULL DEFAULT 0,
2130 date INTEGER NOT NULL DEFAULT 0,
2131 seq INTEGER NOT NULL DEFAULT 0
2132 );
2133 CREATE TABLE IF NOT EXISTS channel_pts (
2134 channel_id INTEGER PRIMARY KEY,
2135 pts INTEGER NOT NULL
2136 );
2137 CREATE TABLE IF NOT EXISTS peers (
2138 id INTEGER PRIMARY KEY,
2139 access_hash INTEGER NOT NULL,
2140 is_channel INTEGER NOT NULL DEFAULT 0,
2141 is_chat INTEGER NOT NULL DEFAULT 0,
2142 channel_kind INTEGER
2143 );
2144 CREATE TABLE IF NOT EXISTS min_peers (
2145 user_id INTEGER PRIMARY KEY,
2146 peer_id INTEGER NOT NULL,
2147 msg_id INTEGER NOT NULL
2148 );
2149 ";
2150
2151 fn block<F, T>(fut: F) -> io::Result<T>
2152 where
2153 F: std::future::Future<Output = Result<T, libsql::Error>>,
2154 {
2155 tokio::runtime::Handle::current()
2156 .block_on(fut)
2157 .map_err(io::Error::other)
2158 }
2159
2160 async fn apply_schema(conn: &libsql::Connection) -> Result<(), libsql::Error> {
2161 conn.execute_batch(Self::SCHEMA).await?;
2162 let _ = conn
2166 .execute_batch("ALTER TABLE peers ADD COLUMN channel_kind INTEGER;")
2167 .await;
2168 Ok(())
2169 }
2170
2171 pub fn open_local(path: impl Into<PathBuf>) -> io::Result<Self> {
2173 let path = path.into();
2174 let label = path.display().to_string();
2175 let db = Self::block(async { libsql::Builder::new_local(path).build().await })?;
2176 let conn = Self::block(async { db.connect() }).map_err(io::Error::other)?;
2177 Self::block(Self::apply_schema(&conn))?;
2178 Ok(Self { conn, label })
2179 }
2180
2181 pub fn in_memory() -> io::Result<Self> {
2183 let db = Self::block(async { libsql::Builder::new_local(":memory:").build().await })?;
2184 let conn = Self::block(async { db.connect() }).map_err(io::Error::other)?;
2185 Self::block(Self::apply_schema(&conn))?;
2186 Ok(Self {
2187 conn,
2188 label: ":memory:".into(),
2189 })
2190 }
2191
2192 pub fn open_remote(url: impl Into<String>, auth_token: impl Into<String>) -> io::Result<Self> {
2194 let url = url.into();
2195 let label = url.clone();
2196 let db = Self::block(async {
2197 libsql::Builder::new_remote(url, auth_token.into())
2198 .build()
2199 .await
2200 })?;
2201 let conn = Self::block(async { db.connect() }).map_err(io::Error::other)?;
2202 Self::block(Self::apply_schema(&conn))?;
2203 Ok(Self { conn, label })
2204 }
2205
2206 pub fn open_replica(
2208 path: impl Into<PathBuf>,
2209 url: impl Into<String>,
2210 auth_token: impl Into<String>,
2211 ) -> io::Result<Self> {
2212 let path = path.into();
2213 let url = url.into();
2214 let auth_token = auth_token.into();
2215 let label = format!("{} (replica of {})", path.display(), url);
2216 let db = Self::block(async {
2217 libsql::Builder::new_remote_replica(path, url, auth_token)
2218 .build()
2219 .await
2220 })?;
2221 let conn = Self::block(async { db.connect() }).map_err(io::Error::other)?;
2222 Self::block(Self::apply_schema(&conn))?;
2223 Ok(Self { conn, label })
2224 }
2225
2226 async fn read_session_async(
2227 conn: &libsql::Connection,
2228 ) -> Result<PersistedSession, libsql::Error> {
2229 let home_dc_id: i32 = conn
2231 .query("SELECT value FROM meta WHERE key = 'home_dc_id'", ())
2232 .await?
2233 .next()
2234 .await?
2235 .map(|r| r.get::<i32>(0))
2236 .transpose()?
2237 .unwrap_or(0);
2238
2239 let future_auth_token: Option<Vec<u8>> = conn
2241 .query("SELECT value FROM meta WHERE key = 'future_auth_token'", ())
2242 .await?
2243 .next()
2244 .await?
2245 .map(|r| r.get::<Vec<u8>>(0))
2246 .transpose()?;
2247
2248 let mut rows = conn
2250 .query(
2251 "SELECT dc_id, flags, addr, auth_key, first_salt, time_offset FROM dcs",
2252 (),
2253 )
2254 .await?;
2255 let mut dcs = Vec::new();
2256 while let Some(row) = rows.next().await? {
2257 let dc_id: i32 = row.get(0)?;
2258 let flags_raw: u8 = row.get::<i64>(1)? as u8;
2259 let addr: String = row.get(2)?;
2260 let key_blob: Option<Vec<u8>> = row.get(3)?;
2261 let first_salt: i64 = row.get(4)?;
2262 let time_offset: i32 = row.get(5)?;
2263 let auth_key = match key_blob {
2264 Some(b) if b.len() == 256 => {
2265 let mut k = [0u8; 256];
2266 k.copy_from_slice(&b);
2267 Some(k)
2268 }
2269 Some(b) => {
2270 return Err(libsql::Error::Misuse(format!(
2271 "auth_key blob must be 256 bytes, got {}",
2272 b.len()
2273 )));
2274 }
2275 None => None,
2276 };
2277 dcs.push(DcEntry {
2278 dc_id,
2279 addr,
2280 auth_key,
2281 first_salt,
2282 time_offset,
2283 flags: DcFlags(flags_raw),
2284 });
2285 }
2286
2287 let mut us_row = conn
2289 .query(
2290 "SELECT pts, qts, date, seq FROM update_state WHERE id = 1",
2291 (),
2292 )
2293 .await?;
2294 let updates_state = if let Some(r) = us_row.next().await? {
2295 UpdatesStateSnap {
2296 pts: r.get(0)?,
2297 qts: r.get(1)?,
2298 date: r.get(2)?,
2299 seq: r.get(3)?,
2300 channels: vec![],
2301 }
2302 } else {
2303 UpdatesStateSnap::default()
2304 };
2305
2306 let mut ch_rows = conn
2308 .query("SELECT channel_id, pts FROM channel_pts", ())
2309 .await?;
2310 let mut channels = Vec::new();
2311 while let Some(r) = ch_rows.next().await? {
2312 channels.push((r.get::<i64>(0)?, r.get::<i32>(1)?));
2313 }
2314
2315 let mut peer_rows = conn
2317 .query(
2318 "SELECT id, access_hash, is_channel, is_chat, channel_kind FROM peers",
2319 (),
2320 )
2321 .await?;
2322 let mut peers = Vec::new();
2323 while let Some(r) = peer_rows.next().await? {
2324 let kind_raw: Option<i32> = r.get(4).ok();
2325 peers.push(CachedPeer {
2326 id: r.get(0)?,
2327 access_hash: r.get(1)?,
2328 is_channel: r.get::<i32>(2)? != 0,
2329 is_chat: r.get::<i32>(3)? != 0,
2330 channel_kind: kind_raw.and_then(|k| ChannelKind::from_byte(k as u8)),
2331 });
2332 }
2333
2334 let mut min_rows = conn
2336 .query("SELECT user_id, peer_id, msg_id FROM min_peers", ())
2337 .await?;
2338 let mut min_peers = Vec::new();
2339 while let Some(r) = min_rows.next().await? {
2340 min_peers.push(CachedMinPeer {
2341 user_id: r.get(0)?,
2342 peer_id: r.get(1)?,
2343 msg_id: r.get(2)?,
2344 });
2345 }
2346
2347 Ok(PersistedSession {
2348 home_dc_id,
2349 dcs,
2350 updates_state: UpdatesStateSnap {
2351 channels,
2352 ..updates_state
2353 },
2354 peers,
2355 min_peers,
2356 future_auth_token,
2357 })
2358 }
2359
2360 async fn write_session_async(
2361 conn: &libsql::Connection,
2362 s: &PersistedSession,
2363 ) -> Result<(), libsql::Error> {
2364 conn.execute_batch("BEGIN IMMEDIATE").await.map(|_| ())?;
2365
2366 conn.execute(
2367 "INSERT INTO meta (key, value) VALUES ('home_dc_id', ?1)
2368 ON CONFLICT(key) DO UPDATE SET value = excluded.value",
2369 libsql::params![s.home_dc_id],
2370 )
2371 .await?;
2372
2373 match &s.future_auth_token {
2374 Some(token) => {
2375 conn.execute(
2376 "INSERT INTO meta (key, value) VALUES ('future_auth_token', ?1)
2377 ON CONFLICT(key) DO UPDATE SET value = excluded.value",
2378 libsql::params![token.clone()],
2379 )
2380 .await?
2381 }
2382 None => {
2383 conn.execute("DELETE FROM meta WHERE key = 'future_auth_token'", ())
2384 .await?
2385 }
2386 };
2387
2388 conn.execute("DELETE FROM dcs", ()).await?;
2389 for d in &s.dcs {
2390 conn.execute(
2391 "INSERT INTO dcs (dc_id, flags, addr, auth_key, first_salt, time_offset)
2392 VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
2393 libsql::params![
2394 d.dc_id,
2395 d.flags.0 as i64,
2396 d.addr.clone(),
2397 d.auth_key.map(|k| k.to_vec()),
2398 d.first_salt,
2399 d.time_offset,
2400 ],
2401 )
2402 .await?;
2403 }
2404
2405 let us = &s.updates_state;
2406 conn.execute(
2407 "INSERT INTO update_state (id, pts, qts, date, seq) VALUES (1,?1,?2,?3,?4)
2408 ON CONFLICT(id) DO UPDATE SET
2409 pts = MAX(excluded.pts, update_state.pts),
2410 qts = MAX(excluded.qts, update_state.qts),
2411 date = excluded.date,
2412 seq = excluded.seq",
2413 libsql::params![us.pts, us.qts, us.date, us.seq],
2414 )
2415 .await?;
2416
2417 conn.execute("DELETE FROM channel_pts", ()).await?;
2418 for &(cid, cpts) in &us.channels {
2419 conn.execute(
2420 "INSERT INTO channel_pts (channel_id, pts) VALUES (?1,?2)",
2421 libsql::params![cid, cpts],
2422 )
2423 .await?;
2424 }
2425
2426 conn.execute("DELETE FROM peers", ()).await?;
2427 for p in &s.peers {
2428 conn.execute(
2429 "INSERT INTO peers (id, access_hash, is_channel, is_chat, channel_kind) VALUES (?1,?2,?3,?4,?5)",
2430 libsql::params![
2431 p.id,
2432 p.access_hash,
2433 p.is_channel as i32,
2434 p.is_chat as i32,
2435 p.channel_kind.map(|k| k.to_byte() as i32),
2436 ],
2437 )
2438 .await?;
2439 }
2440
2441 conn.execute("DELETE FROM min_peers", ()).await?;
2442 for m in &s.min_peers {
2443 conn.execute(
2444 "INSERT INTO min_peers (user_id, peer_id, msg_id) VALUES (?1,?2,?3)",
2445 libsql::params![m.user_id, m.peer_id, m.msg_id],
2446 )
2447 .await?;
2448 }
2449
2450 conn.execute_batch("COMMIT").await.map(|_| ())
2451 }
2452}
2453
2454#[cfg(feature = "libsql-session")]
2455impl SessionBackend for LibSqlBackend {
2456 fn save(&self, session: &PersistedSession) -> io::Result<()> {
2457 let conn = self.conn.clone();
2458 let session = session.clone();
2459 Self::block(async move { Self::write_session_async(&conn, &session).await })
2460 }
2461
2462 fn load(&self) -> io::Result<Option<PersistedSession>> {
2463 let conn = self.conn.clone();
2464 let count: i64 = Self::block(async move {
2465 let mut rows = conn.query("SELECT COUNT(*) FROM meta", ()).await?;
2466 Ok::<i64, libsql::Error>(rows.next().await?.and_then(|r| r.get(0).ok()).unwrap_or(0))
2467 })?;
2468 if count == 0 {
2469 return Ok(None);
2470 }
2471 let conn = self.conn.clone();
2472 Self::block(async move { Self::read_session_async(&conn).await }).map(Some)
2473 }
2474
2475 fn delete(&self) -> io::Result<()> {
2476 let conn = self.conn.clone();
2477 Self::block(async move {
2478 conn.execute_batch(
2479 "BEGIN IMMEDIATE;
2480 DELETE FROM meta;
2481 DELETE FROM dcs;
2482 DELETE FROM update_state;
2483 DELETE FROM channel_pts;
2484 DELETE FROM peers;
2485 DELETE FROM min_peers;
2486 COMMIT;",
2487 )
2488 .await
2489 .map(|_| ())
2490 })
2491 }
2492
2493 fn name(&self) -> &str {
2494 &self.label
2495 }
2496
2497 fn update_dc(&self, entry: &DcEntry) -> io::Result<()> {
2500 let conn = self.conn.clone();
2501 let (dc_id, addr, key, salt, off, flags) = (
2502 entry.dc_id,
2503 entry.addr.clone(),
2504 entry.auth_key.map(|k| k.to_vec()),
2505 entry.first_salt,
2506 entry.time_offset,
2507 entry.flags.0 as i64,
2508 );
2509 Self::block(async move {
2510 conn.execute(
2511 "INSERT INTO dcs (dc_id, flags, addr, auth_key, first_salt, time_offset)
2512 VALUES (?1,?6,?2,?3,?4,?5)
2513 ON CONFLICT(dc_id, flags) DO UPDATE SET
2514 addr=excluded.addr, auth_key=excluded.auth_key,
2515 first_salt=excluded.first_salt, time_offset=excluded.time_offset",
2516 libsql::params![dc_id, addr, key, salt, off, flags],
2517 )
2518 .await
2519 .map(|_| ())
2520 })
2521 }
2522
2523 fn set_home_dc(&self, dc_id: i32) -> io::Result<()> {
2524 let conn = self.conn.clone();
2525 Self::block(async move {
2526 conn.execute(
2527 "INSERT INTO meta (key, value) VALUES ('home_dc_id',?1)
2528 ON CONFLICT(key) DO UPDATE SET value=excluded.value",
2529 libsql::params![dc_id],
2530 )
2531 .await
2532 .map(|_| ())
2533 })
2534 }
2535
2536 fn apply_update_state(&self, update: UpdateStateChange) -> io::Result<()> {
2537 let conn = self.conn.clone();
2538 Self::block(async move {
2539 match update {
2540 UpdateStateChange::All(snap) => {
2541 conn.execute(
2542 "INSERT INTO update_state (id,pts,qts,date,seq) VALUES (1,?1,?2,?3,?4)
2543 ON CONFLICT(id) DO UPDATE SET pts=excluded.pts,qts=excluded.qts,
2544 date=excluded.date,seq=excluded.seq",
2545 libsql::params![snap.pts, snap.qts, snap.date, snap.seq],
2546 )
2547 .await?;
2548 conn.execute("DELETE FROM channel_pts", ()).await?;
2549 for &(cid, cpts) in &snap.channels {
2550 conn.execute(
2551 "INSERT INTO channel_pts (channel_id,pts) VALUES (?1,?2)",
2552 libsql::params![cid, cpts],
2553 )
2554 .await?;
2555 }
2556 Ok(())
2557 }
2558 UpdateStateChange::Primary { pts, date, seq } => conn
2559 .execute(
2560 "INSERT INTO update_state (id,pts,qts,date,seq) VALUES (1,?1,0,?2,?3)
2561 ON CONFLICT(id) DO UPDATE SET pts=excluded.pts,date=excluded.date,
2562 seq=excluded.seq",
2563 libsql::params![pts, date, seq],
2564 )
2565 .await
2566 .map(|_| ()),
2567 UpdateStateChange::Secondary { qts } => conn
2568 .execute(
2569 "INSERT INTO update_state (id,pts,qts,date,seq) VALUES (1,0,?1,0,0)
2570 ON CONFLICT(id) DO UPDATE SET qts=excluded.qts",
2571 libsql::params![qts],
2572 )
2573 .await
2574 .map(|_| ()),
2575 UpdateStateChange::Channel { id, pts } => conn
2576 .execute(
2577 "INSERT INTO channel_pts (channel_id,pts) VALUES (?1,?2)
2578 ON CONFLICT(channel_id) DO UPDATE SET pts=excluded.pts",
2579 libsql::params![id, pts],
2580 )
2581 .await
2582 .map(|_| ()),
2583 }
2584 })
2585 }
2586
2587 fn cache_peer(&self, peer: &CachedPeer) -> io::Result<()> {
2588 let conn = self.conn.clone();
2589 let (id, hash, is_ch, is_ct, kind) = (
2590 peer.id,
2591 peer.access_hash,
2592 peer.is_channel as i32,
2593 peer.is_chat as i32,
2594 peer.channel_kind.map(|k| k.to_byte() as i32),
2595 );
2596 Self::block(async move {
2597 conn.execute(
2598 "INSERT INTO peers (id,access_hash,is_channel,is_chat,channel_kind) VALUES (?1,?2,?3,?4,?5)
2599 ON CONFLICT(id) DO UPDATE SET
2600 access_hash=excluded.access_hash,
2601 is_channel=excluded.is_channel,
2602 is_chat=excluded.is_chat,
2603 channel_kind=excluded.channel_kind",
2604 libsql::params![id, hash, is_ch, is_ct, kind],
2605 )
2606 .await
2607 .map(|_| ())
2608 })
2609 }
2610}
2611
2612#[cfg(all(test, feature = "libsql-session"))]
2613mod libsql_backend_tests {
2614 use super::*;
2615
2616 async fn open_in_memory() -> libsql::Connection {
2617 let db = libsql::Builder::new_local(":memory:")
2618 .build()
2619 .await
2620 .unwrap();
2621 let conn = db.connect().unwrap();
2622 LibSqlBackend::apply_schema(&conn).await.unwrap();
2623 conn
2624 }
2625
2626 #[tokio::test]
2627 async fn future_auth_token_roundtrip() {
2628 let conn = open_in_memory().await;
2629 let mut s = PersistedSession::default();
2630 s.home_dc_id = 2;
2631 s.future_auth_token = Some(vec![1, 2, 3, 4, 5]);
2632 LibSqlBackend::write_session_async(&conn, &s).await.unwrap();
2633
2634 let loaded = LibSqlBackend::read_session_async(&conn).await.unwrap();
2635 assert_eq!(loaded.future_auth_token, Some(vec![1, 2, 3, 4, 5]));
2636 }
2637
2638 #[tokio::test]
2639 async fn future_auth_token_cleared_on_none() {
2640 let conn = open_in_memory().await;
2641 let mut s = PersistedSession::default();
2642 s.home_dc_id = 2;
2643 s.future_auth_token = Some(vec![9, 9, 9]);
2644 LibSqlBackend::write_session_async(&conn, &s).await.unwrap();
2645 let loaded = LibSqlBackend::read_session_async(&conn).await.unwrap();
2646 assert_eq!(loaded.future_auth_token, Some(vec![9, 9, 9]));
2647
2648 s.future_auth_token = None;
2649 LibSqlBackend::write_session_async(&conn, &s).await.unwrap();
2650 let loaded = LibSqlBackend::read_session_async(&conn).await.unwrap();
2651 assert_eq!(loaded.future_auth_token, None);
2652 }
2653}