From a0bf5c7579c2b0021ed219b9c4a3782694d0313a Mon Sep 17 00:00:00 2001
From: Gabriel Huber <mail@gabrielhuber.at>
Date: Wed, 25 Nov 2020 00:35:29 +0100
Subject: [PATCH] Implement request padding to mitigate amplification attacks

Proposed in "RFC: Changes to the A2S_INFO protocol"
https://steamcommunity.com/discussions/forum/14/2989789048633291344/
---
 README.md        | 4 ++++
 a2s/a2sasync.py  | 4 ++++
 a2s/a2sstream.py | 4 ++++
 setup.py         | 2 +-
 4 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index b814d8c..17fbff2 100644
--- a/README.md
+++ b/README.md
@@ -5,6 +5,10 @@ Impliments [Valve's Server Query Protocol](https://developer.valvesoftware.com/w
 Rewrite of the [python-valve](https://github.com/serverstf/python-valve) module.
 Supports both synchronous and asyncronous applications.
 
+**This branch implements the 1200 Bytes padding proposed by Valve in
+[a recent forum thread](https://steamcommunity.com/discussions/forum/14/2989789048633291344/).
+Only use this for testing as not all servers are compatible yet!**
+
 Official demo application: [Sourcequery](https://sourcequery.yepoleb.at)
 
 ## Requirements
diff --git a/a2s/a2sasync.py b/a2s/a2sasync.py
index 8eb3180..c5d7662 100644
--- a/a2s/a2sasync.py
+++ b/a2s/a2sasync.py
@@ -8,6 +8,7 @@ from a2s.a2sfragment import decode_fragment
 
 HEADER_SIMPLE = b"\xFF\xFF\xFF\xFF"
 HEADER_MULTI = b"\xFE\xFF\xFF\xFF"
+MIN_PACKET_SIZE = 1200
 
 logger = logging.getLogger("a2s")
 
@@ -71,6 +72,9 @@ class A2SStreamAsync:
 
     def send(self, payload):
         packet = HEADER_SIMPLE + payload
+        padding_count = MIN_PACKET_SIZE - len(packet)
+        if padding_count > 0:
+            packet += b"\x00" * padding_count
         self.transport.sendto(packet)
 
     async def recv(self):
diff --git a/a2s/a2sstream.py b/a2s/a2sstream.py
index 5d8c593..5ddf9b9 100644
--- a/a2s/a2sstream.py
+++ b/a2s/a2sstream.py
@@ -8,6 +8,7 @@ from a2s.a2sfragment import decode_fragment
 
 HEADER_SIMPLE = b"\xFF\xFF\xFF\xFF"
 HEADER_MULTI = b"\xFE\xFF\xFF\xFF"
+MIN_PACKET_SIZE = 1200
 
 logger = logging.getLogger("a2s")
 
@@ -22,6 +23,9 @@ class A2SStream:
 
     def send(self, data):
         packet = HEADER_SIMPLE + data
+        padding_count = MIN_PACKET_SIZE - len(packet)
+        if padding_count > 0:
+            packet += b"\x00" * padding_count
         self._socket.sendto(packet, self.address)
 
     def recv(self):
diff --git a/setup.py b/setup.py
index d5380fc..7fb2fb0 100644
--- a/setup.py
+++ b/setup.py
@@ -7,7 +7,7 @@ with open("README.md", "r") as readme:
 
 setuptools.setup(
     name="python-a2s",
-    version="1.2.1",
+    version="1.3.0",
     author="Gabriel Huber",
     author_email="mail@gabrielhuber.at",
     description="Query Source and GoldSource servers for name, map, players and more.",