fix: accept bytearray in ctr256_encrypt/decrypt (pyrogram compatibility)

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"

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.

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"

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,14 +169,30 @@ 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__)
         self.assertIn("Encrypt or decrypt the next chunk", tgcrypto.Ctr256.update.__doc__)
         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()
 

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<u8> {
     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<u8>,
+        source: Bound<'py, PyByteArray>,
+    },
+}
+
+impl<'py> BufferInput<'py> {
+    fn from_any(ob: &Bound<'py, PyAny>, label: &str) -> PyResult<Self> {
+        if let Ok(bytes) = ob.cast::<PyBytes>() {
+            return Ok(BufferInput::Bytes(bytes.clone()));
+        }
+        if let Ok(array) = ob.cast::<PyByteArray>() {
+            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<Bound<'py, PyByteArray>> {
+        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<Bound<'py, PyBytes>> {
-    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<Bound<'py, PyBytes>> {
     ctr256_encrypt(py, data, key, iv, state)
 }