1#![deny(unsafe_code)]
16#![cfg_attr(docsrs, feature(doc_cfg))]
17#![doc(html_root_url = "https://docs.rs/ferogram-session/0.6.5")]
18use std::collections::HashMap;
114use std::io::{self, ErrorKind};
115use std::path::Path;
116
117#[cfg(feature = "serde")]
118mod auth_key_serde {
119 use serde::{Deserialize, Deserializer, Serializer};
120
121 pub fn serialize<S>(value: &Option<[u8; 256]>, s: S) -> Result<S::Ok, S::Error>
124 where
125 S: Serializer,
126 {
127 match value {
128 Some(k) => s.serialize_some(k.as_slice()),
129 None => s.serialize_none(),
130 }
131 }
132
133 pub fn deserialize<'de, D>(d: D) -> Result<Option<[u8; 256]>, D::Error>
136 where
137 D: Deserializer<'de>,
138 {
139 let opt: Option<Vec<u8>> = Option::deserialize(d)?;
140 match opt {
141 None => Ok(None),
142 Some(v) => {
143 let arr: [u8; 256] = v
144 .try_into()
145 .map_err(|_| serde::de::Error::custom("auth_key must be exactly 256 bytes"))?;
146 Ok(Some(arr))
147 }
148 }
149 }
150}
151
152#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
156#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
157pub struct DcFlags(pub u8);
158
159impl DcFlags {
160 pub const NONE: DcFlags = DcFlags(0);
161 pub const IPV6: DcFlags = DcFlags(1 << 0);
162 pub const MEDIA_ONLY: DcFlags = DcFlags(1 << 1);
163 pub const TCPO_ONLY: DcFlags = DcFlags(1 << 2);
164 pub const CDN: DcFlags = DcFlags(1 << 3);
165 pub const STATIC: DcFlags = DcFlags(1 << 4);
166
167 pub fn contains(self, other: DcFlags) -> bool {
169 self.0 & other.0 == other.0
170 }
171
172 pub fn set(&mut self, flag: DcFlags) {
174 self.0 |= flag.0;
175 }
176}
177
178impl std::ops::BitOr for DcFlags {
179 type Output = DcFlags;
180 fn bitor(self, rhs: DcFlags) -> DcFlags {
181 DcFlags(self.0 | rhs.0)
182 }
183}
184
185#[derive(Clone, Debug)]
187#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
188pub struct DcEntry {
189 pub dc_id: i32,
190 pub addr: String,
191 #[cfg_attr(feature = "serde", serde(with = "auth_key_serde"))]
192 pub auth_key: Option<[u8; 256]>,
193 pub first_salt: i64,
194 pub time_offset: i32,
195 pub flags: DcFlags,
197}
198
199impl DcEntry {
200 #[inline]
202 pub fn is_ipv6(&self) -> bool {
203 self.flags.contains(DcFlags::IPV6)
204 }
205
206 pub fn socket_addr(&self) -> io::Result<std::net::SocketAddr> {
213 self.addr.parse::<std::net::SocketAddr>().map_err(|_| {
214 io::Error::new(
215 io::ErrorKind::InvalidData,
216 format!("invalid DC address: {:?}", self.addr),
217 )
218 })
219 }
220
221 pub fn from_parts(dc_id: i32, ip: &str, port: u16, flags: DcFlags) -> Self {
234 let addr = if ip.contains(':') {
236 format!("[{ip}]:{port}")
237 } else {
238 format!("{ip}:{port}")
239 };
240 Self {
241 dc_id,
242 addr,
243 auth_key: None,
244 first_salt: 0,
245 time_offset: 0,
246 flags,
247 }
248 }
249}
250
251#[derive(Clone, Debug, Default)]
254#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
255pub struct UpdatesStateSnap {
256 pub pts: i32,
258 pub qts: i32,
260 pub date: i32,
262 pub seq: i32,
264 pub channels: Vec<(i64, i32)>,
266}
267
268impl UpdatesStateSnap {
269 #[inline]
271 pub fn is_initialised(&self) -> bool {
272 self.pts > 0
273 }
274
275 pub fn set_channel_pts(&mut self, channel_id: i64, pts: i32) {
277 if let Some(entry) = self.channels.iter_mut().find(|c| c.0 == channel_id) {
278 entry.1 = pts;
279 } else {
280 self.channels.push((channel_id, pts));
281 }
282 }
283
284 pub fn channel_pts(&self, channel_id: i64) -> i32 {
286 self.channels
287 .iter()
288 .find(|c| c.0 == channel_id)
289 .map(|c| c.1)
290 .unwrap_or(0)
291 }
292}
293
294#[derive(Clone, Copy, Debug, PartialEq, Eq)]
297#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
298#[repr(u8)]
299pub enum ChannelKind {
300 Broadcast = 0,
302 Megagroup = 1,
304 Gigagroup = 2,
306}
307
308impl ChannelKind {
309 pub fn from_byte(b: u8) -> Option<Self> {
312 match b {
313 0 => Some(Self::Broadcast),
314 1 => Some(Self::Megagroup),
315 2 => Some(Self::Gigagroup),
316 _ => None,
317 }
318 }
319
320 pub fn to_byte(self) -> u8 {
322 self as u8
323 }
324}
325
326#[derive(Clone, Debug)]
329#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
330pub struct CachedPeer {
331 pub id: i64,
333 pub access_hash: i64,
336 pub is_channel: bool,
338 pub is_chat: bool,
341 pub channel_kind: Option<ChannelKind>,
344 pub is_community: bool,
349}
350
351#[derive(Clone, Debug)]
355#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
356pub struct CachedMinPeer {
357 pub user_id: i64,
359 pub peer_id: i64,
361 pub msg_id: i32,
363}
364
365#[derive(Clone, Debug, Default)]
367#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
368pub struct PersistedSession {
369 pub home_dc_id: i32,
370 pub dcs: Vec<DcEntry>,
371 pub updates_state: UpdatesStateSnap,
373 pub peers: Vec<CachedPeer>,
376 pub min_peers: Vec<CachedMinPeer>,
379 pub future_auth_token: Option<Vec<u8>>,
386}
387
388impl PersistedSession {
389 pub fn to_bytes(&self) -> Vec<u8> {
391 let mut b = Vec::with_capacity(512);
392
393 b.push(0x08u8); b.extend_from_slice(&self.home_dc_id.to_le_bytes());
396
397 b.push(self.dcs.len() as u8);
398 for d in &self.dcs {
399 b.extend_from_slice(&d.dc_id.to_le_bytes());
400 match &d.auth_key {
401 Some(k) => {
402 b.push(1);
403 b.extend_from_slice(k);
404 }
405 None => {
406 b.push(0);
407 }
408 }
409 b.extend_from_slice(&d.first_salt.to_le_bytes());
410 b.extend_from_slice(&d.time_offset.to_le_bytes());
411 let ab = d.addr.as_bytes();
412 b.push(ab.len() as u8);
413 b.extend_from_slice(ab);
414 b.push(d.flags.0);
415 }
416
417 b.extend_from_slice(&self.updates_state.pts.to_le_bytes());
418 b.extend_from_slice(&self.updates_state.qts.to_le_bytes());
419 b.extend_from_slice(&self.updates_state.date.to_le_bytes());
420 b.extend_from_slice(&self.updates_state.seq.to_le_bytes());
421 let ch = &self.updates_state.channels;
422 b.extend_from_slice(&(ch.len() as u16).to_le_bytes());
423 for &(cid, cpts) in ch {
424 b.extend_from_slice(&cid.to_le_bytes());
425 b.extend_from_slice(&cpts.to_le_bytes());
426 }
427
428 b.extend_from_slice(&(self.peers.len() as u16).to_le_bytes());
431 for p in &self.peers {
432 b.extend_from_slice(&p.id.to_le_bytes());
433 b.extend_from_slice(&p.access_hash.to_le_bytes());
434 let peer_type: u8 = if p.is_chat {
435 2
436 } else if p.is_channel {
437 1
438 } else {
439 0
440 };
441 b.push(peer_type);
442 b.push(p.channel_kind.map(|k| k.to_byte()).unwrap_or(0xFF));
444 b.push(p.is_community as u8);
446 }
447
448 b.extend_from_slice(&(self.min_peers.len() as u16).to_le_bytes());
449 for m in &self.min_peers {
450 b.extend_from_slice(&m.user_id.to_le_bytes());
451 b.extend_from_slice(&m.peer_id.to_le_bytes());
452 b.extend_from_slice(&m.msg_id.to_le_bytes());
453 }
454
455 match &self.future_auth_token {
457 Some(t) => {
458 b.push(1);
459 b.extend_from_slice(&(t.len() as u16).to_le_bytes());
460 b.extend_from_slice(t);
461 }
462 None => b.push(0),
463 }
464
465 b
466 }
467
468 pub fn save(&self, path: &Path) -> io::Result<()> {
475 use std::sync::atomic::{AtomicU64, Ordering};
476 static SEQ: AtomicU64 = AtomicU64::new(0);
477 let n = SEQ.fetch_add(1, Ordering::Relaxed);
478 let tmp = path.with_extension(format!("{n}.tmp"));
479 std::fs::write(&tmp, self.to_bytes())?;
480 std::fs::rename(&tmp, path).inspect_err(|_e| {
481 let _ = std::fs::remove_file(&tmp);
482 })
483 }
484
485 pub fn from_bytes(buf: &[u8]) -> io::Result<Self> {
487 if buf.is_empty() {
488 return Err(io::Error::new(ErrorKind::InvalidData, "empty session data"));
489 }
490
491 let mut p = 0usize;
492
493 macro_rules! r {
494 ($n:expr) => {{
495 if p + $n > buf.len() {
496 return Err(io::Error::new(ErrorKind::InvalidData, "truncated session"));
497 }
498 let s = &buf[p..p + $n];
499 p += $n;
500 s
501 }};
502 }
503 macro_rules! r_i32 {
504 () => {
505 i32::from_le_bytes(r!(4).try_into().unwrap())
506 };
507 }
508 macro_rules! r_i64 {
509 () => {
510 i64::from_le_bytes(r!(8).try_into().unwrap())
511 };
512 }
513 macro_rules! r_u8 {
514 () => {
515 r!(1)[0]
516 };
517 }
518 macro_rules! r_u16 {
519 () => {
520 u16::from_le_bytes(r!(2).try_into().unwrap())
521 };
522 }
523
524 let first_byte = r_u8!();
525
526 let (home_dc_id, version) = if first_byte == 0x08 {
527 (r_i32!(), 8u8)
528 } else if first_byte == 0x07 {
529 (r_i32!(), 7u8)
530 } else if first_byte == 0x06 {
531 (r_i32!(), 6u8)
532 } else if first_byte == 0x05 {
533 (r_i32!(), 5u8)
534 } else if first_byte == 0x04 {
535 (r_i32!(), 4u8)
536 } else if first_byte == 0x03 {
537 (r_i32!(), 3u8)
538 } else if first_byte == 0x02 {
539 (r_i32!(), 2u8)
540 } else {
541 let rest = r!(3);
542 let mut bytes = [0u8; 4];
543 bytes[0] = first_byte;
544 bytes[1..4].copy_from_slice(rest);
545 (i32::from_le_bytes(bytes), 1u8)
546 };
547
548 let dc_count = r_u8!() as usize;
549 let mut dcs = Vec::with_capacity(dc_count);
550 for _ in 0..dc_count {
551 let dc_id = r_i32!();
552 let has_key = r_u8!();
553 let auth_key = if has_key == 1 {
554 let mut k = [0u8; 256];
555 k.copy_from_slice(r!(256));
556 Some(k)
557 } else {
558 None
559 };
560 let first_salt = r_i64!();
561 let time_offset = r_i32!();
562 let al = r_u8!() as usize;
563 let addr = String::from_utf8_lossy(r!(al)).into_owned();
564 let flags = if version >= 3 {
565 DcFlags(r_u8!())
566 } else {
567 DcFlags::NONE
568 };
569 dcs.push(DcEntry {
570 dc_id,
571 addr,
572 auth_key,
573 first_salt,
574 time_offset,
575 flags,
576 });
577 }
578
579 if version < 2 {
580 return Ok(Self {
581 home_dc_id,
582 dcs,
583 updates_state: UpdatesStateSnap::default(),
584 peers: Vec::new(),
585 min_peers: Vec::new(),
586 future_auth_token: None,
587 });
588 }
589
590 let pts = r_i32!();
591 let qts = r_i32!();
592 let date = r_i32!();
593 let seq = r_i32!();
594 let ch_count = r_u16!() as usize;
595 let mut channels = Vec::with_capacity(ch_count);
596 for _ in 0..ch_count {
597 let cid = r_i64!();
598 let cpts = r_i32!();
599 channels.push((cid, cpts));
600 }
601
602 let peer_count = r_u16!() as usize;
603 let mut peers = Vec::with_capacity(peer_count);
604 for _ in 0..peer_count {
605 let id = r_i64!();
606 let access_hash = r_i64!();
607 let peer_type = r_u8!();
609 let is_channel = peer_type == 1;
610 let is_chat = peer_type == 2;
611 let channel_kind = if version >= 6 {
613 let kb = r_u8!();
614 if kb == 0xFF {
615 None
616 } else {
617 ChannelKind::from_byte(kb)
618 }
619 } else {
620 None
621 };
622 let is_community = if version >= 8 { r_u8!() != 0 } else { false };
627 peers.push(CachedPeer {
628 id,
629 access_hash,
630 is_channel,
631 is_chat,
632 channel_kind,
633 is_community,
634 });
635 }
636
637 let min_peers = if version >= 4 {
639 let count = r_u16!() as usize;
640 let mut v = Vec::with_capacity(count);
641 for _ in 0..count {
642 let user_id = r_i64!();
643 let peer_id = r_i64!();
644 let msg_id = r_i32!();
645 v.push(CachedMinPeer {
646 user_id,
647 peer_id,
648 msg_id,
649 });
650 }
651 v
652 } else {
653 Vec::new()
654 };
655
656 let future_auth_token = if version >= 7 {
658 let has_token = r_u8!();
659 if has_token == 1 {
660 let tl = r_u16!() as usize;
661 Some(r!(tl).to_vec())
662 } else {
663 None
664 }
665 } else {
666 None
667 };
668 let _ = p;
672
673 Ok(Self {
674 home_dc_id,
675 dcs,
676 updates_state: UpdatesStateSnap {
677 pts,
678 qts,
679 date,
680 seq,
681 channels,
682 },
683 peers,
684 min_peers,
685 future_auth_token,
686 })
687 }
688
689 pub fn from_string(s: &str) -> io::Result<Self> {
691 use base64::Engine as _;
692 let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
693 .decode(s.trim())
694 .map_err(|e| io::Error::new(ErrorKind::InvalidData, e))?;
695 Self::from_bytes(&bytes)
696 }
697
698 pub fn load(path: &Path) -> io::Result<Self> {
701 let buf = std::fs::read(path)?;
702 Self::from_bytes(&buf)
703 }
704
705 pub fn dc_for(&self, dc_id: i32, prefer_ipv6: bool) -> Option<&DcEntry> {
717 let mut candidates = self.dcs.iter().filter(|d| d.dc_id == dc_id).peekable();
718 candidates.peek()?;
719 let cands: Vec<&DcEntry> = self.dcs.iter().filter(|d| d.dc_id == dc_id).collect();
721 cands
723 .iter()
724 .copied()
725 .find(|d| d.is_ipv6() == prefer_ipv6)
726 .or_else(|| cands.first().copied())
727 }
728
729 pub fn all_dcs_for(&self, dc_id: i32) -> impl Iterator<Item = &DcEntry> {
734 self.dcs.iter().filter(move |d| d.dc_id == dc_id)
735 }
736
737 pub fn stats(&self) -> SessionStats {
746 let users = self
747 .peers
748 .iter()
749 .filter(|p| !p.is_channel && !p.is_chat)
750 .count();
751 let channels = self.peers.iter().filter(|p| p.is_channel).count();
752 let chats = self.peers.iter().filter(|p| p.is_chat).count();
753
754 SessionStats {
755 home_dc_id: self.home_dc_id,
756 dcs_total: self.dcs.len(),
757 dcs_with_auth_key: self.dcs.iter().filter(|d| d.auth_key.is_some()).count(),
758 peers_total: self.peers.len(),
759 peers_users: users,
760 peers_channels: channels,
761 peers_chats: chats,
762 min_peers: self.min_peers.len(),
763 channel_pts_entries: self.updates_state.channels.len(),
764 updates_initialised: self.updates_state.is_initialised(),
765 approx_bytes: self.to_bytes().len(),
766 }
767 }
768}
769
770#[derive(Clone, Debug)]
774pub struct SessionStats {
775 pub home_dc_id: i32,
777 pub dcs_total: usize,
779 pub dcs_with_auth_key: usize,
781 pub peers_total: usize,
783 pub peers_users: usize,
785 pub peers_channels: usize,
787 pub peers_chats: usize,
789 pub min_peers: usize,
792 pub channel_pts_entries: usize,
794 pub updates_initialised: bool,
796 pub approx_bytes: usize,
798}
799
800impl std::fmt::Display for PersistedSession {
801 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
802 use base64::Engine as _;
803 f.write_str(&base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(self.to_bytes()))
804 }
805}
806
807pub fn default_dc_addresses() -> HashMap<i32, String> {
809 [
810 (1, "149.154.175.53:443"),
811 (2, "149.154.167.51:443"),
812 (3, "149.154.175.100:443"),
813 (4, "149.154.167.91:443"),
814 (5, "91.108.56.130:443"),
815 ]
816 .into_iter()
817 .map(|(id, addr)| (id, addr.to_string()))
818 .collect()
819}
820
821use std::path::PathBuf;
826
827pub trait SessionBackend: Send + Sync {
835 fn save(&self, session: &PersistedSession) -> io::Result<()>;
836 fn load(&self) -> io::Result<Option<PersistedSession>>;
837 fn delete(&self) -> io::Result<()>;
838
839 fn name(&self) -> &str;
841
842 fn update_dc(&self, entry: &DcEntry) -> io::Result<()> {
855 let mut s = self.load()?.unwrap_or_default();
856 if let Some(existing) = s
858 .dcs
859 .iter_mut()
860 .find(|d| d.dc_id == entry.dc_id && d.is_ipv6() == entry.is_ipv6())
861 {
862 *existing = entry.clone();
863 } else {
864 s.dcs.push(entry.clone());
865 }
866 self.save(&s)
867 }
868
869 fn set_home_dc(&self, dc_id: i32) -> io::Result<()> {
875 let mut s = self.load()?.unwrap_or_default();
876 s.home_dc_id = dc_id;
877 self.save(&s)
878 }
879
880 fn apply_update_state(&self, update: UpdateStateChange) -> io::Result<()> {
885 let mut s = self.load()?.unwrap_or_default();
886 update.apply_to(&mut s.updates_state);
887 self.save(&s)
888 }
889
890 fn cache_peer(&self, peer: &CachedPeer) -> io::Result<()> {
896 let mut s = self.load()?.unwrap_or_default();
897 if let Some(existing) = s.peers.iter_mut().find(|p| p.id == peer.id) {
898 *existing = peer.clone();
899 } else {
900 s.peers.push(peer.clone());
901 }
902 self.save(&s)
903 }
904}
905
906#[derive(Debug, Clone)]
921pub enum UpdateStateChange {
922 All(UpdatesStateSnap),
924 Primary { pts: i32, date: i32, seq: i32 },
926 Secondary { qts: i32 },
928 Channel { id: i64, pts: i32 },
930}
931
932impl UpdateStateChange {
933 pub fn apply_to(&self, snap: &mut UpdatesStateSnap) {
935 match self {
936 Self::All(new_snap) => *snap = new_snap.clone(),
937 Self::Primary { pts, date, seq } => {
938 snap.pts = *pts;
939 snap.date = *date;
940 snap.seq = *seq;
941 }
942 Self::Secondary { qts } => {
943 snap.qts = *qts;
944 }
945 Self::Channel { id, pts } => {
946 if let Some(existing) = snap.channels.iter_mut().find(|c| c.0 == *id) {
948 existing.1 = *pts;
949 } else {
950 snap.channels.push((*id, *pts));
951 }
952 }
953 }
954 }
955}
956
957pub struct BinaryFileBackend {
961 path: PathBuf,
962 write_lock: std::sync::Mutex<()>,
967}
968
969impl BinaryFileBackend {
970 pub fn new(path: impl Into<PathBuf>) -> Self {
973 Self {
974 path: path.into(),
975 write_lock: std::sync::Mutex::new(()),
976 }
977 }
978
979 pub fn path(&self) -> &std::path::Path {
981 &self.path
982 }
983}
984
985impl SessionBackend for BinaryFileBackend {
986 fn save(&self, session: &PersistedSession) -> io::Result<()> {
987 let _guard = self.write_lock.lock().unwrap();
988 session.save(&self.path)
989 }
990
991 fn load(&self) -> io::Result<Option<PersistedSession>> {
992 if !self.path.exists() {
993 return Ok(None);
994 }
995 match PersistedSession::load(&self.path) {
996 Ok(s) => Ok(Some(s)),
997 Err(e) => {
998 let bak = self.path.with_extension("bak");
999 tracing::warn!(
1000 "[ferogram::session] session file {:?} could not be read ({e}); \
1001 backing it up to {:?} and starting a new session",
1002 self.path,
1003 bak
1004 );
1005 let _ = std::fs::rename(&self.path, &bak);
1006 Ok(None)
1007 }
1008 }
1009 }
1010
1011 fn delete(&self) -> io::Result<()> {
1012 if self.path.exists() {
1013 std::fs::remove_file(&self.path)?;
1014 }
1015 Ok(())
1016 }
1017
1018 fn name(&self) -> &str {
1019 "binary-file"
1020 }
1021
1022 }
1025
1026#[derive(Default)]
1034pub struct InMemoryBackend {
1035 data: std::sync::Mutex<Option<PersistedSession>>,
1036}
1037
1038impl InMemoryBackend {
1039 pub fn new() -> Self {
1041 Self::default()
1042 }
1043
1044 pub fn snapshot(&self) -> Option<PersistedSession> {
1046 self.data.lock().unwrap().clone()
1047 }
1048}
1049
1050impl SessionBackend for InMemoryBackend {
1051 fn save(&self, s: &PersistedSession) -> io::Result<()> {
1052 *self.data.lock().unwrap() = Some(s.clone());
1053 Ok(())
1054 }
1055
1056 fn load(&self) -> io::Result<Option<PersistedSession>> {
1057 Ok(self.data.lock().unwrap().clone())
1058 }
1059
1060 fn delete(&self) -> io::Result<()> {
1061 *self.data.lock().unwrap() = None;
1062 Ok(())
1063 }
1064
1065 fn name(&self) -> &str {
1066 "in-memory"
1067 }
1068
1069 fn update_dc(&self, entry: &DcEntry) -> io::Result<()> {
1072 let mut guard = self.data.lock().unwrap();
1073 let s = guard.get_or_insert_with(PersistedSession::default);
1074 if let Some(existing) = s
1075 .dcs
1076 .iter_mut()
1077 .find(|d| d.dc_id == entry.dc_id && d.is_ipv6() == entry.is_ipv6())
1078 {
1079 *existing = entry.clone();
1080 } else {
1081 s.dcs.push(entry.clone());
1082 }
1083 Ok(())
1084 }
1085
1086 fn set_home_dc(&self, dc_id: i32) -> io::Result<()> {
1087 let mut guard = self.data.lock().unwrap();
1088 let s = guard.get_or_insert_with(PersistedSession::default);
1089 s.home_dc_id = dc_id;
1090 Ok(())
1091 }
1092
1093 fn apply_update_state(&self, update: UpdateStateChange) -> io::Result<()> {
1094 let mut guard = self.data.lock().unwrap();
1095 let s = guard.get_or_insert_with(PersistedSession::default);
1096 update.apply_to(&mut s.updates_state);
1097 Ok(())
1098 }
1099
1100 fn cache_peer(&self, peer: &CachedPeer) -> io::Result<()> {
1101 let mut guard = self.data.lock().unwrap();
1102 let s = guard.get_or_insert_with(PersistedSession::default);
1103 if let Some(existing) = s.peers.iter_mut().find(|p| p.id == peer.id) {
1104 *existing = peer.clone();
1105 } else {
1106 s.peers.push(peer.clone());
1107 }
1108 Ok(())
1109 }
1110}
1111
1112pub struct StringSessionBackend {
1116 data: std::sync::Mutex<String>,
1117}
1118
1119impl StringSessionBackend {
1120 pub fn new(s: impl Into<String>) -> Self {
1123 Self {
1124 data: std::sync::Mutex::new(s.into()),
1125 }
1126 }
1127
1128 pub fn current(&self) -> String {
1130 self.data.lock().unwrap().clone()
1131 }
1132}
1133
1134impl SessionBackend for StringSessionBackend {
1135 fn save(&self, session: &PersistedSession) -> io::Result<()> {
1136 *self.data.lock().unwrap() = session.to_string();
1137 Ok(())
1138 }
1139
1140 fn load(&self) -> io::Result<Option<PersistedSession>> {
1141 let s = self.data.lock().unwrap().clone();
1142 if s.trim().is_empty() {
1143 return Ok(None);
1144 }
1145 PersistedSession::from_string(&s).map(Some)
1146 }
1147
1148 fn delete(&self) -> io::Result<()> {
1149 *self.data.lock().unwrap() = String::new();
1150 Ok(())
1151 }
1152
1153 fn name(&self) -> &str {
1154 "string-session"
1155 }
1156}
1157
1158pub mod string_session;
1166pub use string_session::{FullSession, Session, StringSession, StringSessionError};
1167
1168#[cfg(test)]
1169mod tests {
1170 use super::*;
1171
1172 fn make_dc(id: i32) -> DcEntry {
1173 DcEntry {
1174 dc_id: id,
1175 addr: format!("1.2.3.{id}:443"),
1176 auth_key: None,
1177 first_salt: 0,
1178 time_offset: 0,
1179 flags: DcFlags::NONE,
1180 }
1181 }
1182
1183 fn make_peer(id: i64, hash: i64) -> CachedPeer {
1184 CachedPeer {
1185 id,
1186 access_hash: hash,
1187 is_channel: false,
1188 is_chat: false,
1189 channel_kind: None,
1190 is_community: false,
1191 }
1192 }
1193
1194 #[test]
1197 fn inmemory_load_returns_none_when_empty() {
1198 let b = InMemoryBackend::new();
1199 assert!(b.load().unwrap().is_none());
1200 }
1201
1202 #[test]
1203 fn inmemory_save_then_load_round_trips() {
1204 let b = InMemoryBackend::new();
1205 let mut s = PersistedSession::default();
1206 s.home_dc_id = 3;
1207 s.dcs.push(make_dc(3));
1208 b.save(&s).unwrap();
1209
1210 let loaded = b.load().unwrap().unwrap();
1211 assert_eq!(loaded.home_dc_id, 3);
1212 assert_eq!(loaded.dcs.len(), 1);
1213 }
1214
1215 #[test]
1216 fn inmemory_delete_clears_state() {
1217 let b = InMemoryBackend::new();
1218 let mut s = PersistedSession::default();
1219 s.home_dc_id = 2;
1220 b.save(&s).unwrap();
1221 b.delete().unwrap();
1222 assert!(b.load().unwrap().is_none());
1223 }
1224
1225 #[test]
1228 fn inmemory_update_dc_inserts_new() {
1229 let b = InMemoryBackend::new();
1230 b.update_dc(&make_dc(4)).unwrap();
1231 let s = b.snapshot().unwrap();
1232 assert_eq!(s.dcs.len(), 1);
1233 assert_eq!(s.dcs[0].dc_id, 4);
1234 }
1235
1236 #[test]
1237 fn inmemory_update_dc_replaces_existing() {
1238 let b = InMemoryBackend::new();
1239 b.update_dc(&make_dc(2)).unwrap();
1240 let mut updated = make_dc(2);
1241 updated.addr = "9.9.9.9:443".to_string();
1242 b.update_dc(&updated).unwrap();
1243
1244 let s = b.snapshot().unwrap();
1245 assert_eq!(s.dcs.len(), 1);
1246 assert_eq!(s.dcs[0].addr, "9.9.9.9:443");
1247 }
1248
1249 #[test]
1250 fn inmemory_set_home_dc() {
1251 let b = InMemoryBackend::new();
1252 b.set_home_dc(5).unwrap();
1253 assert_eq!(b.snapshot().unwrap().home_dc_id, 5);
1254 }
1255
1256 #[test]
1257 fn inmemory_cache_peer_inserts() {
1258 let b = InMemoryBackend::new();
1259 b.cache_peer(&make_peer(100, 0xdeadbeef)).unwrap();
1260 let s = b.snapshot().unwrap();
1261 assert_eq!(s.peers.len(), 1);
1262 assert_eq!(s.peers[0].id, 100);
1263 }
1264
1265 #[test]
1266 fn inmemory_cache_peer_updates_existing() {
1267 let b = InMemoryBackend::new();
1268 b.cache_peer(&make_peer(100, 111)).unwrap();
1269 b.cache_peer(&make_peer(100, 222)).unwrap();
1270 let s = b.snapshot().unwrap();
1271 assert_eq!(s.peers.len(), 1);
1272 assert_eq!(s.peers[0].access_hash, 222);
1273 }
1274
1275 #[test]
1278 fn update_state_primary() {
1279 let mut snap = UpdatesStateSnap {
1280 pts: 0,
1281 qts: 0,
1282 date: 0,
1283 seq: 0,
1284 channels: vec![],
1285 };
1286 UpdateStateChange::Primary {
1287 pts: 10,
1288 date: 20,
1289 seq: 30,
1290 }
1291 .apply_to(&mut snap);
1292 assert_eq!(snap.pts, 10);
1293 assert_eq!(snap.date, 20);
1294 assert_eq!(snap.seq, 30);
1295 assert_eq!(snap.qts, 0); }
1297
1298 #[test]
1299 fn update_state_secondary() {
1300 let mut snap = UpdatesStateSnap {
1301 pts: 5,
1302 qts: 0,
1303 date: 0,
1304 seq: 0,
1305 channels: vec![],
1306 };
1307 UpdateStateChange::Secondary { qts: 99 }.apply_to(&mut snap);
1308 assert_eq!(snap.qts, 99);
1309 assert_eq!(snap.pts, 5); }
1311
1312 #[test]
1313 fn update_state_channel_inserts() {
1314 let mut snap = UpdatesStateSnap {
1315 pts: 0,
1316 qts: 0,
1317 date: 0,
1318 seq: 0,
1319 channels: vec![],
1320 };
1321 UpdateStateChange::Channel { id: 12345, pts: 42 }.apply_to(&mut snap);
1322 assert_eq!(snap.channels, vec![(12345, 42)]);
1323 }
1324
1325 #[test]
1326 fn update_state_channel_updates_existing() {
1327 let mut snap = UpdatesStateSnap {
1328 pts: 0,
1329 qts: 0,
1330 date: 0,
1331 seq: 0,
1332 channels: vec![(12345, 10), (67890, 5)],
1333 };
1334 UpdateStateChange::Channel { id: 12345, pts: 99 }.apply_to(&mut snap);
1335 assert_eq!(snap.channels[0], (12345, 99));
1337 assert_eq!(snap.channels[1], (67890, 5));
1338 }
1339
1340 #[test]
1341 fn apply_update_state_via_backend() {
1342 let b = InMemoryBackend::new();
1343 b.apply_update_state(UpdateStateChange::Primary {
1344 pts: 7,
1345 date: 8,
1346 seq: 9,
1347 })
1348 .unwrap();
1349 let s = b.snapshot().unwrap();
1350 assert_eq!(s.updates_state.pts, 7);
1351 }
1352
1353 #[test]
1356 fn default_update_dc_via_trait_object() {
1357 let b: Box<dyn SessionBackend> = Box::new(InMemoryBackend::new());
1358 b.update_dc(&make_dc(1)).unwrap();
1359 b.update_dc(&make_dc(2)).unwrap();
1360 let loaded = b.load().unwrap().unwrap();
1362 assert_eq!(loaded.dcs.len(), 2);
1363 }
1364
1365 fn make_dc_v6(id: i32) -> DcEntry {
1368 DcEntry {
1369 dc_id: id,
1370 addr: format!("[2001:b28:f23d:f00{}::a]:443", id),
1371 auth_key: None,
1372 first_salt: 0,
1373 time_offset: 0,
1374 flags: DcFlags::IPV6,
1375 }
1376 }
1377
1378 #[test]
1379 fn dc_entry_from_parts_ipv4() {
1380 let dc = DcEntry::from_parts(1, "149.154.175.53", 443, DcFlags::NONE);
1381 assert_eq!(dc.addr, "149.154.175.53:443");
1382 assert!(!dc.is_ipv6());
1383 let sa = dc.socket_addr().unwrap();
1384 assert_eq!(sa.port(), 443);
1385 }
1386
1387 #[test]
1388 fn dc_entry_from_parts_ipv6() {
1389 let dc = DcEntry::from_parts(2, "2001:b28:f23d:f001::a", 443, DcFlags::IPV6);
1390 assert_eq!(dc.addr, "[2001:b28:f23d:f001::a]:443");
1391 assert!(dc.is_ipv6());
1392 let sa = dc.socket_addr().unwrap();
1393 assert_eq!(sa.port(), 443);
1394 }
1395
1396 #[test]
1397 fn persisted_session_dc_for_prefers_ipv6() {
1398 let mut s = PersistedSession::default();
1399 s.dcs.push(make_dc(2)); s.dcs.push(make_dc_v6(2)); let v6 = s.dc_for(2, true).unwrap();
1403 assert!(v6.is_ipv6());
1404
1405 let v4 = s.dc_for(2, false).unwrap();
1406 assert!(!v4.is_ipv6());
1407 }
1408
1409 #[test]
1410 fn persisted_session_dc_for_falls_back_when_only_ipv4() {
1411 let mut s = PersistedSession::default();
1412 s.dcs.push(make_dc(3)); let dc = s.dc_for(3, true).unwrap();
1416 assert!(!dc.is_ipv6());
1417 }
1418
1419 #[test]
1420 fn persisted_session_all_dcs_for_returns_both() {
1421 let mut s = PersistedSession::default();
1422 s.dcs.push(make_dc(1));
1423 s.dcs.push(make_dc_v6(1));
1424 s.dcs.push(make_dc(2));
1425
1426 assert_eq!(s.all_dcs_for(1).count(), 2);
1427 assert_eq!(s.all_dcs_for(2).count(), 1);
1428 assert_eq!(s.all_dcs_for(5).count(), 0);
1429 }
1430
1431 #[test]
1432 fn inmemory_ipv4_and_ipv6_coexist() {
1433 let b = InMemoryBackend::new();
1434 b.update_dc(&make_dc(2)).unwrap(); b.update_dc(&make_dc_v6(2)).unwrap(); let s = b.snapshot().unwrap();
1438 assert_eq!(s.dcs.iter().filter(|d| d.dc_id == 2).count(), 2);
1440 }
1441
1442 #[test]
1443 fn binary_roundtrip_ipv4_and_ipv6() {
1444 let mut s = PersistedSession::default();
1445 s.home_dc_id = 2;
1446 s.dcs.push(make_dc(2));
1447 s.dcs.push(make_dc_v6(2));
1448
1449 let bytes = s.to_bytes();
1450 let loaded = PersistedSession::from_bytes(&bytes).unwrap();
1451 assert_eq!(loaded.dcs.len(), 2);
1452 assert_eq!(loaded.dcs.iter().filter(|d| d.is_ipv6()).count(), 1);
1453 assert_eq!(loaded.dcs.iter().filter(|d| !d.is_ipv6()).count(), 1);
1454 }
1455
1456 #[test]
1457 fn v6_channel_kind_roundtrip_all_variants() {
1458 let mut s = PersistedSession::default();
1459 s.home_dc_id = 1;
1460 s.peers.push(CachedPeer {
1461 id: 1001,
1462 access_hash: 0xaaaa,
1463 is_channel: true,
1464 is_chat: false,
1465 channel_kind: Some(ChannelKind::Broadcast),
1466 is_community: false,
1467 });
1468 s.peers.push(CachedPeer {
1469 id: 1002,
1470 access_hash: 0xbbbb,
1471 is_channel: true,
1472 is_chat: false,
1473 channel_kind: Some(ChannelKind::Megagroup),
1474 is_community: false,
1475 });
1476 s.peers.push(CachedPeer {
1477 id: 1003,
1478 access_hash: 0xcccc,
1479 is_channel: true,
1480 is_chat: false,
1481 channel_kind: Some(ChannelKind::Gigagroup),
1482 is_community: false,
1483 });
1484 s.peers.push(CachedPeer {
1485 id: 1004,
1486 access_hash: 0xdddd,
1487 is_channel: false,
1488 is_chat: false,
1489 channel_kind: None,
1490 is_community: false,
1491 });
1492
1493 let bytes = s.to_bytes();
1494 assert_eq!(bytes[0], 0x08, "version byte must be 8");
1495
1496 let loaded = PersistedSession::from_bytes(&bytes).unwrap();
1497 assert_eq!(loaded.peers.len(), 4);
1498
1499 let p = &loaded.peers[0];
1500 assert_eq!(p.id, 1001);
1501 assert_eq!(p.channel_kind, Some(ChannelKind::Broadcast));
1502
1503 let p = &loaded.peers[1];
1504 assert_eq!(p.id, 1002);
1505 assert_eq!(p.channel_kind, Some(ChannelKind::Megagroup));
1506
1507 let p = &loaded.peers[2];
1508 assert_eq!(p.id, 1003);
1509 assert_eq!(p.channel_kind, Some(ChannelKind::Gigagroup));
1510
1511 let p = &loaded.peers[3];
1512 assert_eq!(p.id, 1004);
1513 assert_eq!(p.channel_kind, None);
1514 }
1515
1516 #[test]
1517 fn v7_future_auth_token_roundtrip() {
1518 let mut s = PersistedSession::default();
1519 s.home_dc_id = 2;
1520 s.future_auth_token = Some(vec![1, 2, 3, 4, 5]);
1521
1522 let bytes = s.to_bytes();
1523 assert_eq!(bytes[0], 0x08, "version byte must be 8");
1524
1525 let loaded = PersistedSession::from_bytes(&bytes).unwrap();
1526 assert_eq!(loaded.future_auth_token, Some(vec![1, 2, 3, 4, 5]));
1527 }
1528
1529 #[test]
1530 fn v7_future_auth_token_absent() {
1531 let s = PersistedSession {
1532 home_dc_id: 2,
1533 ..Default::default()
1534 };
1535
1536 let bytes = s.to_bytes();
1537 let loaded = PersistedSession::from_bytes(&bytes).unwrap();
1538 assert_eq!(loaded.future_auth_token, None);
1539 }
1540
1541 #[test]
1542 fn v6_file_without_token_loads_as_none() {
1543 let mut s = PersistedSession::default();
1546 s.home_dc_id = 3;
1547 let mut bytes = s.to_bytes();
1548 bytes[0] = 0x06; bytes.truncate(bytes.len() - 1); let loaded = PersistedSession::from_bytes(&bytes).unwrap();
1552 assert_eq!(loaded.home_dc_id, 3);
1553 assert_eq!(loaded.future_auth_token, None);
1554 }
1555
1556 #[test]
1557 fn v6_channel_kind_absent_sentinel_roundtrip() {
1558 let mut s = PersistedSession::default();
1560 s.home_dc_id = 1;
1561 s.peers.push(CachedPeer {
1562 id: 555,
1563 access_hash: 0x1234,
1564 is_channel: false,
1565 is_chat: false,
1566 channel_kind: None,
1567 is_community: false,
1568 });
1569 let loaded = PersistedSession::from_bytes(&s.to_bytes()).unwrap();
1570 assert_eq!(loaded.peers[0].channel_kind, None);
1571 }
1572
1573 #[test]
1574 fn v8_community_flag_roundtrip() {
1575 let mut s = PersistedSession::default();
1576 s.home_dc_id = 1;
1577 s.peers.push(CachedPeer {
1578 id: 2001,
1579 access_hash: 0xeeee,
1580 is_channel: false,
1581 is_chat: false,
1582 channel_kind: None,
1583 is_community: true,
1584 });
1585 s.peers.push(CachedPeer {
1586 id: 2002,
1587 access_hash: 0xffff,
1588 is_channel: true,
1589 is_chat: false,
1590 channel_kind: Some(ChannelKind::Megagroup),
1591 is_community: false,
1592 });
1593
1594 let bytes = s.to_bytes();
1595 assert_eq!(bytes[0], 0x08, "version byte must be 8");
1596
1597 let loaded = PersistedSession::from_bytes(&bytes).unwrap();
1598 assert_eq!(loaded.peers.len(), 2);
1599 assert!(loaded.peers[0].is_community);
1600 assert!(!loaded.peers[1].is_community);
1601 }
1602
1603 #[test]
1604 fn v7_file_without_community_byte_defaults_false() {
1605 let mut bytes = Vec::new();
1610 bytes.push(0x07u8); bytes.extend_from_slice(&4i32.to_le_bytes()); bytes.push(0); bytes.extend_from_slice(&0i32.to_le_bytes()); bytes.extend_from_slice(&0i32.to_le_bytes()); bytes.extend_from_slice(&0i32.to_le_bytes()); bytes.extend_from_slice(&0i32.to_le_bytes()); bytes.extend_from_slice(&0u16.to_le_bytes()); bytes.extend_from_slice(&1u16.to_le_bytes()); bytes.extend_from_slice(&777i64.to_le_bytes()); bytes.extend_from_slice(&0x99i64.to_le_bytes()); bytes.push(0); bytes.push(0xFF); bytes.extend_from_slice(&0u16.to_le_bytes()); bytes.push(0); let loaded = PersistedSession::from_bytes(&bytes).unwrap();
1628 assert_eq!(loaded.peers.len(), 1);
1629 assert_eq!(loaded.peers[0].id, 777);
1630 assert!(!loaded.peers[0].is_community);
1631
1632 let resaved = loaded.to_bytes();
1634 assert_eq!(resaved[0], 0x08, "resave must upgrade to v8");
1635 }
1636}
1637
1638#[cfg(feature = "sqlite-session")]
1664pub struct SqliteBackend {
1665 conn: std::sync::Mutex<rusqlite::Connection>,
1666 label: String,
1667}
1668
1669#[cfg(feature = "sqlite-session")]
1670impl SqliteBackend {
1671 const SCHEMA: &'static str = "
1672 PRAGMA journal_mode = WAL;
1673 PRAGMA synchronous = NORMAL;
1674
1675 CREATE TABLE IF NOT EXISTS meta (
1676 key TEXT PRIMARY KEY,
1677 value INTEGER NOT NULL DEFAULT 0
1678 );
1679
1680 CREATE TABLE IF NOT EXISTS dcs (
1681 dc_id INTEGER NOT NULL,
1682 flags INTEGER NOT NULL DEFAULT 0,
1683 addr TEXT NOT NULL,
1684 auth_key BLOB,
1685 first_salt INTEGER NOT NULL DEFAULT 0,
1686 time_offset INTEGER NOT NULL DEFAULT 0,
1687 PRIMARY KEY (dc_id, flags)
1688 );
1689
1690 CREATE TABLE IF NOT EXISTS update_state (
1691 id INTEGER PRIMARY KEY CHECK (id = 1),
1692 pts INTEGER NOT NULL DEFAULT 0,
1693 qts INTEGER NOT NULL DEFAULT 0,
1694 date INTEGER NOT NULL DEFAULT 0,
1695 seq INTEGER NOT NULL DEFAULT 0
1696 );
1697
1698 CREATE TABLE IF NOT EXISTS channel_pts (
1699 channel_id INTEGER PRIMARY KEY,
1700 pts INTEGER NOT NULL
1701 );
1702
1703 CREATE TABLE IF NOT EXISTS peers (
1704 id INTEGER PRIMARY KEY,
1705 access_hash INTEGER NOT NULL,
1706 is_channel INTEGER NOT NULL DEFAULT 0,
1707 is_chat INTEGER NOT NULL DEFAULT 0,
1708 channel_kind INTEGER,
1709 is_community INTEGER NOT NULL DEFAULT 0
1710 );
1711
1712 CREATE TABLE IF NOT EXISTS min_peers (
1713 user_id INTEGER PRIMARY KEY,
1714 peer_id INTEGER NOT NULL,
1715 msg_id INTEGER NOT NULL
1716 );
1717 ";
1718
1719 pub fn open(path: impl Into<PathBuf>) -> io::Result<Self> {
1721 let path = path.into();
1722 let label = path.display().to_string();
1723 let conn = rusqlite::Connection::open(&path).map_err(io::Error::other)?;
1724 conn.execute_batch(Self::SCHEMA).map_err(io::Error::other)?;
1725 Self::migrate_legacy_sqlite_schema(&conn)?;
1726 Ok(Self {
1727 conn: std::sync::Mutex::new(conn),
1728 label,
1729 })
1730 }
1731
1732 pub fn in_memory() -> io::Result<Self> {
1734 let conn = rusqlite::Connection::open_in_memory().map_err(io::Error::other)?;
1735 conn.execute_batch(Self::SCHEMA).map_err(io::Error::other)?;
1736 Self::migrate_legacy_sqlite_schema(&conn)?;
1737 Ok(Self {
1738 conn: std::sync::Mutex::new(conn),
1739 label: ":memory:".into(),
1740 })
1741 }
1742
1743 fn map_err(e: rusqlite::Error) -> io::Error {
1744 io::Error::other(e)
1745 }
1746
1747 fn migrate_legacy_sqlite_schema(conn: &rusqlite::Connection) -> io::Result<()> {
1751 let mut has_is_chat = false;
1752 let mut stmt = conn
1753 .prepare("PRAGMA table_info(peers)")
1754 .map_err(Self::map_err)?;
1755 let cols = stmt
1756 .query_map([], |row| row.get::<_, String>(1))
1757 .map_err(Self::map_err)?;
1758 for col in cols.filter_map(|r| r.ok()) {
1759 if col == "is_chat" {
1760 has_is_chat = true;
1761 break;
1762 }
1763 }
1764 if !has_is_chat {
1765 conn.execute_batch("ALTER TABLE peers ADD COLUMN is_chat INTEGER NOT NULL DEFAULT 0;")
1766 .map_err(Self::map_err)?;
1767 }
1768 let mut has_channel_kind = false;
1770 let mut stmt2 = conn
1771 .prepare("PRAGMA table_info(peers)")
1772 .map_err(Self::map_err)?;
1773 let cols2 = stmt2
1774 .query_map([], |row| row.get::<_, String>(1))
1775 .map_err(Self::map_err)?;
1776 for col in cols2.filter_map(|r| r.ok()) {
1777 if col == "channel_kind" {
1778 has_channel_kind = true;
1779 break;
1780 }
1781 }
1782 if !has_channel_kind {
1783 conn.execute_batch("ALTER TABLE peers ADD COLUMN channel_kind INTEGER;")
1784 .map_err(Self::map_err)?;
1785 }
1786 let mut has_is_community = false;
1788 let mut stmt3 = conn
1789 .prepare("PRAGMA table_info(peers)")
1790 .map_err(Self::map_err)?;
1791 let cols3 = stmt3
1792 .query_map([], |row| row.get::<_, String>(1))
1793 .map_err(Self::map_err)?;
1794 for col in cols3.filter_map(|r| r.ok()) {
1795 if col == "is_community" {
1796 has_is_community = true;
1797 break;
1798 }
1799 }
1800 if !has_is_community {
1801 conn.execute_batch(
1802 "ALTER TABLE peers ADD COLUMN is_community INTEGER NOT NULL DEFAULT 0;",
1803 )
1804 .map_err(Self::map_err)?;
1805 }
1806 conn.execute_batch(
1807 "CREATE TABLE IF NOT EXISTS min_peers (
1808 user_id INTEGER PRIMARY KEY,
1809 peer_id INTEGER NOT NULL,
1810 msg_id INTEGER NOT NULL
1811 );",
1812 )
1813 .map_err(Self::map_err)?;
1814 Ok(())
1815 }
1816
1817 fn read_session(conn: &rusqlite::Connection) -> io::Result<PersistedSession> {
1819 let home_dc_id: i32 = conn
1821 .query_row("SELECT value FROM meta WHERE key = 'home_dc_id'", [], |r| {
1822 r.get(0)
1823 })
1824 .unwrap_or(0);
1825
1826 let future_auth_token: Option<Vec<u8>> = conn
1828 .query_row(
1829 "SELECT value FROM meta WHERE key = 'future_auth_token'",
1830 [],
1831 |r| r.get::<_, Vec<u8>>(0),
1832 )
1833 .ok();
1834
1835 let mut stmt = conn
1837 .prepare("SELECT dc_id, flags, addr, auth_key, first_salt, time_offset FROM dcs")
1838 .map_err(Self::map_err)?;
1839 let dcs = stmt
1840 .query_map([], |row| {
1841 let dc_id: i32 = row.get(0)?;
1842 let flags_raw: u8 = row.get(1)?;
1843 let addr: String = row.get(2)?;
1844 let key_blob: Option<Vec<u8>> = row.get(3)?;
1845 let first_salt: i64 = row.get(4)?;
1846 let time_offset: i32 = row.get(5)?;
1847 Ok((dc_id, addr, key_blob, first_salt, time_offset, flags_raw))
1848 })
1849 .map_err(Self::map_err)?
1850 .filter_map(|r| r.ok())
1851 .map(
1852 |(dc_id, addr, key_blob, first_salt, time_offset, flags_raw)| {
1853 let auth_key = key_blob.and_then(|b| {
1854 if b.len() == 256 {
1855 let mut k = [0u8; 256];
1856 k.copy_from_slice(&b);
1857 Some(k)
1858 } else {
1859 None
1860 }
1861 });
1862 DcEntry {
1863 dc_id,
1864 addr,
1865 auth_key,
1866 first_salt,
1867 time_offset,
1868 flags: DcFlags(flags_raw),
1869 }
1870 },
1871 )
1872 .collect();
1873
1874 let updates_state = conn
1876 .query_row(
1877 "SELECT pts, qts, date, seq FROM update_state WHERE id = 1",
1878 [],
1879 |r| {
1880 Ok(UpdatesStateSnap {
1881 pts: r.get(0)?,
1882 qts: r.get(1)?,
1883 date: r.get(2)?,
1884 seq: r.get(3)?,
1885 channels: vec![],
1886 })
1887 },
1888 )
1889 .unwrap_or_default();
1890
1891 let mut ch_stmt = conn
1893 .prepare("SELECT channel_id, pts FROM channel_pts")
1894 .map_err(Self::map_err)?;
1895 let channels: Vec<(i64, i32)> = ch_stmt
1896 .query_map([], |r| Ok((r.get::<_, i64>(0)?, r.get::<_, i32>(1)?)))
1897 .map_err(Self::map_err)?
1898 .filter_map(|r| r.ok())
1899 .collect();
1900
1901 let mut peer_stmt = conn
1903 .prepare(
1904 "SELECT id, access_hash, is_channel, is_chat, channel_kind, is_community FROM peers",
1905 )
1906 .map_err(Self::map_err)?;
1907 let peers: Vec<CachedPeer> = peer_stmt
1908 .query_map([], |r| {
1909 let kind_raw: Option<i32> = r.get(4)?;
1910 Ok(CachedPeer {
1911 id: r.get(0)?,
1912 access_hash: r.get(1)?,
1913 is_channel: r.get::<_, i32>(2)? != 0,
1914 is_chat: r.get::<_, i32>(3)? != 0,
1915 channel_kind: kind_raw.and_then(|k| ChannelKind::from_byte(k as u8)),
1916 is_community: r.get::<_, i32>(5)? != 0,
1917 })
1918 })
1919 .map_err(Self::map_err)?
1920 .filter_map(|r| r.ok())
1921 .collect();
1922
1923 let mut min_stmt = conn
1925 .prepare("SELECT user_id, peer_id, msg_id FROM min_peers")
1926 .map_err(Self::map_err)?;
1927 let min_peers: Vec<CachedMinPeer> = min_stmt
1928 .query_map([], |r| {
1929 Ok(CachedMinPeer {
1930 user_id: r.get(0)?,
1931 peer_id: r.get(1)?,
1932 msg_id: r.get(2)?,
1933 })
1934 })
1935 .map_err(Self::map_err)?
1936 .filter_map(|r| r.ok())
1937 .collect();
1938
1939 Ok(PersistedSession {
1940 home_dc_id,
1941 dcs,
1942 updates_state: UpdatesStateSnap {
1943 channels,
1944 ..updates_state
1945 },
1946 peers,
1947 min_peers,
1948 future_auth_token,
1949 })
1950 }
1951
1952 fn write_session(conn: &rusqlite::Connection, s: &PersistedSession) -> io::Result<()> {
1954 conn.execute_batch("BEGIN IMMEDIATE")
1955 .map_err(Self::map_err)?;
1956
1957 conn.execute(
1958 "INSERT INTO meta (key, value) VALUES ('home_dc_id', ?1)
1959 ON CONFLICT(key) DO UPDATE SET value = excluded.value",
1960 rusqlite::params![s.home_dc_id],
1961 )
1962 .map_err(Self::map_err)?;
1963
1964 match &s.future_auth_token {
1965 Some(token) => conn.execute(
1966 "INSERT INTO meta (key, value) VALUES ('future_auth_token', ?1)
1967 ON CONFLICT(key) DO UPDATE SET value = excluded.value",
1968 rusqlite::params![token],
1969 ),
1970 None => conn.execute("DELETE FROM meta WHERE key = 'future_auth_token'", []),
1971 }
1972 .map_err(Self::map_err)?;
1973
1974 conn.execute("DELETE FROM dcs", []).map_err(Self::map_err)?;
1976 for d in &s.dcs {
1977 conn.execute(
1978 "INSERT INTO dcs (dc_id, flags, addr, auth_key, first_salt, time_offset)
1979 VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
1980 rusqlite::params![
1981 d.dc_id,
1982 d.flags.0,
1983 d.addr,
1984 d.auth_key.as_ref().map(|k| k.as_ref()),
1985 d.first_salt,
1986 d.time_offset,
1987 ],
1988 )
1989 .map_err(Self::map_err)?;
1990 }
1991
1992 let us = &s.updates_state;
1996 conn.execute(
1997 "INSERT INTO update_state (id, pts, qts, date, seq) VALUES (1, ?1, ?2, ?3, ?4)
1998 ON CONFLICT(id) DO UPDATE SET
1999 pts = MAX(excluded.pts, update_state.pts),
2000 qts = MAX(excluded.qts, update_state.qts),
2001 date = excluded.date,
2002 seq = excluded.seq",
2003 rusqlite::params![us.pts, us.qts, us.date, us.seq],
2004 )
2005 .map_err(Self::map_err)?;
2006
2007 conn.execute("DELETE FROM channel_pts", [])
2008 .map_err(Self::map_err)?;
2009 for &(cid, cpts) in &us.channels {
2010 conn.execute(
2011 "INSERT INTO channel_pts (channel_id, pts) VALUES (?1, ?2)",
2012 rusqlite::params![cid, cpts],
2013 )
2014 .map_err(Self::map_err)?;
2015 }
2016
2017 conn.execute("DELETE FROM peers", [])
2019 .map_err(Self::map_err)?;
2020 for p in &s.peers {
2021 conn.execute(
2022 "INSERT INTO peers (id, access_hash, is_channel, is_chat, channel_kind, is_community) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
2023 rusqlite::params![
2024 p.id,
2025 p.access_hash,
2026 p.is_channel as i32,
2027 p.is_chat as i32,
2028 p.channel_kind.map(|k| k.to_byte() as i32),
2029 p.is_community as i32,
2030 ],
2031 )
2032 .map_err(Self::map_err)?;
2033 }
2034
2035 conn.execute("DELETE FROM min_peers", [])
2037 .map_err(Self::map_err)?;
2038 for m in &s.min_peers {
2039 conn.execute(
2040 "INSERT INTO min_peers (user_id, peer_id, msg_id) VALUES (?1, ?2, ?3)",
2041 rusqlite::params![m.user_id, m.peer_id, m.msg_id],
2042 )
2043 .map_err(Self::map_err)?;
2044 }
2045
2046 conn.execute_batch("COMMIT").map_err(Self::map_err)
2047 }
2048}
2049
2050#[cfg(feature = "sqlite-session")]
2051impl SessionBackend for SqliteBackend {
2052 fn save(&self, session: &PersistedSession) -> io::Result<()> {
2053 let conn = self.conn.lock().unwrap();
2054 Self::write_session(&conn, session)
2055 }
2056
2057 fn load(&self) -> io::Result<Option<PersistedSession>> {
2058 let conn = self.conn.lock().unwrap();
2059 let count: i64 = conn
2061 .query_row("SELECT COUNT(*) FROM meta", [], |r| r.get(0))
2062 .map_err(Self::map_err)?;
2063 if count == 0 {
2064 return Ok(None);
2065 }
2066 Self::read_session(&conn).map(Some)
2067 }
2068
2069 fn delete(&self) -> io::Result<()> {
2070 let conn = self.conn.lock().unwrap();
2071 conn.execute_batch(
2072 "BEGIN IMMEDIATE;
2073 DELETE FROM meta;
2074 DELETE FROM dcs;
2075 DELETE FROM update_state;
2076 DELETE FROM channel_pts;
2077 DELETE FROM peers;
2078 DELETE FROM min_peers;
2079 COMMIT;",
2080 )
2081 .map_err(Self::map_err)
2082 }
2083
2084 fn name(&self) -> &str {
2085 &self.label
2086 }
2087
2088 fn update_dc(&self, entry: &DcEntry) -> io::Result<()> {
2091 let conn = self.conn.lock().unwrap();
2092 conn.execute(
2093 "INSERT INTO dcs (dc_id, flags, addr, auth_key, first_salt, time_offset)
2094 VALUES (?1, ?6, ?2, ?3, ?4, ?5)
2095 ON CONFLICT(dc_id, flags) DO UPDATE SET
2096 addr = excluded.addr,
2097 auth_key = excluded.auth_key,
2098 first_salt = excluded.first_salt,
2099 time_offset = excluded.time_offset",
2100 rusqlite::params![
2101 entry.dc_id,
2102 entry.addr,
2103 entry.auth_key.as_ref().map(|k| k.as_ref()),
2104 entry.first_salt,
2105 entry.time_offset,
2106 entry.flags.0,
2107 ],
2108 )
2109 .map(|_| ())
2110 .map_err(Self::map_err)
2111 }
2112
2113 fn set_home_dc(&self, dc_id: i32) -> io::Result<()> {
2114 let conn = self.conn.lock().unwrap();
2115 conn.execute(
2116 "INSERT INTO meta (key, value) VALUES ('home_dc_id', ?1)
2117 ON CONFLICT(key) DO UPDATE SET value = excluded.value",
2118 rusqlite::params![dc_id],
2119 )
2120 .map(|_| ())
2121 .map_err(Self::map_err)
2122 }
2123
2124 fn apply_update_state(&self, update: UpdateStateChange) -> io::Result<()> {
2125 let conn = self.conn.lock().unwrap();
2126 match update {
2127 UpdateStateChange::All(snap) => {
2128 conn.execute(
2129 "INSERT INTO update_state (id, pts, qts, date, seq) VALUES (1,?1,?2,?3,?4)
2130 ON CONFLICT(id) DO UPDATE SET
2131 pts=excluded.pts, qts=excluded.qts,
2132 date=excluded.date, seq=excluded.seq",
2133 rusqlite::params![snap.pts, snap.qts, snap.date, snap.seq],
2134 )
2135 .map_err(Self::map_err)?;
2136 conn.execute("DELETE FROM channel_pts", [])
2137 .map_err(Self::map_err)?;
2138 for &(cid, cpts) in &snap.channels {
2139 conn.execute(
2140 "INSERT INTO channel_pts (channel_id, pts) VALUES (?1, ?2)",
2141 rusqlite::params![cid, cpts],
2142 )
2143 .map_err(Self::map_err)?;
2144 }
2145 Ok(())
2146 }
2147 UpdateStateChange::Primary { pts, date, seq } => conn
2148 .execute(
2149 "INSERT INTO update_state (id, pts, qts, date, seq) VALUES (1,?1,0,?2,?3)
2150 ON CONFLICT(id) DO UPDATE SET pts=excluded.pts, date=excluded.date,
2151 seq=excluded.seq",
2152 rusqlite::params![pts, date, seq],
2153 )
2154 .map(|_| ())
2155 .map_err(Self::map_err),
2156 UpdateStateChange::Secondary { qts } => conn
2157 .execute(
2158 "INSERT INTO update_state (id, pts, qts, date, seq) VALUES (1,0,?1,0,0)
2159 ON CONFLICT(id) DO UPDATE SET qts = excluded.qts",
2160 rusqlite::params![qts],
2161 )
2162 .map(|_| ())
2163 .map_err(Self::map_err),
2164 UpdateStateChange::Channel { id, pts } => conn
2165 .execute(
2166 "INSERT INTO channel_pts (channel_id, pts) VALUES (?1, ?2)
2167 ON CONFLICT(channel_id) DO UPDATE SET pts = excluded.pts",
2168 rusqlite::params![id, pts],
2169 )
2170 .map(|_| ())
2171 .map_err(Self::map_err),
2172 }
2173 }
2174
2175 fn cache_peer(&self, peer: &CachedPeer) -> io::Result<()> {
2176 let conn = self.conn.lock().unwrap();
2177 conn.execute(
2178 "INSERT INTO peers (id, access_hash, is_channel, is_chat, channel_kind, is_community) VALUES (?1, ?2, ?3, ?4, ?5, ?6)
2179 ON CONFLICT(id) DO UPDATE SET
2180 access_hash = excluded.access_hash,
2181 is_channel = excluded.is_channel,
2182 is_chat = excluded.is_chat,
2183 channel_kind = excluded.channel_kind,
2184 is_community = excluded.is_community",
2185 rusqlite::params![
2186 peer.id,
2187 peer.access_hash,
2188 peer.is_channel as i32,
2189 peer.is_chat as i32,
2190 peer.channel_kind.map(|k| k.to_byte() as i32),
2191 peer.is_community as i32,
2192 ],
2193 )
2194 .map(|_| ())
2195 .map_err(Self::map_err)
2196 }
2197}
2198
2199#[cfg(all(test, feature = "sqlite-session"))]
2200mod sqlite_backend_tests {
2201 use super::*;
2202
2203 #[test]
2204 fn future_auth_token_roundtrip() {
2205 let b = SqliteBackend::in_memory().unwrap();
2206 let mut s = PersistedSession::default();
2207 s.home_dc_id = 2;
2208 s.future_auth_token = Some(vec![1, 2, 3, 4, 5]);
2209 b.save(&s).unwrap();
2210
2211 let loaded = b.load().unwrap().unwrap();
2212 assert_eq!(loaded.future_auth_token, Some(vec![1, 2, 3, 4, 5]));
2213 }
2214
2215 #[test]
2216 fn future_auth_token_cleared_on_none() {
2217 let b = SqliteBackend::in_memory().unwrap();
2218 let mut s = PersistedSession::default();
2219 s.home_dc_id = 2;
2220 s.future_auth_token = Some(vec![9, 9, 9]);
2221 b.save(&s).unwrap();
2222 assert_eq!(
2223 b.load().unwrap().unwrap().future_auth_token,
2224 Some(vec![9, 9, 9])
2225 );
2226
2227 s.future_auth_token = None;
2228 b.save(&s).unwrap();
2229 assert_eq!(b.load().unwrap().unwrap().future_auth_token, None);
2230 }
2231
2232 #[test]
2233 fn is_community_roundtrip() {
2234 let b = SqliteBackend::in_memory().unwrap();
2235 let mut s = PersistedSession::default();
2236 s.home_dc_id = 2;
2237 s.peers.push(CachedPeer {
2238 id: 42,
2239 access_hash: 0x1111,
2240 is_channel: false,
2241 is_chat: false,
2242 channel_kind: None,
2243 is_community: true,
2244 });
2245 b.save(&s).unwrap();
2246
2247 let loaded = b.load().unwrap().unwrap();
2248 assert_eq!(loaded.peers.len(), 1);
2249 assert!(loaded.peers[0].is_community);
2250 }
2251
2252 #[test]
2253 fn is_community_migrates_from_legacy_schema() {
2254 let conn = rusqlite::Connection::open_in_memory().unwrap();
2257 conn.execute_batch(
2258 "CREATE TABLE meta (key TEXT PRIMARY KEY, value INTEGER NOT NULL DEFAULT 0);
2259 CREATE TABLE dcs (
2260 dc_id INTEGER NOT NULL, flags INTEGER NOT NULL DEFAULT 0,
2261 addr TEXT NOT NULL, auth_key BLOB,
2262 first_salt INTEGER NOT NULL DEFAULT 0, time_offset INTEGER NOT NULL DEFAULT 0,
2263 PRIMARY KEY (dc_id, flags)
2264 );
2265 CREATE TABLE update_state (
2266 id INTEGER PRIMARY KEY CHECK (id = 1),
2267 pts INTEGER NOT NULL DEFAULT 0, qts INTEGER NOT NULL DEFAULT 0,
2268 date INTEGER NOT NULL DEFAULT 0, seq INTEGER NOT NULL DEFAULT 0
2269 );
2270 CREATE TABLE channel_pts (channel_id INTEGER PRIMARY KEY, pts INTEGER NOT NULL);
2271 CREATE TABLE peers (
2272 id INTEGER PRIMARY KEY, access_hash INTEGER NOT NULL,
2273 is_channel INTEGER NOT NULL DEFAULT 0, is_chat INTEGER NOT NULL DEFAULT 0
2274 );
2275 INSERT INTO meta (key, value) VALUES ('home_dc_id', 5);
2276 INSERT INTO peers (id, access_hash, is_channel, is_chat) VALUES (9, 123, 0, 0);",
2277 )
2278 .unwrap();
2279
2280 SqliteBackend::migrate_legacy_sqlite_schema(&conn).unwrap();
2281
2282 let loaded = SqliteBackend::read_session(&conn).unwrap();
2283 assert_eq!(loaded.peers.len(), 1);
2284 assert!(!loaded.peers[0].is_community);
2285 }
2286}
2287
2288#[cfg(feature = "libsql-session")]
2308pub struct LibSqlBackend {
2309 conn: libsql::Connection,
2310 label: String,
2311}
2312
2313#[cfg(feature = "libsql-session")]
2314impl LibSqlBackend {
2315 const SCHEMA: &'static str = "
2316 CREATE TABLE IF NOT EXISTS meta (
2317 key TEXT PRIMARY KEY,
2318 value INTEGER NOT NULL DEFAULT 0
2319 );
2320 CREATE TABLE IF NOT EXISTS dcs (
2321 dc_id INTEGER NOT NULL,
2322 flags INTEGER NOT NULL DEFAULT 0,
2323 addr TEXT NOT NULL,
2324 auth_key BLOB,
2325 first_salt INTEGER NOT NULL DEFAULT 0,
2326 time_offset INTEGER NOT NULL DEFAULT 0,
2327 PRIMARY KEY (dc_id, flags)
2328 );
2329 CREATE TABLE IF NOT EXISTS update_state (
2330 id INTEGER PRIMARY KEY CHECK (id = 1),
2331 pts INTEGER NOT NULL DEFAULT 0,
2332 qts INTEGER NOT NULL DEFAULT 0,
2333 date INTEGER NOT NULL DEFAULT 0,
2334 seq INTEGER NOT NULL DEFAULT 0
2335 );
2336 CREATE TABLE IF NOT EXISTS channel_pts (
2337 channel_id INTEGER PRIMARY KEY,
2338 pts INTEGER NOT NULL
2339 );
2340 CREATE TABLE IF NOT EXISTS peers (
2341 id INTEGER PRIMARY KEY,
2342 access_hash INTEGER NOT NULL,
2343 is_channel INTEGER NOT NULL DEFAULT 0,
2344 is_chat INTEGER NOT NULL DEFAULT 0,
2345 channel_kind INTEGER,
2346 is_community INTEGER NOT NULL DEFAULT 0
2347 );
2348 CREATE TABLE IF NOT EXISTS min_peers (
2349 user_id INTEGER PRIMARY KEY,
2350 peer_id INTEGER NOT NULL,
2351 msg_id INTEGER NOT NULL
2352 );
2353 ";
2354
2355 fn block<F, T>(fut: F) -> io::Result<T>
2356 where
2357 F: std::future::Future<Output = Result<T, libsql::Error>>,
2358 {
2359 tokio::runtime::Handle::current()
2360 .block_on(fut)
2361 .map_err(io::Error::other)
2362 }
2363
2364 async fn apply_schema(conn: &libsql::Connection) -> Result<(), libsql::Error> {
2365 conn.execute_batch(Self::SCHEMA).await?;
2366 let _ = conn
2370 .execute_batch("ALTER TABLE peers ADD COLUMN channel_kind INTEGER;")
2371 .await;
2372 let _ = conn
2374 .execute_batch("ALTER TABLE peers ADD COLUMN is_community INTEGER NOT NULL DEFAULT 0;")
2375 .await;
2376 Ok(())
2377 }
2378
2379 pub fn open_local(path: impl Into<PathBuf>) -> io::Result<Self> {
2381 let path = path.into();
2382 let label = path.display().to_string();
2383 let db = Self::block(async { libsql::Builder::new_local(path).build().await })?;
2384 let conn = Self::block(async { db.connect() }).map_err(io::Error::other)?;
2385 Self::block(Self::apply_schema(&conn))?;
2386 Ok(Self { conn, label })
2387 }
2388
2389 pub fn in_memory() -> io::Result<Self> {
2391 let db = Self::block(async { libsql::Builder::new_local(":memory:").build().await })?;
2392 let conn = Self::block(async { db.connect() }).map_err(io::Error::other)?;
2393 Self::block(Self::apply_schema(&conn))?;
2394 Ok(Self {
2395 conn,
2396 label: ":memory:".into(),
2397 })
2398 }
2399
2400 #[cfg(feature = "libsql-remote-session")]
2402 pub fn open_remote(url: impl Into<String>, auth_token: impl Into<String>) -> io::Result<Self> {
2403 let url = url.into();
2404 let label = url.clone();
2405 let db = Self::block(async {
2406 libsql::Builder::new_remote(url, auth_token.into())
2407 .build()
2408 .await
2409 })?;
2410 let conn = Self::block(async { db.connect() }).map_err(io::Error::other)?;
2411 Self::block(Self::apply_schema(&conn))?;
2412 Ok(Self { conn, label })
2413 }
2414
2415 #[cfg(feature = "libsql-remote-session")]
2417 pub fn open_replica(
2418 path: impl Into<PathBuf>,
2419 url: impl Into<String>,
2420 auth_token: impl Into<String>,
2421 ) -> io::Result<Self> {
2422 let path = path.into();
2423 let url = url.into();
2424 let auth_token = auth_token.into();
2425 let label = format!("{} (replica of {})", path.display(), url);
2426 let db = Self::block(async {
2427 libsql::Builder::new_remote_replica(path, url, auth_token)
2428 .build()
2429 .await
2430 })?;
2431 let conn = Self::block(async { db.connect() }).map_err(io::Error::other)?;
2432 Self::block(Self::apply_schema(&conn))?;
2433 Ok(Self { conn, label })
2434 }
2435
2436 async fn read_session_async(
2437 conn: &libsql::Connection,
2438 ) -> Result<PersistedSession, libsql::Error> {
2439 let home_dc_id: i32 = conn
2441 .query("SELECT value FROM meta WHERE key = 'home_dc_id'", ())
2442 .await?
2443 .next()
2444 .await?
2445 .map(|r| r.get::<i32>(0))
2446 .transpose()?
2447 .unwrap_or(0);
2448
2449 let future_auth_token: Option<Vec<u8>> = conn
2451 .query("SELECT value FROM meta WHERE key = 'future_auth_token'", ())
2452 .await?
2453 .next()
2454 .await?
2455 .map(|r| r.get::<Vec<u8>>(0))
2456 .transpose()?;
2457
2458 let mut rows = conn
2460 .query(
2461 "SELECT dc_id, flags, addr, auth_key, first_salt, time_offset FROM dcs",
2462 (),
2463 )
2464 .await?;
2465 let mut dcs = Vec::new();
2466 while let Some(row) = rows.next().await? {
2467 let dc_id: i32 = row.get(0)?;
2468 let flags_raw: u8 = row.get::<i64>(1)? as u8;
2469 let addr: String = row.get(2)?;
2470 let key_blob: Option<Vec<u8>> = row.get(3)?;
2471 let first_salt: i64 = row.get(4)?;
2472 let time_offset: i32 = row.get(5)?;
2473 let auth_key = match key_blob {
2474 Some(b) if b.len() == 256 => {
2475 let mut k = [0u8; 256];
2476 k.copy_from_slice(&b);
2477 Some(k)
2478 }
2479 Some(b) => {
2480 return Err(libsql::Error::Misuse(format!(
2481 "auth_key blob must be 256 bytes, got {}",
2482 b.len()
2483 )));
2484 }
2485 None => None,
2486 };
2487 dcs.push(DcEntry {
2488 dc_id,
2489 addr,
2490 auth_key,
2491 first_salt,
2492 time_offset,
2493 flags: DcFlags(flags_raw),
2494 });
2495 }
2496
2497 let mut us_row = conn
2499 .query(
2500 "SELECT pts, qts, date, seq FROM update_state WHERE id = 1",
2501 (),
2502 )
2503 .await?;
2504 let updates_state = if let Some(r) = us_row.next().await? {
2505 UpdatesStateSnap {
2506 pts: r.get(0)?,
2507 qts: r.get(1)?,
2508 date: r.get(2)?,
2509 seq: r.get(3)?,
2510 channels: vec![],
2511 }
2512 } else {
2513 UpdatesStateSnap::default()
2514 };
2515
2516 let mut ch_rows = conn
2518 .query("SELECT channel_id, pts FROM channel_pts", ())
2519 .await?;
2520 let mut channels = Vec::new();
2521 while let Some(r) = ch_rows.next().await? {
2522 channels.push((r.get::<i64>(0)?, r.get::<i32>(1)?));
2523 }
2524
2525 let mut peer_rows = conn
2527 .query(
2528 "SELECT id, access_hash, is_channel, is_chat, channel_kind, is_community FROM peers",
2529 (),
2530 )
2531 .await?;
2532 let mut peers = Vec::new();
2533 while let Some(r) = peer_rows.next().await? {
2534 let kind_raw: Option<i32> = r.get(4).ok();
2535 let is_community: i32 = r.get(5).unwrap_or(0);
2536 peers.push(CachedPeer {
2537 id: r.get(0)?,
2538 access_hash: r.get(1)?,
2539 is_channel: r.get::<i32>(2)? != 0,
2540 is_chat: r.get::<i32>(3)? != 0,
2541 channel_kind: kind_raw.and_then(|k| ChannelKind::from_byte(k as u8)),
2542 is_community: is_community != 0,
2543 });
2544 }
2545
2546 let mut min_rows = conn
2548 .query("SELECT user_id, peer_id, msg_id FROM min_peers", ())
2549 .await?;
2550 let mut min_peers = Vec::new();
2551 while let Some(r) = min_rows.next().await? {
2552 min_peers.push(CachedMinPeer {
2553 user_id: r.get(0)?,
2554 peer_id: r.get(1)?,
2555 msg_id: r.get(2)?,
2556 });
2557 }
2558
2559 Ok(PersistedSession {
2560 home_dc_id,
2561 dcs,
2562 updates_state: UpdatesStateSnap {
2563 channels,
2564 ..updates_state
2565 },
2566 peers,
2567 min_peers,
2568 future_auth_token,
2569 })
2570 }
2571
2572 async fn write_session_async(
2573 conn: &libsql::Connection,
2574 s: &PersistedSession,
2575 ) -> Result<(), libsql::Error> {
2576 conn.execute_batch("BEGIN IMMEDIATE").await.map(|_| ())?;
2577
2578 conn.execute(
2579 "INSERT INTO meta (key, value) VALUES ('home_dc_id', ?1)
2580 ON CONFLICT(key) DO UPDATE SET value = excluded.value",
2581 libsql::params![s.home_dc_id],
2582 )
2583 .await?;
2584
2585 match &s.future_auth_token {
2586 Some(token) => {
2587 conn.execute(
2588 "INSERT INTO meta (key, value) VALUES ('future_auth_token', ?1)
2589 ON CONFLICT(key) DO UPDATE SET value = excluded.value",
2590 libsql::params![token.clone()],
2591 )
2592 .await?
2593 }
2594 None => {
2595 conn.execute("DELETE FROM meta WHERE key = 'future_auth_token'", ())
2596 .await?
2597 }
2598 };
2599
2600 conn.execute("DELETE FROM dcs", ()).await?;
2601 for d in &s.dcs {
2602 conn.execute(
2603 "INSERT INTO dcs (dc_id, flags, addr, auth_key, first_salt, time_offset)
2604 VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
2605 libsql::params![
2606 d.dc_id,
2607 d.flags.0 as i64,
2608 d.addr.clone(),
2609 d.auth_key.map(|k| k.to_vec()),
2610 d.first_salt,
2611 d.time_offset,
2612 ],
2613 )
2614 .await?;
2615 }
2616
2617 let us = &s.updates_state;
2618 conn.execute(
2619 "INSERT INTO update_state (id, pts, qts, date, seq) VALUES (1,?1,?2,?3,?4)
2620 ON CONFLICT(id) DO UPDATE SET
2621 pts = MAX(excluded.pts, update_state.pts),
2622 qts = MAX(excluded.qts, update_state.qts),
2623 date = excluded.date,
2624 seq = excluded.seq",
2625 libsql::params![us.pts, us.qts, us.date, us.seq],
2626 )
2627 .await?;
2628
2629 conn.execute("DELETE FROM channel_pts", ()).await?;
2630 for &(cid, cpts) in &us.channels {
2631 conn.execute(
2632 "INSERT INTO channel_pts (channel_id, pts) VALUES (?1,?2)",
2633 libsql::params![cid, cpts],
2634 )
2635 .await?;
2636 }
2637
2638 conn.execute("DELETE FROM peers", ()).await?;
2639 for p in &s.peers {
2640 conn.execute(
2641 "INSERT INTO peers (id, access_hash, is_channel, is_chat, channel_kind, is_community) VALUES (?1,?2,?3,?4,?5,?6)",
2642 libsql::params![
2643 p.id,
2644 p.access_hash,
2645 p.is_channel as i32,
2646 p.is_chat as i32,
2647 p.channel_kind.map(|k| k.to_byte() as i32),
2648 p.is_community as i32,
2649 ],
2650 )
2651 .await?;
2652 }
2653
2654 conn.execute("DELETE FROM min_peers", ()).await?;
2655 for m in &s.min_peers {
2656 conn.execute(
2657 "INSERT INTO min_peers (user_id, peer_id, msg_id) VALUES (?1,?2,?3)",
2658 libsql::params![m.user_id, m.peer_id, m.msg_id],
2659 )
2660 .await?;
2661 }
2662
2663 conn.execute_batch("COMMIT").await.map(|_| ())
2664 }
2665}
2666
2667#[cfg(feature = "libsql-session")]
2668impl SessionBackend for LibSqlBackend {
2669 fn save(&self, session: &PersistedSession) -> io::Result<()> {
2670 let conn = self.conn.clone();
2671 let session = session.clone();
2672 Self::block(async move { Self::write_session_async(&conn, &session).await })
2673 }
2674
2675 fn load(&self) -> io::Result<Option<PersistedSession>> {
2676 let conn = self.conn.clone();
2677 let count: i64 = Self::block(async move {
2678 let mut rows = conn.query("SELECT COUNT(*) FROM meta", ()).await?;
2679 Ok::<i64, libsql::Error>(rows.next().await?.and_then(|r| r.get(0).ok()).unwrap_or(0))
2680 })?;
2681 if count == 0 {
2682 return Ok(None);
2683 }
2684 let conn = self.conn.clone();
2685 Self::block(async move { Self::read_session_async(&conn).await }).map(Some)
2686 }
2687
2688 fn delete(&self) -> io::Result<()> {
2689 let conn = self.conn.clone();
2690 Self::block(async move {
2691 conn.execute_batch(
2692 "BEGIN IMMEDIATE;
2693 DELETE FROM meta;
2694 DELETE FROM dcs;
2695 DELETE FROM update_state;
2696 DELETE FROM channel_pts;
2697 DELETE FROM peers;
2698 DELETE FROM min_peers;
2699 COMMIT;",
2700 )
2701 .await
2702 .map(|_| ())
2703 })
2704 }
2705
2706 fn name(&self) -> &str {
2707 &self.label
2708 }
2709
2710 fn update_dc(&self, entry: &DcEntry) -> io::Result<()> {
2713 let conn = self.conn.clone();
2714 let (dc_id, addr, key, salt, off, flags) = (
2715 entry.dc_id,
2716 entry.addr.clone(),
2717 entry.auth_key.map(|k| k.to_vec()),
2718 entry.first_salt,
2719 entry.time_offset,
2720 entry.flags.0 as i64,
2721 );
2722 Self::block(async move {
2723 conn.execute(
2724 "INSERT INTO dcs (dc_id, flags, addr, auth_key, first_salt, time_offset)
2725 VALUES (?1,?6,?2,?3,?4,?5)
2726 ON CONFLICT(dc_id, flags) DO UPDATE SET
2727 addr=excluded.addr, auth_key=excluded.auth_key,
2728 first_salt=excluded.first_salt, time_offset=excluded.time_offset",
2729 libsql::params![dc_id, addr, key, salt, off, flags],
2730 )
2731 .await
2732 .map(|_| ())
2733 })
2734 }
2735
2736 fn set_home_dc(&self, dc_id: i32) -> io::Result<()> {
2737 let conn = self.conn.clone();
2738 Self::block(async move {
2739 conn.execute(
2740 "INSERT INTO meta (key, value) VALUES ('home_dc_id',?1)
2741 ON CONFLICT(key) DO UPDATE SET value=excluded.value",
2742 libsql::params![dc_id],
2743 )
2744 .await
2745 .map(|_| ())
2746 })
2747 }
2748
2749 fn apply_update_state(&self, update: UpdateStateChange) -> io::Result<()> {
2750 let conn = self.conn.clone();
2751 Self::block(async move {
2752 match update {
2753 UpdateStateChange::All(snap) => {
2754 conn.execute(
2755 "INSERT INTO update_state (id,pts,qts,date,seq) VALUES (1,?1,?2,?3,?4)
2756 ON CONFLICT(id) DO UPDATE SET pts=excluded.pts,qts=excluded.qts,
2757 date=excluded.date,seq=excluded.seq",
2758 libsql::params![snap.pts, snap.qts, snap.date, snap.seq],
2759 )
2760 .await?;
2761 conn.execute("DELETE FROM channel_pts", ()).await?;
2762 for &(cid, cpts) in &snap.channels {
2763 conn.execute(
2764 "INSERT INTO channel_pts (channel_id,pts) VALUES (?1,?2)",
2765 libsql::params![cid, cpts],
2766 )
2767 .await?;
2768 }
2769 Ok(())
2770 }
2771 UpdateStateChange::Primary { pts, date, seq } => conn
2772 .execute(
2773 "INSERT INTO update_state (id,pts,qts,date,seq) VALUES (1,?1,0,?2,?3)
2774 ON CONFLICT(id) DO UPDATE SET pts=excluded.pts,date=excluded.date,
2775 seq=excluded.seq",
2776 libsql::params![pts, date, seq],
2777 )
2778 .await
2779 .map(|_| ()),
2780 UpdateStateChange::Secondary { qts } => conn
2781 .execute(
2782 "INSERT INTO update_state (id,pts,qts,date,seq) VALUES (1,0,?1,0,0)
2783 ON CONFLICT(id) DO UPDATE SET qts=excluded.qts",
2784 libsql::params![qts],
2785 )
2786 .await
2787 .map(|_| ()),
2788 UpdateStateChange::Channel { id, pts } => conn
2789 .execute(
2790 "INSERT INTO channel_pts (channel_id,pts) VALUES (?1,?2)
2791 ON CONFLICT(channel_id) DO UPDATE SET pts=excluded.pts",
2792 libsql::params![id, pts],
2793 )
2794 .await
2795 .map(|_| ()),
2796 }
2797 })
2798 }
2799
2800 fn cache_peer(&self, peer: &CachedPeer) -> io::Result<()> {
2801 let conn = self.conn.clone();
2802 let (id, hash, is_ch, is_ct, kind, is_comm) = (
2803 peer.id,
2804 peer.access_hash,
2805 peer.is_channel as i32,
2806 peer.is_chat as i32,
2807 peer.channel_kind.map(|k| k.to_byte() as i32),
2808 peer.is_community as i32,
2809 );
2810 Self::block(async move {
2811 conn.execute(
2812 "INSERT INTO peers (id,access_hash,is_channel,is_chat,channel_kind,is_community) VALUES (?1,?2,?3,?4,?5,?6)
2813 ON CONFLICT(id) DO UPDATE SET
2814 access_hash=excluded.access_hash,
2815 is_channel=excluded.is_channel,
2816 is_chat=excluded.is_chat,
2817 channel_kind=excluded.channel_kind,
2818 is_community=excluded.is_community",
2819 libsql::params![id, hash, is_ch, is_ct, kind, is_comm],
2820 )
2821 .await
2822 .map(|_| ())
2823 })
2824 }
2825}
2826
2827#[cfg(all(test, feature = "libsql-session"))]
2828mod libsql_backend_tests {
2829 use super::*;
2830
2831 async fn open_in_memory() -> libsql::Connection {
2832 let db = libsql::Builder::new_local(":memory:")
2833 .build()
2834 .await
2835 .unwrap();
2836 let conn = db.connect().unwrap();
2837 LibSqlBackend::apply_schema(&conn).await.unwrap();
2838 conn
2839 }
2840
2841 #[tokio::test]
2842 async fn future_auth_token_roundtrip() {
2843 let conn = open_in_memory().await;
2844 let mut s = PersistedSession::default();
2845 s.home_dc_id = 2;
2846 s.future_auth_token = Some(vec![1, 2, 3, 4, 5]);
2847 LibSqlBackend::write_session_async(&conn, &s).await.unwrap();
2848
2849 let loaded = LibSqlBackend::read_session_async(&conn).await.unwrap();
2850 assert_eq!(loaded.future_auth_token, Some(vec![1, 2, 3, 4, 5]));
2851 }
2852
2853 #[tokio::test]
2854 async fn future_auth_token_cleared_on_none() {
2855 let conn = open_in_memory().await;
2856 let mut s = PersistedSession::default();
2857 s.home_dc_id = 2;
2858 s.future_auth_token = Some(vec![9, 9, 9]);
2859 LibSqlBackend::write_session_async(&conn, &s).await.unwrap();
2860 let loaded = LibSqlBackend::read_session_async(&conn).await.unwrap();
2861 assert_eq!(loaded.future_auth_token, Some(vec![9, 9, 9]));
2862
2863 s.future_auth_token = None;
2864 LibSqlBackend::write_session_async(&conn, &s).await.unwrap();
2865 let loaded = LibSqlBackend::read_session_async(&conn).await.unwrap();
2866 assert_eq!(loaded.future_auth_token, None);
2867 }
2868
2869 #[tokio::test]
2870 async fn is_community_roundtrip() {
2871 let conn = open_in_memory().await;
2872 let mut s = PersistedSession::default();
2873 s.home_dc_id = 2;
2874 s.peers.push(CachedPeer {
2875 id: 42,
2876 access_hash: 0x1111,
2877 is_channel: false,
2878 is_chat: false,
2879 channel_kind: None,
2880 is_community: true,
2881 });
2882 LibSqlBackend::write_session_async(&conn, &s).await.unwrap();
2883
2884 let loaded = LibSqlBackend::read_session_async(&conn).await.unwrap();
2885 assert_eq!(loaded.peers.len(), 1);
2886 assert!(loaded.peers[0].is_community);
2887 }
2888}