astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-06-18 08:58:30 +03:00
Code Issues Packages Projects Releases Wiki Activity

Implement shared MTProto framing and ME address role separation

Browse Source 
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-06-15 08:50:08 +03:00
parent d81d7dba62
commit 37d0184a0b
13 changed files with 415 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
src/stream/frame_codec.rs
View File

@@ -15,6 +15,7 @@ use crate::crypto::SecureRandom;
use crate::protocol::constants::{
ProtoTag, is_valid_secure_payload_len, secure_padding_len, secure_payload_len_from_wire_len,
};
use crate::protocol::framing::{encode_intermediate_header, parse_intermediate_header};
// ============= Unified Codec =============
@@ -197,13 +198,9 @@ fn decode_intermediate(src: &mut BytesMut, max_size: usize) -> io::Result<Option
}
let mut meta = FrameMeta::new();
let mut len = u32::from_le_bytes([src[0], src[1], src[2], src[3]]) as usize;
// Check QuickACK flag
if len >= 0x80000000 {
meta.quickack = true;
len -= 0x80000000;
}
let header = parse_intermediate_header([src[0], src[1], src[2], src[3]]);
let len = header.wire_len;
meta.quickack = header.quickack;
// Validate size
if len > max_size {
@@ -239,10 +236,12 @@ fn encode_intermediate(frame: &Frame, dst: &mut BytesMut) -> io::Result<()> {
dst.reserve(4 + data.len());
let mut len = data.len() as u32;
if frame.meta.quickack {
len |= 0x80000000;
}
let len = encode_intermediate_header(data.len(), frame.meta.quickack).ok_or_else(|| {
Error::new(
ErrorKind::InvalidInput,
format!("frame too large: {} bytes", data.len()),
)
})?;
dst.extend_from_slice(&len.to_le_bytes());
dst.extend_from_slice(data);
@@ -258,13 +257,9 @@ fn decode_secure(src: &mut BytesMut, max_size: usize) -> io::Result<Option<Frame
}
let mut meta = FrameMeta::new();
let mut len = u32::from_le_bytes([src[0], src[1], src[2], src[3]]) as usize;
// Check QuickACK flag
if len >= 0x80000000 {
meta.quickack = true;
len -= 0x80000000;
}
let header = parse_intermediate_header([src[0], src[1], src[2], src[3]]);
let len = header.wire_len;
meta.quickack = header.quickack;
// Validate size
if len > max_size {
@@ -323,10 +318,12 @@ fn encode_secure(frame: &Frame, dst: &mut BytesMut, rng: &SecureRandom) -> io::R
let total_len = data.len() + padding_len;
dst.reserve(4 + total_len);
let mut len = total_len as u32;
if frame.meta.quickack {
len |= 0x80000000;
}
let len = encode_intermediate_header(total_len, frame.meta.quickack).ok_or_else(|| {
Error::new(
ErrorKind::InvalidInput,
format!("frame too large: {} bytes", total_len),
)
})?;
dst.extend_from_slice(&len.to_le_bytes());
dst.extend_from_slice(data);

83
src/stream/frame_stream.rs
View File

@@ -5,6 +5,7 @@
use super::traits::{FrameMeta, LayeredStream};
use crate::crypto::{SecureRandom, crc32};
use crate::protocol::constants::*;
use crate::protocol::framing::{encode_intermediate_header, parse_intermediate_header};
use bytes::Bytes;
use std::io::{Error, ErrorKind, Result};
use std::sync::Arc;
@@ -105,10 +106,17 @@ impl<W: AsyncWrite + Unpin> AbridgedFrameWriter<W> {
if len_div_4 < 0x7f {
// Short length (1 byte)
self.upstream.write_all(&[len_div_4 as u8]).await?;
let mut first = len_div_4 as u8;
if meta.quickack {
first |= 0x80;
}
self.upstream.write_all(&[first]).await?;
} else if len_div_4 < (1 << 24) {
// Long length (4 bytes: 0x7f + 3 bytes)
let mut header = [0x7f, 0, 0, 0];
if meta.quickack {
header[0] |= 0x80;
}
header[1..4].copy_from_slice(&(len_div_4 as u32).to_le_bytes()[..3]);
self.upstream.write_all(&header).await?;
} else {
@@ -160,13 +168,9 @@ impl<R: AsyncRead + Unpin> IntermediateFrameReader<R> {
let mut len_bytes = [0u8; 4];
self.upstream.read_exact(&mut len_bytes).await?;
let mut len = u32::from_le_bytes(len_bytes) as usize;
// Check QuickACK flag (high bit)
if len > 0x80000000 {
meta.quickack = true;
len -= 0x80000000;
}
let header = parse_intermediate_header(len_bytes);
let len = header.wire_len;
meta.quickack = header.quickack;
// Read data
let mut data = vec![0u8; len];
@@ -204,7 +208,13 @@ impl<W: AsyncWrite + Unpin> IntermediateFrameWriter<W> {
if meta.simple_ack {
self.upstream.write_all(data).await?;
} else {
let len_bytes = (data.len() as u32).to_le_bytes();
let len = encode_intermediate_header(data.len(), meta.quickack).ok_or_else(|| {
Error::new(
ErrorKind::InvalidInput,
format!("Frame too large: {} bytes", data.len()),
)
})?;
let len_bytes = len.to_le_bytes();
self.upstream.write_all(&len_bytes).await?;
self.upstream.write_all(data).await?;
}
@@ -249,13 +259,9 @@ impl<R: AsyncRead + Unpin> SecureIntermediateFrameReader<R> {
let mut len_bytes = [0u8; 4];
self.upstream.read_exact(&mut len_bytes).await?;
let mut len = u32::from_le_bytes(len_bytes) as usize;
// Check QuickACK flag
if len > 0x80000000 {
meta.quickack = true;
len -= 0x80000000;
}
let header = parse_intermediate_header(len_bytes);
let len = header.wire_len;
meta.quickack = header.quickack;
// Read data (including padding)
let mut data = vec![0u8; len];
@@ -316,7 +322,13 @@ impl<W: AsyncWrite + Unpin> SecureIntermediateFrameWriter<W> {
let padding = self.rng.bytes(padding_len);
let total_len = data.len() + padding_len;
let len_bytes = (total_len as u32).to_le_bytes();
let len = encode_intermediate_header(total_len, meta.quickack).ok_or_else(|| {
Error::new(
ErrorKind::InvalidInput,
format!("Frame too large: {total_len} bytes"),
)
})?;
let len_bytes = len.to_le_bytes();
self.upstream.write_all(&len_bytes).await?;
self.upstream.write_all(data).await?;
@@ -623,6 +635,43 @@ mod tests {
assert_eq!(&received[..], &data[..]);
}
#[tokio::test]
async fn test_intermediate_quickack_zero_length_roundtrip() {
let (client, server) = duplex(1024);
let mut writer = IntermediateFrameWriter::new(client);
let mut reader = IntermediateFrameReader::new(server);
writer
.write_frame(&[], &FrameMeta::new().with_quickack())
.await
.unwrap();
writer.flush().await.unwrap();
let (received, meta) = reader.read_frame().await.unwrap();
assert!(received.is_empty());
assert!(meta.quickack);
}
#[tokio::test]
async fn test_abridged_quickack_roundtrip() {
let (client, server) = duplex(1024);
let mut writer = AbridgedFrameWriter::new(client);
let mut reader = AbridgedFrameReader::new(server);
let data = vec![1u8, 2, 3, 4];
writer
.write_frame(&data, &FrameMeta::new().with_quickack())
.await
.unwrap();
writer.flush().await.unwrap();
let (received, meta) = reader.read_frame().await.unwrap();
assert_eq!(&received[..], &data[..]);
assert!(meta.quickack);
}
#[tokio::test]
async fn test_secure_intermediate_padding() {
let (client, server) = duplex(1024);