Repository: Zapix/mtpylon (language: python)

Issue:
Add auth key for session manager
Need store session associated with auth_key. It stores without auth_key now that could be used for attack. 


Discussion / hints:

Fix patch:
diff --git a/mtpylon/middlewares/set_session_id.py b/mtpylon/middlewares/set_session_id.py
index 0fe7996..abb1fd3 100644
--- a/mtpylon/middlewares/set_session_id.py
+++ b/mtpylon/middlewares/set_session_id.py
@@ -5,7 +5,11 @@
 from aiohttp.web import Request
 
 from mtpylon.constants import SESSION_SUBJECT_RESOURCE_NAME
-from mtpylon.contextvars import income_message_var, session_id_var
+from mtpylon.contextvars import (
+    income_message_var,
+    session_id_var,
+    auth_key_var,
+)
 from mtpylon.messages import EncryptedMessage
 
 from .types import Handler
@@ -22,8 +26,15 @@ async def set_session_id(
     income_message = income_message_var.get()
 
     if isinstance(income_message, EncryptedMessage):
-        if not await session_subject.has_session(income_message.session_id):
-            await session_subject.create_session(income_message.session_id)
+        auth_key = auth_key_var.get()
+        if not await session_subject.has_session(
+            auth_key,
+            income_message.session_id
+        ):
+            await session_subject.create_session(
+                auth_key,
+                income_message.session_id
+            )
 
         logger.info(f'Set session id: {income_message.session_id}')
 
diff --git a/mtpylon/sessions/in_memory_session_storage.py b/mtpylon/sessions/in_memory_session_storage.py
index c1de05d..b4a03a7 100644
--- a/mtpylon/sessions/in_memory_session_storage.py
+++ b/mtpylon/sessions/in_memory_session_storage.py
@@ -1,28 +1,35 @@
 # -*- coding: utf-8 -*-
-from typing import Set
+from typing import Set, Tuple
 from asyncio import Lock
 
 from mtpylon.types import long
+from mtpylon.crypto import AuthKey
 from .session_storage_protocol import SessionStorageProtocol
 
 
+AuthKeySessionTuple = Tuple[AuthKey, long]
+
+
 class InMemorySessionStorage(SessionStorageProtocol):
     """
     Stores available session in session set
     """
 
     def __init__(self):
-        self._storage: Set[long] = set()
+        self._storage: Set[AuthKeySessionTuple] = set()
         self._lock = Lock()
 
-    async def create_session(self, session_id: long):
+    async def create_session(self, auth_key: AuthKey, session_id: long):
+        auth_key_session_tuple = (auth_key, session_id)
         async with self._lock:
-            self._storage.add(session_id)
+            self._storage.add(auth_key_session_tuple)
 
-    async def has_session(self, session_id: long) -> bool:
-        return session_id in self._storage
+    async def has_session(self, auth_key: AuthKey, session_id: long) -> bool:
+        auth_key_session_tuple = (auth_key, session_id)
+        return auth_key_session_tuple in self._storage
 
-    async def destroy_session(self, session_id: long):
-        if session_id in self._storage:
+    async def destroy_session(self, auth_key: AuthKey, session_id: long):
+        auth_key_session_tuple = (auth_key, session_id)
+        if auth_key_session_tuple in self._storage:
             async with self._lock:
-                self._storage.remove(session_id)
+                self._storage.remove(auth_key_session_tuple)
diff --git a/mtpylon/sessions/session_event.py b/mtpylon/sessions/session_event.py
index 3f63154..eadab12 100644
--- a/mtpylon/sessions/session_event.py
+++ b/mtpylon/sessions/session_event.py
@@ -3,6 +3,7 @@
 from dataclasses import dataclass
 
 from mtpylon.types import long
+from mtpylon.crypto import AuthKey
 
 
 EventType = Literal['created', 'destroyed']
@@ -12,3 +13,4 @@
 class SessionEvent:
     type: EventType
     session_id: long
+    auth_key: AuthKey
diff --git a/mtpylon/sessions/session_storage_protocol.py b/mtpylon/sessions/session_storage_protocol.py
index 6de98fc..1829ef8 100644
--- a/mtpylon/sessions/session_storage_protocol.py
+++ b/mtpylon/sessions/session_storage_protocol.py
@@ -1,24 +1,37 @@
 # -*- coding: utf-8 -*-
 from typing import Protocol
 
-from mtpylon import long
+from mtpylon.types import long
+from mtpylon.crypto import AuthKey
 
 
 class SessionStorageProtocol(Protocol):
 
-    async def create_session(self, session_id: long):  # pragma: nocover
+    async def create_session(
+        self,
+        auth_key: AuthKey,
+        session_id: long
+    ):  # pragma: nocover
         """
         Saves info about new session(session_id)
         """
         ...
 
-    async def has_session(self, session_id: long) -> bool:  # pragma: nocover
+    async def has_session(
+        self,
+        auth_key: AuthKey,
+        session_id: long
+    ) -> bool:  # pragma: nocover
         """
         Checks has session_id been stored or not
         """
         ...
 
-    async def destroy_session(self, session_id: long):  # pragma: nocover
+    async def destroy_session(
+        self,
+        auth_key: AuthKey,
+        session_id: long
+    ):  # pragma: nocover
         """
         Removes info about stored
         """
diff --git a/mtpylon/sessions/session_subject.py b/mtpylon/sessions/session_subject.py
index 9dae38b..552e03f 100644
--- a/mtpylon/sessions/session_subject.py
+++ b/mtpylon/sessions/session_subject.py
@@ -3,6 +3,7 @@
 from typing import Callable
 
 from mtpylon.types import long
+from mtpylon.crypto import AuthKey
 from .session_subject_protocol import SessionSubjectProtocol
 from .session_storage_protocol import SessionStorageProtocol
 from .session_oberver_protocol import SessionObserverProtocol
@@ -30,22 +31,24 @@ async def _notify(self, event: SessionEvent):
             for observer in self.observers
         ])
 
