Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WASM-1] Implement the obfuscated MTProto transport #298

Merged
merged 2 commits into from
Dec 21, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions lib/grammers-crypto/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ pbkdf2 = "0.12.2"
sha1 = "0.10.6"
sha2 = "0.10.8"
num-traits = "0.2.19"
ctr = "0.9.2"

[dev-dependencies]
bencher = "0.1.5"
Expand Down
4 changes: 4 additions & 0 deletions lib/grammers-crypto/DEPS.md
Original file line number Diff line number Diff line change
Expand Up @@ -48,3 +48,7 @@ Used for methods relied on by the 2-factor offered by Telegram.
## toml

Used to test that this file lists all dependencies from `Cargo.toml`.

## ctr

Used for the AES-CTR mode needed for the obfuscated MTProto transport.
1 change: 1 addition & 0 deletions lib/grammers-crypto/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ mod auth_key;
pub mod deque_buffer;
pub mod factorize;
pub mod hex;
pub mod obfuscated;
pub mod rsa;
pub mod sha;
pub mod two_factor_auth;
Expand Down
43 changes: 43 additions & 0 deletions lib/grammers-crypto/src/obfuscated.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
// Copyright 2020 - developers of the `grammers` project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use aes::cipher::{generic_array::GenericArray, KeyIvInit, StreamCipher};

/// This implements the AES-256-CTR cipher used by Telegram to encrypt data
/// when using the obfuscated transport.
///
/// You're not supposed to use this directly, You're probably looking for the
/// actual implementation in `grammers-mtproto`.
pub struct ObfuscatedCipher {
rx: ctr::Ctr128BE<aes::Aes256>,
tx: ctr::Ctr128BE<aes::Aes256>,
}

impl ObfuscatedCipher {
pub fn new(init: &[u8; 64]) -> Self {
let init_rev = init.iter().copied().rev().collect::<Vec<_>>();
Self {
rx: ctr::Ctr128BE::<aes::Aes256>::new(
GenericArray::from_slice(&init_rev[8..40]),
GenericArray::from_slice(&init_rev[40..56]),
),
tx: ctr::Ctr128BE::<aes::Aes256>::new(
GenericArray::from_slice(&init[8..40]),
GenericArray::from_slice(&init[40..56]),
),
}
}

pub fn encrypt(&mut self, buffer: &mut [u8]) {
self.tx.apply_keystream(buffer);
}

pub fn decrypt(&mut self, buffer: &mut [u8]) {
self.rx.apply_keystream(buffer);
}
}
23 changes: 15 additions & 8 deletions lib/grammers-mtproto/src/transport/abridged.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
use super::{Error, Transport, UnpackedOffset};
use super::{Error, Tagged, Transport, UnpackedOffset};
use grammers_crypto::DequeBuffer;

/// The lightest MTProto transport protocol available. This is an
Expand Down Expand Up @@ -63,7 +63,7 @@ impl Transport for Abridged {
}
}

