diff --git a/telethon/tl/tlobject.py b/telethon/tl/tlobject.py index 451fad17..cd563857 100644 --- a/telethon/tl/tlobject.py +++ b/telethon/tl/tlobject.py @@ -186,6 +186,19 @@ class TLObject: return json.dumps(d, default=default, **kwargs) def __bytes__(self): + try: + return self._bytes() + except AttributeError: + # If a type is wrong (e.g. expected `TLObject` but `int` was + # provided) it will try to access `._bytes()`, which will fail + # with `AttributeError`. This occurs in fact because the type + # was wrong, so raise the correct error type. + raise TypeError('a TLObject was expected but found something else') + + # Custom objects will call `(...)._bytes()` and not `bytes(...)` so that + # if the wrong type is used (e.g. `int`) we won't try allocating a huge + # amount of data, which would cause a `MemoryError`. + def _bytes(self): raise NotImplementedError @classmethod diff --git a/telethon_generator/generators/tlobject.py b/telethon_generator/generators/tlobject.py index cd2b99d7..9ef23861 100644 --- a/telethon_generator/generators/tlobject.py +++ b/telethon_generator/generators/tlobject.py @@ -330,7 +330,7 @@ def _write_to_dict(tlobject, builder): def _write_to_bytes(tlobject, builder): - builder.writeln('def __bytes__(self):') + builder.writeln('def _bytes(self):') # Some objects require more than one flag parameter to be set # at the same time. In this case, add an assertion. @@ -509,7 +509,7 @@ def _write_arg_to_bytes(builder, arg, args, name=None): else: # Else it may be a custom type - builder.write('bytes({})', name) + builder.write('{}._bytes()', name) # If the type is not boxed (i.e. starts with lowercase) we should # not serialize the constructor ID (so remove its first 4 bytes). diff --git a/tests/telethon/test_utils.py b/tests/telethon/test_utils.py index 89391e36..b0cfb33f 100644 --- a/tests/telethon/test_utils.py +++ b/tests/telethon/test_utils.py @@ -1,6 +1,8 @@ import io import pathlib +import pytest + from telethon import utils from telethon.tl.types import ( MessageMediaGame, Game, PhotoEmpty diff --git a/tests/telethon/tl/__init__.py b/tests/telethon/tl/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/telethon/tl/test_serialization.py b/tests/telethon/tl/test_serialization.py new file mode 100644 index 00000000..7dbb067d --- /dev/null +++ b/tests/telethon/tl/test_serialization.py @@ -0,0 +1,13 @@ +import pytest + +from telethon.tl import types, functions + + +def test_nested_invalid_serialization(): + large_long = 2**62 + request = functions.account.SetPrivacyRequest( + key=types.InputPrivacyKeyChatInvite(), + rules=[types.InputPrivacyValueDisallowUsers(users=[large_long])] + ) + with pytest.raises(TypeError): + bytes(request)