diff --git a/Cargo.lock b/Cargo.lock index 07de357..661aee3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -559,7 +559,7 @@ checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" [[package]] name = "tgcryptors-core" -version = "1.3.0" +version = "1.3.1" dependencies = [ "criterion", "rayon", @@ -568,7 +568,7 @@ dependencies = [ [[package]] name = "tgcryptors-python" -version = "1.3.0" +version = "1.3.1" dependencies = [ "pyo3", "tgcryptors-core", diff --git a/Cargo.toml b/Cargo.toml index 4780bff..4d1f870 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,7 +6,7 @@ members = [ resolver = "2" [workspace.package] -version = "1.3.0" +version = "1.3.1" edition = "2021" license = "MIT" rust-version = "1.86" diff --git a/README.md b/README.md index 2f35dc6..fff3699 100644 --- a/README.md +++ b/README.md @@ -95,6 +95,27 @@ decrypted = tgcryptors.ctr256_decrypt(encrypted, key, iv, state) assert decrypted == data ``` +TgCrypto-compatible `bytearray` mode: `data`, `key`, `iv`, and `state` may also +be passed as `bytearray`. When `iv` or `state` are `bytearray`, they are +updated **in place** — after the call they hold the advanced counter and +residual offset, so passing the same objects to the next call continues the +stream across chunks. This is how pyrogram and its forks (kurigram, pyrofork, +...) drive the obfuscated TCP transports and CDN file downloads: + +```python +import os +import tgcrypto + +key = os.urandom(32) +iv = bytearray(os.urandom(16)) +state = bytearray(1) + +first = tgcrypto.ctr256_encrypt(data[:100], key, iv, state) # iv/state advance in place +second = tgcrypto.ctr256_encrypt(data[100:], key, iv, state) # continues the stream +``` + +Passing plain `bytes` behaves as a stateless one-shot call. + ### CBC Mode **Note**: Data must be padded to a multiple of the block size (16 bytes). @@ -237,6 +258,15 @@ uv build --wheel ## Changelog +### 1.3.1 + +- `ctr256_encrypt`/`ctr256_decrypt` now accept `bytearray` (and `bytes`) for + `data`, `key`, `iv`, and `state`. When `iv` or `state` are passed as + `bytearray`, the advanced counter and residual byte offset are written back + in place, matching TgCrypto's stateful semantics. This fixes a crash at + connect time with pyrogram and its forks, whose obfuscated TCP transports + and CDN downloads pass `bytearray` and rely on the in-place carry (issue #4). + ### 1.3.0 - Bumped MSRV from Rust 1.83 to 1.86. diff --git a/pyproject.toml b/pyproject.toml index a4fa0f4..4782b6c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "TgCryptoRust" -version = "1.3.0" +version = "1.3.1" description = "Rust-powered drop-in replacement for TgCrypto" readme = "README.md" requires-python = ">=3.9,<3.15" diff --git a/tests/test_python_api.py b/tests/test_python_api.py index 252135e..128f139 100644 --- a/tests/test_python_api.py +++ b/tests/test_python_api.py @@ -43,6 +43,87 @@ def test_ctr_stream_matches_one_shot(self) -> None: self.assertEqual(actual, expected) + def test_ctr_accepts_bytearray(self) -> None: + expected = tgcrypto.ctr256_encrypt(self.data, self.key, self.iv_cbc, b"\x00") + + ciphertext = tgcrypto.ctr256_encrypt( + bytearray(self.data), + bytearray(self.key), + bytearray(self.iv_cbc), + bytearray(1), + ) + + self.assertEqual(ciphertext, expected) + + def test_ctr_bytearray_state_carries_across_calls(self) -> None: + data = self.data + b"payload" * 40 + expected = tgcrypto.ctr256_encrypt(data, self.key, self.iv_cbc, b"\x00") + + enc_iv = bytearray(self.iv_cbc) + enc_state = bytearray(1) + ciphertext = ( + tgcrypto.ctr256_encrypt(data[:100], self.key, enc_iv, enc_state) + + tgcrypto.ctr256_encrypt(data[100:333], self.key, enc_iv, enc_state) + + tgcrypto.ctr256_encrypt(data[333:], self.key, enc_iv, enc_state) + ) + + self.assertEqual(ciphertext, expected) + + dec_iv = bytearray(self.iv_cbc) + dec_state = bytearray(1) + plaintext = ( + tgcrypto.ctr256_decrypt(ciphertext[:5], self.key, dec_iv, dec_state) + + tgcrypto.ctr256_decrypt(ciphertext[5:700], self.key, dec_iv, dec_state) + + tgcrypto.ctr256_decrypt(ciphertext[700:], self.key, dec_iv, dec_state) + ) + + self.assertEqual(plaintext, data) + + def test_ctr_bytearray_large_chunked_roundtrip(self) -> None: + data = bytes(range(256)) * 2048 + key = self.key + iv0 = self.iv_cbc + + enc_iv = bytearray(iv0) + enc_state = bytearray(1) + ciphertext = b"" + for i in range(0, len(data), 65553): + ciphertext += tgcrypto.ctr256_decrypt( + data[i : i + 65553], key, enc_iv, enc_state + ) + + dec_iv = bytearray(iv0) + dec_state = bytearray(1) + plaintext = b"" + for i in range(0, len(ciphertext), 65553): + plaintext += tgcrypto.ctr256_decrypt( + ciphertext[i : i + 65553], key, dec_iv, dec_state + ) + + self.assertEqual(plaintext, data) + + def test_ctr_bytearray_mutates_iv_and_state_in_place(self) -> None: + data = self.data + b"xyz" # 67 bytes, not block aligned + iv = bytearray(self.iv_cbc) + state = bytearray(1) + + tgcrypto.ctr256_encrypt(data, self.key, iv, state) + + self.assertEqual(state[0], len(data) % 16) + self.assertNotEqual(iv, bytearray(self.iv_cbc)) + + def test_ctr_bytes_do_not_mutate_iv_or_state(self) -> None: + data = self.data + b"xyz" # 67 bytes, not block aligned + iv = self.iv_cbc + state = b"\x00" + + ciphertext1 = tgcrypto.ctr256_encrypt(data, self.key, iv, state) + ciphertext2 = tgcrypto.ctr256_encrypt(data, self.key, iv, state) + + self.assertEqual(iv, self.iv_cbc) + self.assertEqual(state, b"\x00") + self.assertEqual(ciphertext1, ciphertext2) + def test_ige_stream_matches_one_shot(self) -> None: expected = tgcrypto.ige256_encrypt(self.data, self.key, self.iv_ige) stream = tgcrypto.Ige256(self.key, self.iv_ige) @@ -88,6 +169,22 @@ def test_validation_errors_are_explicit(self) -> None: with self.assertRaisesRegex(ValueError, "State value must be in the range \\[0, 15\\]"): tgcrypto.ctr256_encrypt(self.data, self.key, self.iv_cbc, b"\x10") + def test_ctr_bytearray_validation_errors_are_explicit(self) -> None: + with self.assertRaisesRegex(ValueError, "Key must be exactly 32 bytes"): + tgcrypto.ctr256_encrypt(self.data, bytearray(31), self.iv_cbc, b"\x00") + + with self.assertRaisesRegex(ValueError, "IV must be exactly 16 bytes"): + tgcrypto.ctr256_encrypt(self.data, self.key, bytearray(15), b"\x00") + + with self.assertRaisesRegex(ValueError, "State value must be in the range \\[0, 15\\]"): + tgcrypto.ctr256_encrypt(self.data, self.key, self.iv_cbc, bytearray(b"\x10")) + + with self.assertRaises(TypeError): + tgcrypto.ctr256_encrypt(self.data, "not bytes", self.iv_cbc, b"\x00") + + with self.assertRaises(TypeError): + tgcrypto.ctr256_encrypt(self.data, self.key, "not bytes", b"\x00") + def test_docstrings_are_available(self) -> None: self.assertIn("Encrypt bytes with AES-256-CTR", tgcrypto.ctr256_encrypt.__doc__) self.assertIn("Stateful AES-256-CTR stream cipher", tgcrypto.Ctr256.__doc__) @@ -95,7 +192,7 @@ def test_docstrings_are_available(self) -> None: self.assertIn("Stateful AES-256-IGE stream cipher", tgcrypto.Ige256.__doc__) def test_runtime_metadata_is_available(self) -> None: - self.assertEqual(tgcrypto.__version__, "1.3.0") + self.assertEqual(tgcrypto.__version__, "1.3.1") info = tgcrypto.runtime_info() diff --git a/tgcryptors-python/src/lib.rs b/tgcryptors-python/src/lib.rs index 045e033..92e05eb 100644 --- a/tgcryptors-python/src/lib.rs +++ b/tgcryptors-python/src/lib.rs @@ -10,10 +10,11 @@ //! - `cbc256_decrypt(data, key, iv) -> bytes` use core::slice; -use pyo3::exceptions::{PyOverflowError, PyRuntimeError, PyValueError}; +use pyo3::exceptions::{PyOverflowError, PyRuntimeError, PyTypeError, PyValueError}; use pyo3::ffi; use pyo3::prelude::*; -use pyo3::types::{PyBytes, PyDict}; +use pyo3::sync::critical_section::with_critical_section; +use pyo3::types::{PyAny, PyByteArray, PyBytes, PyDict}; use std::panic::{catch_unwind, AssertUnwindSafe}; use tgcryptors_core::AES_BLOCK_SIZE; @@ -47,6 +48,75 @@ fn validate_ctr_state(state: &[u8]) -> PyResult { Ok(state) } +/// A `bytes` or `bytearray` argument. +/// +/// `bytes` is read zero-copy and treated as stateless input. `bytearray` +/// contents are copied out up front, because the buffer may be mutated by +/// Python code while the GIL is released during the operation; the original +/// object is kept so mutated state can be written back afterwards (see +/// [`BufferInput::into_source`]). +enum BufferInput<'py> { + Bytes(Bound<'py, PyBytes>), + ByteArray { + data: Vec, + source: Bound<'py, PyByteArray>, + }, +} + +impl<'py> BufferInput<'py> { + fn from_any(ob: &Bound<'py, PyAny>, label: &str) -> PyResult { + if let Ok(bytes) = ob.cast::() { + return Ok(BufferInput::Bytes(bytes.clone())); + } + if let Ok(array) = ob.cast::() { + return Ok(BufferInput::ByteArray { + data: array.to_vec(), + source: array.clone(), + }); + } + Err(PyTypeError::new_err(format!( + "{label} must be bytes or bytearray" + ))) + } + + fn as_bytes(&self) -> &[u8] { + match self { + BufferInput::Bytes(bytes) => bytes.as_bytes(), + BufferInput::ByteArray { data, .. } => data, + } + } + + /// Take the underlying `bytearray` to write state back into, if any. + fn into_source(self) -> Option> { + match self { + BufferInput::Bytes(_) => None, + BufferInput::ByteArray { source, .. } => Some(source), + } + } +} + +/// Overwrite the contents of a `bytearray` with `value`. +/// +/// The buffer length is guaranteed to match `value` because extraction copied +/// the exact same length and Python code cannot legally resize the object +/// while this function holds it inside the critical section. +fn write_back(byte_array: &Bound<'_, PyByteArray>, value: &[u8]) { + with_critical_section(byte_array.as_any(), || { + debug_assert_eq!( + byte_array.len(), + value.len(), + "write_back: bytearray length ({}) must match value length ({})", + byte_array.len(), + value.len() + ); + + // SAFETY: the critical section prevents concurrent mutation of the + // buffer, and the buffer was not resized since extraction; the + // lengths are asserted to match. + unsafe { byte_array.as_bytes_mut() }.copy_from_slice(value); + }); +} + /// Zero-copy PyBytes allocation with GIL release and panic isolation. /// /// Allocates uninitialized Python bytes, releases the GIL, @@ -156,48 +226,83 @@ fn ige256_decrypt<'py>( /// Encrypt bytes with AES-256-CTR. /// /// Args: -/// data: Plaintext bytes of any length. -/// key: AES-256 key, exactly 32 bytes. -/// iv: Counter block, exactly 16 bytes. -/// state: Residual CTR byte offset encoded as a one-byte `bytes` object. +/// data: Plaintext bytes of any length. May be `bytes` or `bytearray`. +/// key: AES-256 key, exactly 32 bytes. May be `bytes` or `bytearray`. +/// iv: Counter block, exactly 16 bytes. May be `bytes` or `bytearray`. +/// state: Residual CTR byte offset encoded as a one-byte `bytes` or +/// `bytearray`. /// /// Returns: /// The encrypted ciphertext as `bytes`. /// +/// Notes: +/// When `iv` or `state` are passed as `bytearray`, they are updated in +/// place: after the call they hold the advanced counter and residual +/// offset, so passing the same objects to the next call continues the +/// keystream (TgCrypto-compatible, required by pyrogram and its forks). +/// `bytes` inputs are treated as stateless one-shot calls. +/// /// Raises: -/// ValueError: If `key`, `iv`, or `state` are invalid. +/// TypeError: If `data`, `key`, `iv`, or `state` are not `bytes` or +/// `bytearray`. +/// ValueError: If `data`, `key`, `iv`, or `state` have invalid lengths. /// OverflowError: If the requested output would exceed Python's `bytes` size limit. /// RuntimeError: If an unexpected internal error occurs. #[pyfunction] #[pyo3(signature = (data, key, iv, state))] fn ctr256_encrypt<'py>( py: Python<'py>, - data: &[u8], - key: &[u8], - iv: &[u8], - state: &[u8], + data: &Bound<'py, PyAny>, + key: &Bound<'py, PyAny>, + iv: &Bound<'py, PyAny>, + state: &Bound<'py, PyAny>, ) -> PyResult> { - let key_arr = copy_array::<32>(key, "Key")?; - let mut iv_arr = copy_array::<16>(iv, "IV")?; - let mut state_val = validate_ctr_state(state)?; - - let (bytes, _) = execute_zerocopy(py, data.len(), move |dest| { - tgcryptors_core::ctr256_encrypt_into(data, &key_arr, &mut iv_arr, &mut state_val, dest); + let data = BufferInput::from_any(data, "Data")?; + let key = BufferInput::from_any(key, "Key")?; + let iv = BufferInput::from_any(iv, "IV")?; + let state = BufferInput::from_any(state, "State")?; + + let key_arr = copy_array::<32>(key.as_bytes(), "Key")?; + let mut iv_arr = copy_array::<16>(iv.as_bytes(), "IV")?; + let mut state_val = validate_ctr_state(state.as_bytes())?; + let data_slice = data.as_bytes(); + let iv_source = iv.into_source(); + let state_source = state.into_source(); + + let (bytes, (next_iv, next_state)) = execute_zerocopy(py, data_slice.len(), move |dest| { + tgcryptors_core::ctr256_encrypt_into( + data_slice, + &key_arr, + &mut iv_arr, + &mut state_val, + dest, + ); + (iv_arr, state_val) })?; + + if let Some(source) = iv_source { + write_back(&source, &next_iv); + } + if let Some(source) = state_source { + write_back(&source, &[next_state]); + } + Ok(bytes) } /// Decrypt bytes with AES-256-CTR. /// -/// CTR is symmetric, so decryption delegates to `ctr256_encrypt`. +/// CTR is symmetric, so decryption delegates to `ctr256_encrypt`. See +/// [`ctr256_encrypt`] for the accepted argument types and the in-place +/// `bytearray` carry semantics. #[pyfunction] #[pyo3(signature = (data, key, iv, state))] fn ctr256_decrypt<'py>( py: Python<'py>, - data: &[u8], - key: &[u8], - iv: &[u8], - state: &[u8], + data: &Bound<'py, PyAny>, + key: &Bound<'py, PyAny>, + iv: &Bound<'py, PyAny>, + state: &Bound<'py, PyAny>, ) -> PyResult> { ctr256_encrypt(py, data, key, iv, state) } diff --git a/uv.lock b/uv.lock index 7c3bae4..bc11a5a 100644 --- a/uv.lock +++ b/uv.lock @@ -4,5 +4,5 @@ requires-python = ">=3.9, <3.15" [[package]] name = "tgcryptorust" -version = "1.3.0" +version = "1.3.1" source = { editable = "." }