From cdba0f708803e9084976ca51e387ad1c05b951c9 Mon Sep 17 00:00:00 2001 From: Wizou <11647984+wiz0u@users.noreply.github.com> Date: Thu, 27 Jan 2022 17:34:16 +0100 Subject: [PATCH] optimize session writes --- src/Client.cs | 6 ++-- src/Session.cs | 83 +++++++++++++++++++++++++++----------------------- 2 files changed, 48 insertions(+), 41 deletions(-) diff --git a/src/Client.cs b/src/Client.cs index fa6ec17..73ca7ba 100644 --- a/src/Client.cs +++ b/src/Client.cs @@ -364,7 +364,7 @@ namespace WTelegram } finally { - _session.Save(); + lock (_session) _session.Save(); } Helpers.Log(2, $"Connected to {(TLConfig.test_mode ? "Test DC" : "DC")} {TLConfig.this_dc}... {TLConfig.flags & (Config.Flags)~0xE00}"); } @@ -907,7 +907,7 @@ namespace WTelegram else if (rpcError.error_code == 500 && rpcError.error_message == "AUTH_RESTART") { _session.UserId = 0; // force a full login authorization flow, next time - _session.Save(); + lock (_session) _session.Save(); } throw new RpcException(rpcError.error_code, rpcError.error_message); case ReactorError: @@ -1186,7 +1186,7 @@ namespace WTelegram if (authorization is not Auth_Authorization { user: User user }) throw new ApplicationException("Failed to get Authorization: " + authorization.GetType().Name); _session.UserId = _dcSession.UserId = user.id; - _session.Save(); + lock (_session) _session.Save(); return user; } diff --git a/src/Session.cs b/src/Session.cs index 6b1e2aa..643b9c6 100644 --- a/src/Session.cs +++ b/src/Session.cs @@ -9,7 +9,7 @@ using System.Text.Json; namespace WTelegram { - internal class Session + internal class Session : IDisposable { public int ApiId; public long UserId; @@ -39,11 +39,24 @@ namespace WTelegram private readonly DateTime _sessionStart = DateTime.UtcNow; private readonly SHA256 _sha256 = SHA256.Create(); private Stream _store; - private byte[] _rgbKey; // 32-byte encryption key - private static readonly Aes aes = Aes.Create(); + private byte[] _encrypted = new byte[16]; + private ICryptoTransform _encryptor; + private Utf8JsonWriter _jsonWriter; + private readonly MemoryStream _jsonStream = new(4096); + + public void Dispose() + { + _sha256.Dispose(); + _store.Dispose(); + _encryptor.Dispose(); + _jsonWriter.Dispose(); + _jsonStream.Dispose(); + } internal static Session LoadOrCreate(Stream store, byte[] rgbKey) { + using var aes = Aes.Create(); + Session session = null; try { var length = (int)store.Length; @@ -52,11 +65,13 @@ namespace WTelegram var input = new byte[length]; if (store.Read(input, 0, length) != length) throw new ApplicationException($"Can't read session block ({store.Position}, {length})"); - var session = Load(input, rgbKey); - session._store = store; - session._rgbKey = rgbKey; + using var sha256 = SHA256.Create(); + using var decryptor = aes.CreateDecryptor(rgbKey, input[0..16]); + var utf8Json = decryptor.TransformFinalBlock(input, 16, input.Length - 16); + if (!sha256.ComputeHash(utf8Json, 32, utf8Json.Length - 32).SequenceEqual(utf8Json[0..32])) + throw new ApplicationException("Integrity check failed in session loading"); + session = JsonSerializer.Deserialize(utf8Json.AsSpan(32), Helpers.JsonOptions); Helpers.Log(2, "Loaded previous session"); - return session; } } catch (Exception ex) @@ -64,48 +79,40 @@ namespace WTelegram store.Dispose(); throw new ApplicationException($"Exception while reading session file: {ex.Message}\nDelete the file to start a new session", ex); } - return new Session { _store = store, _rgbKey = rgbKey }; + session ??= new Session(); + session._store = store; + Encryption.RNG.GetBytes(session._encrypted, 0, 16); + session._encryptor = aes.CreateEncryptor(rgbKey, session._encrypted); + session._jsonWriter = new Utf8JsonWriter(session._jsonStream, default); + return session; } - internal void Dispose() => _store.Dispose(); - - internal static Session Load(byte[] input, byte[] rgbKey) + internal void Save() // must be called with lock(session) { - using var sha256 = SHA256.Create(); - using var decryptor = aes.CreateDecryptor(rgbKey, input[0..16]); - var utf8Json = decryptor.TransformFinalBlock(input, 16, input.Length - 16); - if (!sha256.ComputeHash(utf8Json, 32, utf8Json.Length - 32).SequenceEqual(utf8Json[0..32])) - throw new ApplicationException("Integrity check failed in session loading"); - return JsonSerializer.Deserialize(utf8Json.AsSpan(32), Helpers.JsonOptions); - } - - internal void Save() - { - var utf8Json = JsonSerializer.SerializeToUtf8Bytes(this, Helpers.JsonOptions); - var finalBlock = new byte[16]; - var output = new byte[(16 + 32 + utf8Json.Length + 16) & ~15]; - Encryption.RNG.GetBytes(output, 0, 16); - using var encryptor = aes.CreateEncryptor(_rgbKey, output[0..16]); - encryptor.TransformBlock(_sha256.ComputeHash(utf8Json), 0, 32, output, 16); - encryptor.TransformBlock(utf8Json, 0, utf8Json.Length & ~15, output, 48); - utf8Json.AsSpan(utf8Json.Length & ~15).CopyTo(finalBlock); - encryptor.TransformFinalBlock(finalBlock, 0, utf8Json.Length & 15).CopyTo(output.AsMemory(48 + utf8Json.Length & ~15)); - lock (_store) - { - _store.Position = 0; - _store.Write(output, 0, output.Length); - _store.SetLength(output.Length); - } + JsonSerializer.Serialize(_jsonWriter, this, Helpers.JsonOptions); + var utf8Json = _jsonStream.GetBuffer(); + var utf8JsonLen = (int)_jsonStream.Position; + int encryptedLen = 64 + (utf8JsonLen & ~15); + if (encryptedLen > _encrypted.Length) + Array.Copy(_encrypted, _encrypted = new byte[encryptedLen + 256], 16); + _encryptor.TransformBlock(_sha256.ComputeHash(utf8Json, 0, utf8JsonLen), 0, 32, _encrypted, 16); + _encryptor.TransformBlock(utf8Json, 0, encryptedLen - 64, _encrypted, 48); + _encryptor.TransformFinalBlock(utf8Json, encryptedLen - 64, utf8JsonLen & 15).CopyTo(_encrypted, encryptedLen - 16); + _store.Position = 0; + _store.Write(_encrypted, 0, encryptedLen); + _store.SetLength(encryptedLen); + _jsonStream.Position = 0; + _jsonWriter.Reset(); } } internal class SessionStore : FileStream { public override long Length { get; } - private readonly byte[] _header = new byte[8]; - private int _nextPosition = 8; public override long Position { get => base.Position; set { } } public override void SetLength(long value) { } + private readonly byte[] _header = new byte[8]; + private int _nextPosition = 8; public SessionStore(string pathname) : base(pathname, FileMode.OpenOrCreate, FileAccess.ReadWrite, FileShare.None, 1) // no buffering