Repository: Zapix/mtpylon (language: python)

Issue:
Set server salt into context
Set server salt value to current context


Discussion / hints:

Fix patch:
diff --git a/mtpylon/contextvars.py b/mtpylon/contextvars.py
index 2cf665a..5d31a2e 100644
--- a/mtpylon/contextvars.py
+++ b/mtpylon/contextvars.py
@@ -4,7 +4,7 @@
 """
 from contextvars import ContextVar
 
-from .types import int128, int256
+from .types import int128, int256, long
 from .crypto import AuthKey
 from .messages.types import MtprotoMessage
 
@@ -33,10 +33,14 @@
 dh_prime_var: ContextVar[int] = ContextVar('dh_prime')
 
 
-"Stores auth key for current context"
+"""Stores auth key for current context"""
 auth_key_var: ContextVar[AuthKey] = ContextVar('auth_key')
 
 
+"""Stores server salt for current context"""
+server_salt_var: ContextVar[long] = ContextVar('server_salt_var')
+
+
 """
 Stores message that has been received
 """
diff --git a/mtpylon/message_handler/strategies/handle_rpc_query_message.py b/mtpylon/message_handler/strategies/handle_rpc_query_message.py
index 6c7b83d..a2e945a 100644
--- a/mtpylon/message_handler/strategies/handle_rpc_query_message.py
+++ b/mtpylon/message_handler/strategies/handle_rpc_query_message.py
@@ -7,7 +7,10 @@
 
 from mtpylon.exceptions import RpcCallError
 from mtpylon.messages import MtprotoMessage, Message
-from mtpylon.contextvars import income_message_var
+from mtpylon.contextvars import (
+    income_message_var,
+    server_salt_var,
+)
 from mtpylon.serialization import CallableFunc
 from mtpylon.middlewares import apply_middleware, MiddleWareFunc, Handler
 from mtpylon.message_sender import MessageSender
@@ -76,9 +79,11 @@ async def run_rpc_query(
 
     result = await handler(request, **value.params)
 
+    server_salt = server_salt_var.get()
+
     await sender.send_encrypted_message(
         request,
-        message.salt,
+        server_salt,
         message.session_id,
         result,
         True
diff --git a/mtpylon/middlewares/set_server_salt.py b/mtpylon/middlewares/set_server_salt.py
index 759f8bf..b5392f0 100644
--- a/mtpylon/middlewares/set_server_salt.py
+++ b/mtpylon/middlewares/set_server_salt.py
@@ -6,7 +6,11 @@
 from mtpylon.messages import Message
 from mtpylon.salts import ServerSaltManagerProtocol, Salt
 from mtpylon.service_schema.constructors import BadServerSalt
-from mtpylon.contextvars import auth_key_var, income_message_var
+from mtpylon.contextvars import (
+    auth_key_var,
+    income_message_var,
+    server_salt_var,
+)
 from .types import Handler
 
 
@@ -39,4 +43,6 @@ async def set_server_salt(
                 new_server_salt=new_salt.salt
             )
 
+        server_salt_var.set(message.salt)
+
     return await handler(request, **params)


Test patch:
diff --git a/tests/test_message_handler/test_strategies/test_handle_rpc_query_message.py b/tests/test_message_handler/test_strategies/test_handle_rpc_query_message.py
index a0a418f..046b83c 100644
--- a/tests/test_message_handler/test_strategies/test_handle_rpc_query_message.py
+++ b/tests/test_message_handler/test_strategies/test_handle_rpc_query_message.py
@@ -11,6 +11,7 @@
     run_rpc_query
 )
 from mtpylon.service_schema.constructors import RpcResult, RpcError
+from mtpylon.contextvars import server_salt_var
 
 from tests.simpleschema import get_task, set_task, Task
 
@@ -62,11 +63,17 @@ async def test_run_rpc_query_success():
         )
     )
 
+    server_salt_var.set(server_salt)
+
     await run_rpc_query([], sender, request, message)
 
     sender.send_encrypted_message.assert_awaited()
 
     args = sender.send_encrypted_message.await_args[0]
+
+    server_salt_encrypt = args[1]
+    assert server_salt_encrypt == server_salt
+
     rpc_result = args[3]
 
     assert isinstance(rpc_result, RpcResult)
@@ -94,6 +101,8 @@ async def test_run_rpc_query_error():
         )
     )
 
+    server_salt_var.set(server_salt)
+
     await run_rpc_query([], sender, request, message)
 
     sender.send_encrypted_message.assert_awaited()
@@ -126,6 +135,8 @@ async def test_run_rpc_unexpected_error():
         )
     )
 
+    server_salt_var.set(server_salt)
+
     await run_rpc_query([], sender, request, message)
 
     sender.send_encrypted_message.assert_awaited()
diff --git a/tests/test_middlewares/test_set_server_salt.py b/tests/test_middlewares/test_set_server_salt.py
index 180f29b..145a325 100644
--- a/tests/test_middlewares/test_set_server_salt.py
+++ b/tests/test_middlewares/test_set_server_salt.py
@@ -10,7 +10,11 @@
 from mtpylon.middlewares.set_server_salt import set_server_salt
 from mtpylon.crypto import AuthKey
 from mtpylon.salts import ServerSaltManager, Salt
-from mtpylon.contextvars import income_message_var, auth_key_var
+from mtpylon.contextvars import (
+    income_message_var,
+    auth_key_var,
+    server_salt_var
+)
 from mtpylon.service_schema.constructors import BadServerSalt
 
 from tests.simpleschema import set_task
@@ -56,6 +60,7 @@ async def test_server_salt_exist():
     await set_server_salt(handler, request, content='hello_world')
 
     handler.assert_awaited()
+    assert server_salt_var.get() == server_salt.salt
 
 
 @pytest.mark.asyncio