-    async def create_session(self, session_id: long):
-        await self.session_storage.create_session(session_id)
+    async def create_session(self, auth_key: AuthKey, session_id: long):
+        await self.session_storage.create_session(auth_key, session_id)
         event = SessionEvent(
             type='created',
-            session_id=session_id
+            session_id=session_id,
+            auth_key=auth_key
         )
         await self._notify(event)
 
-    async def has_session(self, session_id: long) -> bool:
-        return await self.session_storage.has_session(session_id)
+    async def has_session(self, auth_key: AuthKey, session_id: long) -> bool:
+        return await self.session_storage.has_session(auth_key, session_id)
 
-    async def destroy_session(self, session_id: long):
-        if await self.session_storage.has_session(session_id):
-            await self.session_storage.destroy_session(session_id)
+    async def destroy_session(self, auth_key: AuthKey, session_id: long):
+        if await self.session_storage.has_session(auth_key, session_id):
+            await self.session_storage.destroy_session(auth_key, session_id)
             event = SessionEvent(
                 type='destroyed',
-                session_id=session_id
+                session_id=session_id,
+                auth_key=auth_key
             )
             await self._notify(event)


Test patch:
diff --git a/tests/test_middlewares/test_session_id.py b/tests/test_middlewares/test_session_id.py
index 7f82b55..a3fc44b 100644
--- a/tests/test_middlewares/test_session_id.py
+++ b/tests/test_middlewares/test_session_id.py
@@ -4,10 +4,12 @@
 import pytest
 
 from mtpylon.types import long
+from mtpylon.crypto import AuthKey
 from mtpylon.constants import SESSION_SUBJECT_RESOURCE_NAME
 from mtpylon.contextvars import (
     income_message_var,
     session_id_var,
+    auth_key_var
 )
 from mtpylon.serialization import CallableFunc
 from mtpylon.messages import EncryptedMessage
@@ -21,8 +23,16 @@
 from tests.simpleschema import set_task
 
 
+@pytest.fixture
+def auth_key():
+    auth_key_data = 19173138450442901591212846731489158386637947003706705793697466005840359118325984130331040988945126053253996441948760332982664467588522441280332705742850565768210808401324268532424996018690635427668017180714026266980907736921933287244346526265760362799421836958089132713806536144155052702458589251084836839798895173633060038337086228492970997360899822298102887831495037141025890442863060204588167094562664112338589266018632963871940922732022252873979144345421549843044971414444544589457542139689588733663359139549339618779617218500603713731236381263744006515913607287705083142719869116507454233793160540351518647758028  # noqa
+    return AuthKey(auth_key_data)
+
+
 @pytest.mark.asyncio