fn unpack(&mut self, buffer: &[u8]) -> Result<UnpackedOffset, Error> {
fn unpack(&mut self, buffer: &mut [u8]) -> Result<UnpackedOffset, Error> {
if buffer.is_empty() {
return Err(Error::MissingBytes);
}
Expand Down Expand Up @@ -112,6 +112,13 @@ impl Transport for Abridged {
}
}

impl Tagged for Abridged {
fn init_tag(&mut self) -> [u8; 4] {
self.init = true;
[0xef, 0xef, 0xef, 0xef]
}
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down Expand Up @@ -160,7 +167,7 @@ mod tests {
let mut transport = Abridged::new();
let mut buffer = DequeBuffer::with_capacity(1, 0);
buffer.extend([1]);
assert_eq!(transport.unpack(&buffer[..]), Err(Error::MissingBytes));
assert_eq!(transport.unpack(&mut buffer[..]), Err(Error::MissingBytes));
}

#[test]
Expand All @@ -169,7 +176,7 @@ mod tests {
let orig = buffer.clone();
transport.pack(&mut buffer);
let n = 1; // init byte
let offset = transport.unpack(&buffer[n..][..]).unwrap();
let offset = transport.unpack(&mut buffer[n..][..]).unwrap();
assert_eq!(&buffer[n..][offset.data_start..offset.data_end], &orig[..]);
}

Expand All @@ -187,12 +194,12 @@ mod tests {
transport.pack(&mut buffer);
two_buffer.extend(&buffer[..]);

let offset = transport.unpack(&two_buffer[..]).unwrap();
let offset = transport.unpack(&mut two_buffer[..]).unwrap();
assert_eq!(&buffer[offset.data_start..offset.data_end], &orig[..]);
assert_eq!(offset.next_offset, single_size);

let n = offset.next_offset;
let offset = transport.unpack(&two_buffer[n..]).unwrap();
let offset = transport.unpack(&mut two_buffer[n..]).unwrap();
assert_eq!(&buffer[offset.data_start..offset.data_end], &orig[..]);
}

Expand All @@ -202,7 +209,7 @@ mod tests {
let orig = buffer.clone();
transport.pack(&mut buffer);
let n = 1; // init byte
let offset = transport.unpack(&buffer[n..]).unwrap();
let offset = transport.unpack(&mut buffer[n..]).unwrap();
assert_eq!(&buffer[n..][offset.data_start..offset.data_end], &orig[..]);
}

Expand All @@ -214,7 +221,7 @@ mod tests {
buffer.extend(&(-404_i32).to_le_bytes());

assert_eq!(
transport.unpack(&buffer[..]),
transport.unpack(&mut buffer[..]),
Err(Error::BadStatus { status: 404 })
);
}
Expand Down
16 changes: 8 additions & 8 deletions lib/grammers-mtproto/src/transport/full.rs
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ impl Transport for Full {
self.send_seq += 1;
}

fn unpack(&mut self, buffer: &[u8]) -> Result<UnpackedOffset, Error> {
fn unpack(&mut self, buffer: &mut [u8]) -> Result<UnpackedOffset, Error> {
// Need 4 bytes for the initial length
if buffer.len() < 4 {
return Err(Error::MissingBytes);
Expand Down Expand Up @@ -188,15 +188,15 @@ mod tests {
let mut transport = Full::new();
let mut buffer = DequeBuffer::with_capacity(3, 0);
buffer.extend([0, 1, 3]);
assert_eq!(transport.unpack(&buffer[..]), Err(Error::MissingBytes));
assert_eq!(transport.unpack(&mut buffer[..]), Err(Error::MissingBytes));
}

#[test]
fn unpack_normal() {
let (mut transport, mut buffer) = setup_pack(128);
let orig = buffer.clone();
transport.pack(&mut buffer);
let offset = transport.unpack(&buffer[..]).unwrap();
let offset = transport.unpack(&mut buffer[..]).unwrap();
assert_eq!(&buffer[offset.data_start..offset.data_end], &orig[..]);
}

Expand All @@ -214,12 +214,12 @@ mod tests {
transport.pack(&mut buffer);
two_buffer.extend(&buffer[..]);

let offset = transport.unpack(&two_buffer[..]).unwrap();
let offset = transport.unpack(&mut two_buffer[..]).unwrap();
assert_eq!(&buffer[offset.data_start..offset.data_end], &orig[..]);
assert_eq!(offset.next_offset, single_size);

let n = offset.next_offset;
let offset = transport.unpack(&two_buffer[n..]).unwrap();
let offset = transport.unpack(&mut two_buffer[n..]).unwrap();
assert_eq!(&buffer[offset.data_start..offset.data_end], &orig[..]);
}

Expand All @@ -230,7 +230,7 @@ mod tests {
buffer[4] = 1;

assert_eq!(
transport.unpack(&buffer[..]),
transport.unpack(&mut buffer[..]),
Err(Error::BadSeq {
expected: 0,
got: 1,
Expand All @@ -246,7 +246,7 @@ mod tests {
buffer[len - 1] ^= 0xff;

assert_eq!(
transport.unpack(&buffer[..]),
transport.unpack(&mut buffer[..]),
Err(Error::BadCrc {
expected: 932541318,
got: 3365237638,
Expand All @@ -261,7 +261,7 @@ mod tests {
buffer.extend(&(-404_i32).to_le_bytes());

assert_eq!(
transport.unpack(&buffer[..]),
transport.unpack(&mut buffer[..]),
Err(Error::BadStatus { status: 404 })
);
}
Expand Down
25 changes: 17 additions & 8 deletions lib/grammers-mtproto/src/transport/intermediate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
use super::{Error, Transport, UnpackedOffset};
use super::{Error, Tagged, Transport, UnpackedOffset};
use grammers_crypto::DequeBuffer;

/// A light MTProto transport protocol available that guarantees data padded
Expand All @@ -31,6 +31,8 @@ pub struct Intermediate {

#[allow(clippy::new_without_default)]
impl Intermediate {
const TAG: [u8; 4] = 0xee_ee_ee_ee_u32.to_le_bytes();

pub fn new() -> Self {
Self { init: false }
}
Expand All @@ -44,12 +46,12 @@ impl Transport for Intermediate {
buffer.extend_front(&(len as i32).to_le_bytes());

if !self.init {
buffer.extend_front(&0xee_ee_ee_ee_u32.to_le_bytes());
buffer.extend_front(&Self::TAG);
self.init = true;
}
}

fn unpack(&mut self, buffer: &[u8]) -> Result<UnpackedOffset, Error> {
fn unpack(&mut self, buffer: &mut [u8]) -> Result<UnpackedOffset, Error> {
if buffer.len() < 4 {
return Err(Error::MissingBytes);
}
Expand Down Expand Up @@ -84,6 +86,13 @@ impl Transport for Intermediate {
}
}

impl Tagged for Intermediate {
fn init_tag(&mut self) -> [u8; 4] {
self.init = true;
Self::TAG
}
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down Expand Up @@ -123,7 +132,7 @@ mod tests {
let mut transport = Intermediate::new();
let mut buffer = DequeBuffer::with_capacity(1, 0);
buffer.extend([1]);
assert_eq!(transport.unpack(&buffer[..],), Err(Error::MissingBytes));
assert_eq!(transport.unpack(&mut buffer[..],), Err(Error::MissingBytes));
}

#[test]
Expand All @@ -132,7 +141,7 @@ mod tests {
let orig = buffer.clone();
transport.pack(&mut buffer);
let n = 4; // init bytes
let offset = transport.unpack(&buffer[n..]).unwrap();
let offset = transport.unpack(&mut buffer[n..]).unwrap();
assert_eq!(&buffer[n..][offset.data_start..offset.data_end], &orig[..]);
}

Expand All @@ -150,12 +159,12 @@ mod tests {
transport.pack(&mut buffer);
two_buffer.extend(&buffer[..]);

let offset = transport.unpack(&two_buffer[..]).unwrap();
let offset = transport.unpack(&mut two_buffer[..]).unwrap();
assert_eq!(&buffer[offset.data_start..offset.data_end], &orig[..]);
assert_eq!(offset.next_offset, single_size);

let n = offset.next_offset;
let offset = transport.unpack(&two_buffer[n..]).unwrap();
let offset = transport.unpack(&mut two_buffer[n..]).unwrap();
assert_eq!(&buffer[offset.data_start..offset.data_end], &orig[..]);
}

Expand All @@ -167,7 +176,7 @@ mod tests {
buffer.extend(&(-404_i32).to_le_bytes());

assert_eq!(
transport.unpack(&buffer[..]),
transport.unpack(&mut buffer[..]),
Err(Error::BadStatus { status: 404 })
);
}
Expand Down
13 changes: 12 additions & 1 deletion lib/grammers-mtproto/src/transport/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,13 @@
mod abridged;
mod full;
mod intermediate;
mod obfuscated;

pub use abridged::Abridged;
pub use full::Full;
use grammers_crypto::DequeBuffer;
pub use intermediate::Intermediate;
pub use obfuscated::Obfuscated;
use std::fmt;

/// The error type reported by the different transports when something is wrong.
Expand Down Expand Up @@ -91,8 +93,17 @@ pub trait Transport {
fn pack(&mut self, buffer: &mut DequeBuffer<u8>);

/// Unpacks the input buffer in-place.
fn unpack(&mut self, buffer: &[u8]) -> Result<UnpackedOffset, Error>;
/// Subsequent calls to `unpack` should be made with the same buffer,
/// with the data on the ranges from previous `UnpackedOffset` removed.
fn unpack(&mut self, buffer: &mut [u8]) -> Result<UnpackedOffset, Error>;
Copy link
Contributor

@MOZGIII MOZGIII Dec 23, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There should probably be an explainer (in the comments) on why the buffer is mut here... As I was wondering.

Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

pack takes a mutable input so in my view it's reasonable for unpack to do so too.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

By why would we alter the buffer itself? Isn't it supposed to be the read buffer?

Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both packing unpacking occur in-place. None of the previous transports had to modify the buffer during unpack, but obfuscated does. The caller doesn't read from the buffer contents until the calls succeed.

Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(Forgot to add, but that said, I'm happy to have the documentation improved to leave this in a better place than a random GitHub discussion.)


/// Reset the state, as if a new instance was just created.
fn reset(&mut self);
}

/// The trait used by the obfuscated transport to get the transport tags.
pub trait Tagged {
/// Gets the transport tag for use in the obfuscated transport and
/// changes the internal state to avoid sending the tag again.
fn init_tag(&mut self) -> [u8; 4];
}
Loading
Loading