wiz0u/WTelegramClient
1
0
Fork 0
mirror of https://github.com/wiz0u/WTelegramClient.git synced 2025-12-06 06:52:01 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

Optimized reactor frame handling

Browse source
This commit is contained in:
Wizou 2021-09-23 05:37:00 +02:00
parent 223d8984cf
commit fe7bc6f61c
2 changed files with 37 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
src/Client.cs
View file

@ -181,12 +181,24 @@ namespace WTelegram
private async Task Reactor(NetworkStream stream, CancellationTokenSource cts)
{
const int MinBufferSize = 1024;
var data = new byte[MinBufferSize];
while (!cts.IsCancellationRequested)
{
ITLObject obj = null;
try
{
obj = await RecvAsync(stream, cts.Token);
if (await FullReadAsync(stream, data, 4, cts.Token) != 4)
throw new ApplicationException("Could not read payload length : Connection shut down");
int payloadLen = BinaryPrimitives.ReadInt32LittleEndian(data);
if (payloadLen > data.Length)
data = new byte[payloadLen];
else if (Math.Max(payloadLen, MinBufferSize) < data.Length / 4)
data = new byte[Math.Max(payloadLen, MinBufferSize)];
if (await FullReadAsync(stream, data, payloadLen, cts.Token) != payloadLen)
throw new ApplicationException("Could not read frame data : Connection shut down");
obj = ReadFrame(data, payloadLen);
}
catch (Exception ex) // an exception in RecvAsync is always fatal
{
@ -223,27 +235,26 @@ namespace WTelegram
}
}
internal async Task<ITLObject> RecvAsync(NetworkStream stream, CancellationToken ct)
internal ITLObject ReadFrame(byte[] data, int dataLen)
{
var data = await RecvFrameAsync(stream, ct);
if (data.Length == 4 && data[3] == 0xFF)
if (dataLen == 4 && data[3] == 0xFF)
{
int error_code = -BinaryPrimitives.ReadInt32LittleEndian(data);
throw new RpcException(error_code, TransportError(error_code));
}
if (data.Length < 24) // authKeyId+msgId+length+ctorNb | authKeyId+msgKey
throw new ApplicationException($"Packet payload too small: {data.Length}");
if (dataLen < 24) // authKeyId+msgId+length+ctorNb | authKeyId+msgKey
throw new ApplicationException($"Packet payload too small: {dataLen}");
long authKeyId = BinaryPrimitives.ReadInt64LittleEndian(data);
if (authKeyId != _session.AuthKeyID)
throw new ApplicationException($"Received a packet encrypted with unexpected key {authKeyId:X}");
if (authKeyId == 0) // Unencrypted message
{
using var reader = new TL.BinaryReader(new MemoryStream(data, 8, data.Length - 8), this);
using var reader = new TL.BinaryReader(new MemoryStream(data, 8, dataLen - 8), this);
long msgId = _lastRecvMsgId = reader.ReadInt64();
if ((msgId & 1) == 0) throw new ApplicationException($"Invalid server msgId {msgId}");
int length = reader.ReadInt32();
if (length != data.Length - 20) throw new ApplicationException($"Unexpected unencrypted length {length} != {data.Length - 20}");
if (length != dataLen - 20) throw new ApplicationException($"Unexpected unencrypted length {length} != {dataLen - 20}");
var obj = reader.ReadTLObject();
Helpers.Log(1, $"Receiving {obj.GetType().Name,-50} {_session.MsgIdToStamp(msgId):u} {((msgId & 2) == 0 ? "" : "NAR")} unencrypted");
@ -251,12 +262,7 @@ namespace WTelegram
}
else
{
#if MTPROTO1
byte[] msgKeyLarge = data[4..24];
#else
byte[] msgKeyLarge = data[0..24];
#endif
byte[] decrypted_data = EncryptDecryptMessage(data.AsSpan(24), false, _session.AuthKey, msgKeyLarge);
byte[] decrypted_data = EncryptDecryptMessage(data.AsSpan(24, dataLen - 24), false, _session.AuthKey, data, 8);
if (decrypted_data.Length < 36) // header below+ctorNb
throw new ApplicationException($"Decrypted packet too small: {decrypted_data.Length}");
using var reader = new TL.BinaryReader(new MemoryStream(decrypted_data), this);
@ -318,24 +324,9 @@ namespace WTelegram
};
}
private static async Task<byte[]> RecvFrameAsync(NetworkStream stream, CancellationToken ct)
{
byte[] overhead = new byte[4];
if (await FullReadAsync(stream, overhead, 4, ct) != 4)
throw new ApplicationException("Could not read payload length : Connection shut down");
int length = BinaryPrimitives.ReadInt32LittleEndian(overhead);
if (length <= 0)
throw new ApplicationException("Invalid frame_len");
var payload = new byte[length];
if (await FullReadAsync(stream, payload, length, ct) != length)
throw new ApplicationException("Could not read frame data : Connection shut down");
return payload;
}
private static async Task<int> FullReadAsync(Stream stream, byte[] buffer, int length, CancellationToken ct = default)
{
for (int offset = 0; offset != length;)
for (int offset = 0; offset < length;)
{
var read = await stream.ReadAsync(buffer, offset, length - offset, ct);
if (read == 0) return offset;
@ -406,7 +397,7 @@ namespace WTelegram
var msgKeyLarge = Sha256.ComputeHash(clearBuffer, 0, prepend + clearLength + padding);
const int msgKeyOffset = 8; // msg_key = middle 128-bits of SHA256(authkey_part+plaintext+padding)
#endif
byte[] encrypted_data = EncryptDecryptMessage(clearBuffer.AsSpan(prepend, clearLength + padding), true, _session.AuthKey, msgKeyLarge);
byte[] encrypted_data = EncryptDecryptMessage(clearBuffer.AsSpan(prepend, clearLength + padding), true, _session.AuthKey, msgKeyLarge, msgKeyOffset);
writer.Write(_session.AuthKeyID); // int64 auth_key_id
writer.Write(msgKeyLarge, msgKeyOffset, 16); // int128 msg_key

31
src/Encryption.cs
View file

@ -277,29 +277,29 @@ j4WcDuXc2CTHgH8gFTNhp/Y8/SpDOhvn9QIDAQAB
#endif
}
internal static byte[] EncryptDecryptMessage(Span<byte> input, bool encrypt, byte[] authKey, byte[] msgKeyLarge)
internal static byte[] EncryptDecryptMessage(Span<byte> input, bool encrypt, byte[] authKey, byte[] msgKey, int msgKeyOffset)
{
// first, construct AES key & IV
int x = encrypt ? 0 : 8;
byte[] aes_key = new byte[32], aes_iv = new byte[32];
int x = encrypt ? 0 : 8;
#if MTPROTO1
var sha1 = encrypt ? Sha1 : Sha1Recv;
sha1.Initialize();
sha1.TransformBlock(msgKeyLarge, 4, 16, null, 0); // msgKey
sha1.TransformFinalBlock(authKey, x, 32); // authKey[x:32]
sha1.TransformBlock(msgKey, msgKeyOffset, 16, null, 0); // msgKey
sha1.TransformFinalBlock(authKey, x, 32); // authKey[x:32]
var sha1_a = sha1.Hash;
sha1.Initialize();
sha1.TransformBlock(authKey, 32 + x, 16, null, 0); // authKey[32+x:16]
sha1.TransformBlock(msgKeyLarge, 4, 16, null, 0); // msgKey
sha1.TransformFinalBlock(authKey, 48 + x, 16); // authKey[48+x:16]
sha1.TransformBlock(authKey, 32 + x, 16, null, 0); // authKey[32+x:16]
sha1.TransformBlock(msgKey, msgKeyOffset, 16, null, 0); // msgKey
sha1.TransformFinalBlock(authKey, 48 + x, 16); // authKey[48+x:16]
var sha1_b = sha1.Hash;
sha1.Initialize();
sha1.TransformBlock(authKey, 64 + x, 32, null, 0); // authKey[64+x:32]
sha1.TransformFinalBlock(msgKeyLarge, 4, 16); // msgKey
sha1.TransformBlock(authKey, 64 + x, 32, null, 0); // authKey[64+x:32]
sha1.TransformFinalBlock(msgKey, msgKeyOffset, 16); // msgKey
var sha1_c = sha1.Hash;
sha1.Initialize();
sha1.TransformBlock(msgKeyLarge, 4, 16, null, 0); // msgKey
sha1.TransformFinalBlock(authKey, 96 + x, 32); // authKey[96+x:32]
sha1.TransformBlock(msgKey, msgKeyOffset, 16, null, 0); // msgKey
sha1.TransformFinalBlock(authKey, 96 + x, 32); // authKey[96+x:32]
var sha1_d = sha1.Hash;
Array.Copy(sha1_a, 0, aes_key, 0, 8);
Array.Copy(sha1_b, 8, aes_key, 8, 12);
@ -311,12 +311,12 @@ j4WcDuXc2CTHgH8gFTNhp/Y8/SpDOhvn9QIDAQAB
#else
var sha256 = encrypt ? Sha256 : Sha256Recv;
sha256.Initialize();
sha256.TransformBlock(msgKeyLarge, 8, 16, null, 0); // msgKey
sha256.TransformFinalBlock(authKey, x, 36); // authKey[x:36]
sha256.TransformBlock(msgKey, msgKeyOffset, 16, null, 0); // msgKey
sha256.TransformFinalBlock(authKey, x, 36); // authKey[x:36]
var sha256_a = sha256.Hash;
sha256.Initialize();
sha256.TransformBlock(authKey, 40 + x, 36, null, 0); // authKey[40+x:36]
sha256.TransformFinalBlock(msgKeyLarge, 8, 16); // msgKey
sha256.TransformBlock(authKey, 40 + x, 36, null, 0); // authKey[40+x:36]
sha256.TransformFinalBlock(msgKey, msgKeyOffset, 16); // msgKey
var sha256_b = sha256.Hash;
Array.Copy(sha256_a, 0, aes_key, 0, 8);
Array.Copy(sha256_b, 8, aes_key, 8, 16);
@ -325,7 +325,6 @@ j4WcDuXc2CTHgH8gFTNhp/Y8/SpDOhvn9QIDAQAB
Array.Copy(sha256_a, 8, aes_iv, 8, 16);
Array.Copy(sha256_b, 24, aes_iv, 24, 8);
#endif
return AES_IGE_EncryptDecrypt(input, aes_key, aes_iv, encrypt);
}