-async def test_create_session_id():
+async def test_create_session_id(auth_key):
+    auth_key_var.set(auth_key)
+
     msg_id = long(0x51e57ac42770964a)
     session_id = long(123123)
     salt = long(234234)
@@ -61,7 +71,8 @@ async def test_create_session_id():
     awaited_arg = session_observer.update.await_args[0][0]
     expected_event = SessionEvent(
         type='created',
-        session_id=session_id
+        session_id=session_id,
+        auth_key=auth_key
     )
     assert awaited_arg == expected_event
 
@@ -69,7 +80,9 @@ async def test_create_session_id():
 
 
 @pytest.mark.asyncio
-async def test_session_id_created_before():
+async def test_session_id_created_before(auth_key):
+    auth_key_var.set(auth_key)
+
     msg_id = long(0x51e57ac42770964a)
     session_id = long(123123)
     salt = long(234234)
@@ -91,7 +104,7 @@ async def test_session_id_created_before():
 
     session_subject = SessionSubject(InMemorySessionStorage)
 
-    await session_subject.create_session(session_id=session_id)
+    await session_subject.create_session(auth_key, session_id)
 
     request = MagicMock()
     request.app = {
diff --git a/tests/test_sessions/test_in_memory_session_storage.py b/tests/test_sessions/test_in_memory_session_storage.py
index 6437dd0..734f0eb 100644
--- a/tests/test_sessions/test_in_memory_session_storage.py
+++ b/tests/test_sessions/test_in_memory_session_storage.py
@@ -2,6 +2,7 @@
 import pytest
 
 from mtpylon import long
+from mtpylon.crypto import AuthKey
 from mtpylon.sessions import InMemorySessionStorage
 
 
@@ -10,47 +11,52 @@ def session_storage():
     return InMemorySessionStorage()
 
 
+@pytest.fixture
+def auth_key():
+    auth_key_data = 19173138450442901591212846731489158386637947003706705793697466005840359118325984130331040988945126053253996441948760332982664467588522441280332705742850565768210808401324268532424996018690635427668017180714026266980907736921933287244346526265760362799421836958089132713806536144155052702458589251084836839798895173633060038337086228492970997360899822298102887831495037141025890442863060204588167094562664112338589266018632963871940922732022252873979144345421549843044971414444544589457542139689588733663359139549339618779617218500603713731236381263744006515913607287705083142719869116507454233793160540351518647758028  # noqa
+    return AuthKey(auth_key_data)
+
+
 @pytest.mark.asyncio
-async def test_create_session(session_storage):
+async def test_create_session(session_storage, auth_key):
     session_id = long(123123)
 
-    await session_storage.create_session(session_id)
+    await session_storage.create_session(auth_key, session_id)
 
     assert len(session_storage._storage) == 1
-    assert session_id in session_storage._storage
 
 
 @pytest.mark.asyncio
-async def test_has_not_got_session(session_storage):
+async def test_has_not_got_session(session_storage, auth_key):
     session_id = long(123123)
 
-    assert not await session_storage.has_session(session_id)
+    assert not await session_storage.has_session(auth_key, session_id)
 
 
 @pytest.mark.asyncio
-async def test_has_got_session(session_storage):
+async def test_has_got_session(session_storage, auth_key):
     session_id = long(123123)
 
-    await session_storage.create_session(session_id)
+    await session_storage.create_session(auth_key, session_id)
 
-    assert await session_storage.has_session(session_id)
+    assert await session_storage.has_session(auth_key, session_id)
 
 
 @pytest.mark.asyncio
-async def test_destroy_session(session_storage):
+async def test_destroy_session(session_storage, auth_key):
     session_id = long(123123)
 
-    await session_storage.create_session(session_id)
+    await session_storage.create_session(auth_key, session_id)
 
-    await session_storage.destroy_session(session_id)
+    await session_storage.destroy_session(auth_key, session_id)
 
-    assert not await session_storage.has_session(session_id)
+    assert not await session_storage.has_session(auth_key, session_id)
 
 
 @pytest.mark.asyncio
-async def test_destroy_session_no_id(session_storage):
+async def test_destroy_session_no_id(session_storage, auth_key):
     session_id = long(123123)
 
-    await session_storage.destroy_session(session_id)
+    await session_storage.destroy_session(auth_key, session_id)
 
-    assert not await session_storage.has_session(session_id)
+    assert not await session_storage.has_session(auth_key, session_id)
diff --git a/tests/test_sessions/test_session_subject.py b/tests/test_sessions/test_session_subject.py
index e5b81be..9bdfc85 100644
--- a/tests/test_sessions/test_session_subject.py
+++ b/tests/test_sessions/test_session_subject.py
@@ -3,6 +3,7 @@
 
 import pytest
 
+from mtpylon.crypto import AuthKey
 from mtpylon.types import long
 from mtpylon.sessions import (
     SessionEvent,
@@ -16,8 +17,14 @@ def session_subject():
     return SessionSubject(lambda: InMemorySessionStorage())
 
 
+@pytest.fixture
+def auth_key():
+    auth_key_data = 19173138450442901591212846731489158386637947003706705793697466005840359118325984130331040988945126053253996441948760332982664467588522441280332705742850565768210808401324268532424996018690635427668017180714026266980907736921933287244346526265760362799421836958089132713806536144155052702458589251084836839798895173633060038337086228492970997360899822298102887831495037141025890442863060204588167094562664112338589266018632963871940922732022252873979144345421549843044971414444544589457542139689588733663359139549339618779617218500603713731236381263744006515913607287705083142719869116507454233793160540351518647758028  # noqa
+    return AuthKey(auth_key_data)
+
+
 @pytest.mark.asyncio
-async def test_create_session(session_subject):
+async def test_create_session(session_subject, auth_key):
     observer1 = MagicMock(update=AsyncMock())
     observer2 = MagicMock(update=AsyncMock())
 
@@ -26,11 +33,12 @@ async def test_create_session(session_subject):
 
     session_id = long(123123)
 
-    await session_subject.create_session(session_id)
+    await session_subject.create_session(auth_key, session_id)
 
     created_event = SessionEvent(
         type='created',
         session_id=session_id,
+        auth_key=auth_key
     )
 
     observer1.update.assert_awaited()
@@ -41,20 +49,21 @@ async def test_create_session(session_subject):
 
 
 @pytest.mark.asyncio
-async def test_destroy_session(session_subject):
+async def test_destroy_session(session_subject, auth_key):
     session_id = long(123123)
-    await session_subject.create_session(session_id)
+    await session_subject.create_session(auth_key, session_id)
 
     observer1 = MagicMock(update=AsyncMock())
     observer2 = MagicMock(update=AsyncMock())
     session_subject.subscribe(observer1)
     session_subject.subscribe(observer2)
 
-    await session_subject.destroy_session(session_id)
+    await session_subject.destroy_session(auth_key, session_id)
 
     destroyed_event = SessionEvent(
         type='destroyed',
         session_id=session_id,
+        auth_key=auth_key
     )
 
     observer1.update.assert_awaited()
@@ -65,19 +74,19 @@ async def test_destroy_session(session_subject):
 
 
 @pytest.mark.asyncio
-async def test_destory_session_no_session(session_subject):
+async def test_destory_session_no_session(session_subject, auth_key):
     session_id = long(123123)
 
     observer1 = MagicMock(update=AsyncMock())
 
     session_subject.subscribe(observer1)
-    await session_subject.destroy_session(session_id)
+    await session_subject.destroy_session(auth_key, session_id)
 
     observer1.assert_not_called()
 
 
 @pytest.mark.asyncio
-async def test_unsubscribe_observer(session_subject):
+async def test_unsubscribe_observer(session_subject, auth_key):
     session1_id = long(123123)
     session2_id = long(234234)
 
@@ -87,20 +96,20 @@ async def test_unsubscribe_observer(session_subject):
     session_subject.subscribe(observer1)
     session_subject.subscribe(observer2)
 
-    await session_subject.create_session(session1_id)
+    await session_subject.create_session(auth_key, session1_id)
 
     session_subject.unsubscribe(observer1)
 
-    await session_subject.create_session(session2_id)
+    await session_subject.create_session(auth_key, session2_id)
 
     assert observer1.update.await_count == 1
     assert observer2.update.await_count == 2
 
 
 @pytest.mark.asyncio
-async def test_has_got_session(session_subject):
+async def test_has_got_session(session_subject, auth_key):
     session_id = long(123123)
 
-    await session_subject.create_session(session_id)
+    await session_subject.create_session(auth_key, session_id)
 
-    assert await session_subject.has_session(session_id)
+    assert await session_subject.has_session(auth_key, session_id)