From fe9ccb68a878025b971590442b085202a1fe2a1a Mon Sep 17 00:00:00 2001 From: Luna Date: Wed, 30 Jun 2021 18:28:31 -0300 Subject: [PATCH 1/8] add gateway schema module --- litecord/gateway/schemas.py | 144 ++++++++++++++++++++++++++++++++++ litecord/gateway/websocket.py | 12 +-- litecord/schemas.py | 73 ----------------- 3 files changed, 148 insertions(+), 81 deletions(-) create mode 100644 litecord/gateway/schemas.py diff --git a/litecord/gateway/schemas.py b/litecord/gateway/schemas.py new file mode 100644 index 0000000..5b823d5 --- /dev/null +++ b/litecord/gateway/schemas.py @@ -0,0 +1,144 @@ +""" + +Litecord +Copyright (C) 2018-2019 Luna Mendes + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU General Public License as published by +the Free Software Foundation, version 3 of the License. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU General Public License for more details. + +You should have received a copy of the GNU General Public License +along with this program. If not, see . + +""" + +from typing import Dict + +from logbook import Logger + + +from litecord.gateway.errors import DecodeError +from litecord.schemas import LitecordValidator + +log = Logger(__name__) + + +def validate( + reqjson: Dict, + schema: Dict, +) -> Dict: + validator = LitecordValidator(schema) + + try: + valid = validator.validate(reqjson) + except Exception: + log.exception("Error while validating") + raise DecodeError(f"Error while validating: {reqjson}") + + if not valid: + errs = validator.errors + log.warning("Error validating doc {!r}: {!r}", reqjson, errs) + raise DecodeError(f"Error validating message : {errs!r}") + + return validator.document + + +BASE = { + "op": {"type": "number", "required": True}, + "s": {"type": "number", "required": False}, +} + +IDENTIFY_SCHEMA = { + **BASE, + **{ + "d": { + "type": "dict", + "schema": { + "token": {"type": "string", "required": True}, + "compress": {"type": "boolean", "required": False}, + "large_threshold": {"type": "number", "required": False}, + "shard": {"type": "list", "required": False}, + "presence": {"type": "dict", "required": False}, + }, + } + }, +} + + +GW_ACTIVITY = { + "name": {"type": "string", "required": True}, + "type": {"type": "activity_type", "required": True}, + "url": {"type": "string", "required": False, "nullable": True}, + "timestamps": { + "type": "dict", + "required": False, + "schema": { + "start": {"type": "number", "required": False}, + "end": {"type": "number", "required": False}, + }, + }, + "application_id": {"type": "snowflake", "required": False, "nullable": False}, + "details": {"type": "string", "required": False, "nullable": True}, + "state": {"type": "string", "required": False, "nullable": True}, + "party": { + "type": "dict", + "required": False, + "schema": { + "id": {"type": "snowflake", "required": False}, + "size": {"type": "list", "required": False}, + }, + }, + "assets": { + "type": "dict", + "required": False, + "schema": { + "large_image": {"type": "snowflake", "required": False}, + "large_text": {"type": "string", "required": False}, + "small_image": {"type": "snowflake", "required": False}, + "small_text": {"type": "string", "required": False}, + }, + }, + "secrets": { + "type": "dict", + "required": False, + "schema": { + "join": {"type": "string", "required": False}, + "spectate": {"type": "string", "required": False}, + "match": {"type": "string", "required": False}, + }, + }, + "instance": {"type": "boolean", "required": False}, + "flags": {"type": "number", "required": False}, + "emoji": { + "type": "dict", + "required": False, + "nullable": True, + "schema": { + "animated": {"type": "boolean", "required": False, "default": False}, + "id": {"coerce": int, "nullable": True, "default": None}, + "name": {"type": "string", "required": True}, + }, + }, +} + +GW_STATUS_UPDATE = { + "status": {"type": "status_external", "required": False, "default": "online"}, + "activities": { + "type": "list", + "required": False, + "schema": {"type": "dict", "schema": GW_ACTIVITY}, + }, + "afk": {"type": "boolean", "required": False}, + "since": {"type": "number", "required": False, "nullable": True}, + "game": { + "type": "dict", + "required": False, + "nullable": True, + "schema": GW_ACTIVITY, + }, +} diff --git a/litecord/gateway/websocket.py b/litecord/gateway/websocket.py index 681cc56..e91e77b 100644 --- a/litecord/gateway/websocket.py +++ b/litecord/gateway/websocket.py @@ -32,7 +32,6 @@ from quart import current_app as app from litecord.auth import raw_token_check from litecord.enums import RelationshipType, ChannelType, ActivityType -from litecord.schemas import validate, GW_STATUS_UPDATE from litecord.utils import ( task_wrapper, yield_chunks, @@ -59,6 +58,7 @@ from litecord.gateway.encoding import encode_json, decode_json, encode_etf, deco from litecord.gateway.utils import WebsocketFileHandler from litecord.pubsub.guild import GuildFlags from litecord.pubsub.channel import ChannelFlags +from litecord.gateway.schemas import validate, IDENTIFY_SCHEMA, GW_STATUS_UPDATE from litecord.storage import int_ @@ -651,13 +651,9 @@ class GatewayWebsocket: async def handle_2(self, payload: Dict[str, Any]): """Handle the OP 2 Identify packet.""" - try: - data = payload["d"] - token = data["token"] - except KeyError: - raise DecodeError("Invalid identify parameters") - - # TODO proper validation of this payload + payload = validate(payload, IDENTIFY_SCHEMA) + data = payload["d"] + token = data["token"] compress = data.get("compress", False) large = data.get("large_threshold", 50) diff --git a/litecord/schemas.py b/litecord/schemas.py index cbb4781..3ba1f4c 100644 --- a/litecord/schemas.py +++ b/litecord/schemas.py @@ -418,79 +418,6 @@ MESSAGE_CREATE = { } -GW_ACTIVITY = { - "name": {"type": "string", "required": True}, - "type": {"type": "activity_type", "required": True}, - "url": {"type": "string", "required": False, "nullable": True}, - "timestamps": { - "type": "dict", - "required": False, - "schema": { - "start": {"type": "number", "required": False}, - "end": {"type": "number", "required": False}, - }, - }, - "application_id": {"type": "snowflake", "required": False, "nullable": False}, - "details": {"type": "string", "required": False, "nullable": True}, - "state": {"type": "string", "required": False, "nullable": True}, - "party": { - "type": "dict", - "required": False, - "schema": { - "id": {"type": "snowflake", "required": False}, - "size": {"type": "list", "required": False}, - }, - }, - "assets": { - "type": "dict", - "required": False, - "schema": { - "large_image": {"type": "snowflake", "required": False}, - "large_text": {"type": "string", "required": False}, - "small_image": {"type": "snowflake", "required": False}, - "small_text": {"type": "string", "required": False}, - }, - }, - "secrets": { - "type": "dict", - "required": False, - "schema": { - "join": {"type": "string", "required": False}, - "spectate": {"type": "string", "required": False}, - "match": {"type": "string", "required": False}, - }, - }, - "instance": {"type": "boolean", "required": False}, - "flags": {"type": "number", "required": False}, - "emoji": { - "type": "dict", - "required": False, - "nullable": True, - "schema": { - "animated": {"type": "boolean", "required": False, "default": False}, - "id": {"coerce": int, "nullable": True, "default": None}, - "name": {"type": "string", "required": True}, - }, - }, -} - -GW_STATUS_UPDATE = { - "status": {"type": "status_external", "required": False, "default": "online"}, - "activities": { - "type": "list", - "required": False, - "schema": {"type": "dict", "schema": GW_ACTIVITY}, - }, - "afk": {"type": "boolean", "required": False}, - "since": {"type": "number", "required": False, "nullable": True}, - "game": { - "type": "dict", - "required": False, - "nullable": True, - "schema": GW_ACTIVITY, - }, -} - INVITE = { # max_age in seconds # 0 for infinite From ba62403674c548c34e3b8293563434e7ae1f6e11 Mon Sep 17 00:00:00 2001 From: Luna Date: Wed, 30 Jun 2021 18:36:14 -0300 Subject: [PATCH 2/8] add schemas for 'resume' and 'request guild' --- litecord/gateway/schemas.py | 30 ++++++++++++++++++++++++++++++ litecord/gateway/websocket.py | 17 ++++++++++------- 2 files changed, 40 insertions(+), 7 deletions(-) diff --git a/litecord/gateway/schemas.py b/litecord/gateway/schemas.py index 5b823d5..08260b2 100644 --- a/litecord/gateway/schemas.py +++ b/litecord/gateway/schemas.py @@ -69,6 +69,36 @@ IDENTIFY_SCHEMA = { }, } +RESUME_SCHEMA = { + **BASE, + **{ + "d": { + "type": "dict", + "schema": { + "token": {"type": "string", "required": True}, + "session_id": {"type": "string", "required": True}, + "seq": {"type": "number", "required": True}, + }, + } + }, +} + +REQ_GUILD_SCHEMA = { + **BASE, + **{ + "d": { + "type": "dict", + "schema": { + "guild_id": {"type": "string", "required": True}, + "user_ids": {"type": "list", "required": False}, + "query": {"type": "string", "required": False}, + "limit": {"type": "number", "required": False}, + "presences": {"type": "bool", "required": False}, + }, + } + }, +} + GW_ACTIVITY = { "name": {"type": "string", "required": True}, diff --git a/litecord/gateway/websocket.py b/litecord/gateway/websocket.py index e91e77b..091493c 100644 --- a/litecord/gateway/websocket.py +++ b/litecord/gateway/websocket.py @@ -58,7 +58,13 @@ from litecord.gateway.encoding import encode_json, decode_json, encode_etf, deco from litecord.gateway.utils import WebsocketFileHandler from litecord.pubsub.guild import GuildFlags from litecord.pubsub.channel import ChannelFlags -from litecord.gateway.schemas import validate, IDENTIFY_SCHEMA, GW_STATUS_UPDATE +from litecord.gateway.schemas import ( + validate, + IDENTIFY_SCHEMA, + GW_STATUS_UPDATE, + RESUME_SCHEMA, + REQ_GUILD_SCHEMA, +) from litecord.storage import int_ @@ -836,12 +842,9 @@ class GatewayWebsocket: async def handle_6(self, payload: Dict[str, Any]): """Handle OP 6 Resume.""" + payload = validate(payload, RESUME_SCHEMA) data = payload["d"] - - try: - token, sess_id, seq = data["token"], data["session_id"], data["seq"] - except KeyError: - raise DecodeError("Invalid resume payload") + token, sess_id, seq = data["token"], data["session_id"], data["seq"] try: user_id = await raw_token_check(token, self.app.db) @@ -911,6 +914,7 @@ class GatewayWebsocket: async def handle_8(self, payload: Dict): """Handle OP 8 Request Guild Members.""" + payload = validate(payload, REQ_GUILD_SCHEMA) data = payload["d"] gids = data["guild_id"] @@ -949,7 +953,6 @@ class GatewayWebsocket: async def handle_12(self, payload: Dict[str, Any]): """Handle OP 12 Guild Sync.""" data = payload["d"] - gids = await self.user_storage.get_user_guilds(self.state.user_id) for guild_id in data: From e5f8fe7243dfd41289cb937f62ad88e1cc85e0bb Mon Sep 17 00:00:00 2001 From: Luna Date: Wed, 30 Jun 2021 18:42:20 -0300 Subject: [PATCH 3/8] add test for gateway message validation error --- tests/test_websocket.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_websocket.py b/tests/test_websocket.py index ed49838..a197cc2 100644 --- a/tests/test_websocket.py +++ b/tests/test_websocket.py @@ -103,6 +103,25 @@ async def test_ready(test_cli_user): await _close(conn) +@pytest.mark.asyncio +async def test_broken_identify(test_cli_user): + conn = await gw_start(test_cli_user.cli) + + # get the hello frame but ignore it + await _json(conn) + + await _json_send(conn, {"op": OP.IDENTIFY, "d": {"token": True}}) + + # try to get a ready + try: + await _json(conn) + raise AssertionError("Received a JSON message but expected close") + except websockets.ConnectionClosed as exc: + assert exc.code == 4002 + finally: + await _close(conn) + + @pytest.mark.asyncio async def test_ready_fields(test_cli_user): conn = await gw_start(test_cli_user.cli) From fb817e129abad5b29b5542bc0f83bbefff643a2c Mon Sep 17 00:00:00 2001 From: Luna Date: Wed, 30 Jun 2021 22:33:26 -0300 Subject: [PATCH 4/8] ci: debug venv location --- .gitlab-ci.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 263e347..4173da2 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -20,6 +20,7 @@ tests: script: - ls - cp config.ci.py config.py + - pipenv --venv - pipenv run ./manage.py migrate - pipenv run black --check litecord run.py tests manage - pipenv run pyflakes run.py litecord/ From 87f77e65f12b51ceab315f170e1c7254a2294125 Mon Sep 17 00:00:00 2001 From: Luna Date: Wed, 30 Jun 2021 22:33:38 -0300 Subject: [PATCH 5/8] test_guilds: hotfix for missing app --- tests/test_admin_api/test_guilds.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_admin_api/test_guilds.py b/tests/test_admin_api/test_guilds.py index 3ff45fc..9bc2ca3 100644 --- a/tests/test_admin_api/test_guilds.py +++ b/tests/test_admin_api/test_guilds.py @@ -28,9 +28,10 @@ from litecord.errors import GuildNotFound async def _create_guild(test_cli_staff, *, region=None) -> dict: genned_name = secrets.token_hex(6) - resp = await test_cli_staff.post( - "/api/v6/guilds", json={"name": genned_name, "region": region} - ) + async with test_cli_staff.app.app_context(): + resp = await test_cli_staff.post( + "/api/v6/guilds", json={"name": genned_name, "region": region} + ) assert resp.status_code == 200 rjson = await resp.json From e01d1b8465461daaa2517f1ae33a1bc14953d099 Mon Sep 17 00:00:00 2001 From: Luna Date: Sat, 3 Jul 2021 00:00:13 -0300 Subject: [PATCH 6/8] ci: use flake8 instead of pyflakes --- .gitlab-ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 4173da2..44bfdf8 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -23,5 +23,5 @@ tests: - pipenv --venv - pipenv run ./manage.py migrate - pipenv run black --check litecord run.py tests manage - - pipenv run pyflakes run.py litecord/ + - pipenv run flake8 litecord run.py tests manage - pipenv run pytest tests From 173a2330fe984cdba4d894050c00eb60f57c0273 Mon Sep 17 00:00:00 2001 From: Luna Date: Sat, 3 Jul 2021 00:00:47 -0300 Subject: [PATCH 7/8] add validation for guild sync --- litecord/gateway/schemas.py | 10 ++++++++++ litecord/gateway/websocket.py | 2 ++ 2 files changed, 12 insertions(+) diff --git a/litecord/gateway/schemas.py b/litecord/gateway/schemas.py index 08260b2..4928a37 100644 --- a/litecord/gateway/schemas.py +++ b/litecord/gateway/schemas.py @@ -99,6 +99,16 @@ REQ_GUILD_SCHEMA = { }, } +GUILD_SYNC_SCHEMA = { + **BASE, + **{ + "d": { + "type": "list", + "schema": {"type": "snowflake"}, + } + }, +} + GW_ACTIVITY = { "name": {"type": "string", "required": True}, diff --git a/litecord/gateway/websocket.py b/litecord/gateway/websocket.py index 091493c..8e58c28 100644 --- a/litecord/gateway/websocket.py +++ b/litecord/gateway/websocket.py @@ -64,6 +64,7 @@ from litecord.gateway.schemas import ( GW_STATUS_UPDATE, RESUME_SCHEMA, REQ_GUILD_SCHEMA, + GUILD_SYNC_SCHEMA, ) from litecord.storage import int_ @@ -952,6 +953,7 @@ class GatewayWebsocket: async def handle_12(self, payload: Dict[str, Any]): """Handle OP 12 Guild Sync.""" + payload = validate(payload, GUILD_SYNC_SCHEMA) data = payload["d"] gids = await self.user_storage.get_user_guilds(self.state.user_id) From 9293803e11bb25957069a5f92c2ebf5087103a98 Mon Sep 17 00:00:00 2001 From: Luna Date: Fri, 9 Jul 2021 18:50:32 -0300 Subject: [PATCH 8/8] tests: attempt to fix flakey test --- tests/test_admin_api/test_guilds.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_admin_api/test_guilds.py b/tests/test_admin_api/test_guilds.py index 9bc2ca3..4d628ca 100644 --- a/tests/test_admin_api/test_guilds.py +++ b/tests/test_admin_api/test_guilds.py @@ -63,13 +63,13 @@ async def _delete_guild(test_cli, guild_id: int): @pytest.mark.asyncio async def test_guild_fetch(test_cli_staff): """Test the creation and fetching of a guild via the Admin API.""" - rjson = await _create_guild(test_cli_staff) - guild_id = rjson["id"] - - try: - await _fetch_guild(test_cli_staff, guild_id) - finally: - await _delete_guild(test_cli_staff, int(guild_id)) + async with test_cli_staff.app.app_context(): + rjson = await _create_guild(test_cli_staff) + guild_id = rjson["id"] + try: + await _fetch_guild(test_cli_staff, guild_id) + finally: + await _delete_guild(test_cli_staff, int(guild_id)) @pytest.mark.asyncio