From ba62403674c548c34e3b8293563434e7ae1f6e11 Mon Sep 17 00:00:00 2001
From: Luna <git@l4.pm>
Date: Wed, 30 Jun 2021 18:36:14 -0300
Subject: [PATCH] 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: