layer_mtproto/
encrypted.rs

1//! Encrypted MTProto 2.0 session (post auth-key).
2//!
3//! Once you have a `Finished` from [`crate::authentication`], construct an
4//! [`EncryptedSession`] and use it to serialize/deserialize all subsequent
5//! messages.
6
7use std::time::{SystemTime, UNIX_EPOCH};
8
9use layer_crypto::{AuthKey, DequeBuffer, decrypt_data_v2, encrypt_data_v2};
10use layer_tl_types::RemoteCall;
11
12
13/// Errors that can occur when decrypting a server message.
14#[derive(Debug)]
15pub enum DecryptError {
16    /// The underlying crypto layer rejected the message.
17    Crypto(layer_crypto::DecryptError),
18    /// The decrypted inner message was too short to contain a valid header.
19    FrameTooShort,
20    /// Session-ID mismatch (possible replay or wrong connection).
21    SessionMismatch,
22}
23
24impl std::fmt::Display for DecryptError {
25    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        match self {
27            Self::Crypto(e) => write!(f, "crypto: {e}"),
28            Self::FrameTooShort => write!(f, "inner plaintext too short"),
29            Self::SessionMismatch => write!(f, "session_id mismatch"),
30        }
31    }
32}
33impl std::error::Error for DecryptError {}
34
35/// The inner payload extracted from a successfully decrypted server frame.
36pub struct DecryptedMessage {
37    /// `salt` sent by the server.
38    pub salt:       i64,
39    /// The `session_id` from the frame.
40    pub session_id: i64,
41    /// The `msg_id` of the inner message.
42    pub msg_id:     i64,
43    /// `seq_no` of the inner message.
44    pub seq_no:     i32,
45    /// TL-serialized body of the inner message.
46    pub body:       Vec<u8>,
47}
48
49/// MTProto 2.0 encrypted session state.
50///
51/// Wraps an `AuthKey` and tracks per-session counters (session_id, seq_no,
52/// last_msg_id, server salt).  Use [`EncryptedSession::pack`] to encrypt
53/// outgoing requests and [`EncryptedSession::unpack`] to decrypt incoming
54/// server frames.
55pub struct EncryptedSession {
56    auth_key:    AuthKey,
57    session_id:  i64,
58    sequence:    i32,
59    last_msg_id: i64,
60    /// Current server salt to include in outgoing messages.
61    pub salt:    i64,
62    /// Clock skew in seconds vs. server.
63    pub time_offset: i32,
64}
65
66impl EncryptedSession {
67    /// Create a new encrypted session from the output of `authentication::finish`.
68    pub fn new(auth_key: [u8; 256], first_salt: i64, time_offset: i32) -> Self {
69        let mut rnd = [0u8; 8];
70        getrandom::getrandom(&mut rnd).expect("getrandom");
71        Self {
72            auth_key: AuthKey::from_bytes(auth_key),
73            session_id: i64::from_le_bytes(rnd),
74            sequence: 0,
75            last_msg_id: 0,
76            salt: first_salt,
77            time_offset,
78        }
79    }
80
81    /// Compute the next message ID (based on corrected server time).
82    fn next_msg_id(&mut self) -> i64 {
83        let now = SystemTime::now()
84            .duration_since(UNIX_EPOCH).unwrap();
85        let secs = (now.as_secs() as i32).wrapping_add(self.time_offset) as u64;
86        let nanos = now.subsec_nanos() as u64;
87        let mut id = ((secs << 32) | (nanos << 2)) as i64;
88        if self.last_msg_id >= id { id = self.last_msg_id + 4; }
89        self.last_msg_id = id;
90        id
91    }
92
93    /// Next content-related seq_no (odd) and advance the counter.
94    fn next_seq_no(&mut self) -> i32 {
95        let n = self.sequence * 2 + 1;
96        self.sequence += 1;
97        n
98    }
99
100    /// Serialize and encrypt a TL function into a wire-ready byte vector.
101    ///
102    /// Layout of the plaintext before encryption:
103    /// ```text
104    /// salt:       i64
105    /// session_id: i64
106    /// msg_id:     i64
107    /// seq_no:     i32
108    /// body_len:   i32
109    /// body:       [u8; body_len]
110    /// ```
111    /// Like `pack` but only requires `Serializable` (not `RemoteCall`).
112    /// Useful for generic wrapper types like `InvokeWithLayer<InitConnection<X>>`
113    /// where the return type is determined by the inner call, not the wrapper.
114    pub fn pack_serializable<S: layer_tl_types::Serializable>(&mut self, call: &S) -> Vec<u8> {
115        let body = call.to_bytes();
116        let msg_id = self.next_msg_id();
117        let seq_no = self.next_seq_no();
118
119        let inner_len = 8 + 8 + 8 + 4 + 4 + body.len();
120        let mut buf = DequeBuffer::with_capacity(inner_len, 32);
121        buf.extend(self.salt.to_le_bytes());
122        buf.extend(self.session_id.to_le_bytes());
123        buf.extend(msg_id.to_le_bytes());
124        buf.extend(seq_no.to_le_bytes());
125        buf.extend((body.len() as u32).to_le_bytes());
126        buf.extend(body.iter().copied());
127
128        encrypt_data_v2(&mut buf, &self.auth_key);
129        buf.as_ref().to_vec()
130    }
131
132    /// Encrypt and frame a [`RemoteCall`] into a ready-to-send MTProto message.
133    ///
134    /// Returns the encrypted bytes to pass directly to the transport layer.
135    pub fn pack<R: RemoteCall>(&mut self, call: &R) -> Vec<u8> {
136        let body = call.to_bytes();
137        let msg_id = self.next_msg_id();
138        let seq_no = self.next_seq_no();
139
140        // Build plaintext inner payload
141        let inner_len = 8 + 8 + 8 + 4 + 4 + body.len();
142        // Front capacity = 32 for auth_key_id + msg_key
143        let mut buf = DequeBuffer::with_capacity(inner_len, 32);
144        buf.extend(self.salt.to_le_bytes());
145        buf.extend(self.session_id.to_le_bytes());
146        buf.extend(msg_id.to_le_bytes());
147        buf.extend(seq_no.to_le_bytes());
148        buf.extend((body.len() as u32).to_le_bytes());
149        buf.extend(body.iter().copied());
150
151        encrypt_data_v2(&mut buf, &self.auth_key);
152        buf.as_ref().to_vec()
153    }
154
155    /// Decrypt an encrypted server frame.
156    ///
157    /// `frame` should be a raw frame received from the transport (already
158    /// stripped of the abridged-length prefix).
159    pub fn unpack(&self, frame: &mut Vec<u8>) -> Result<DecryptedMessage, DecryptError> {
160        let plaintext = decrypt_data_v2(frame, &self.auth_key)
161            .map_err(DecryptError::Crypto)?;
162
163        // inner: salt(8) + session_id(8) + msg_id(8) + seq_no(4) + len(4) + body
164        if plaintext.len() < 32 {
165            return Err(DecryptError::FrameTooShort);
166        }
167
168        let salt       = i64::from_le_bytes(plaintext[..8].try_into().unwrap());
169        let session_id = i64::from_le_bytes(plaintext[8..16].try_into().unwrap());
170        let msg_id     = i64::from_le_bytes(plaintext[16..24].try_into().unwrap());
171        let seq_no     = i32::from_le_bytes(plaintext[24..28].try_into().unwrap());
172        let body_len   = u32::from_le_bytes(plaintext[28..32].try_into().unwrap()) as usize;
173
174        if session_id != self.session_id {
175            return Err(DecryptError::SessionMismatch);
176        }
177
178        let body = plaintext[32..32 + body_len.min(plaintext.len() - 32)].to_vec();
179
180        Ok(DecryptedMessage { salt, session_id, msg_id, seq_no, body })
181    }
182
183    /// Return the auth_key bytes (for persistence).
184    pub fn auth_key_bytes(&self) -> [u8; 256] { self.auth_key.to_bytes() }
185
186    /// Return the current session_id.
187    pub fn session_id(&self) -> i64 { self.session_id }
188}
layer_mtproto/encrypted.rs

layer_mtproto/
encrypted.rs
