From 43834494397aa0c9bea56ecb9b3295703e1560ef Mon Sep 17 00:00:00 2001
From: Vladislav Yarmak <vladislav@vm-0.com>
Date: Thu, 21 Feb 2019 20:26:35 +0200
Subject: [PATCH] restrict MTA-STS policy server response size to avoid DoS

---
 postfix_mta_sts_resolver/resolver.py | 20 +++++++++++++++++++-
 setup.py                             |  2 +-
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/postfix_mta_sts_resolver/resolver.py b/postfix_mta_sts_resolver/resolver.py
index 3e34d74..8dd1142 100644
--- a/postfix_mta_sts_resolver/resolver.py
+++ b/postfix_mta_sts_resolver/resolver.py
@@ -2,10 +2,15 @@
 import aiodns
 import aiohttp
 import enum
+from io import BytesIO
 from . import defaults
 from .utils import parse_mta_sts_record, parse_mta_sts_policy, is_plaintext
 
 
+HARD_RESP_LIMIT = 64 * 1024
+CHUNK = 4096
+
+
 class BadSTSPolicy(Exception):
     pass
 
@@ -92,7 +97,20 @@ async def resolve(self, domain, last_known_id=None):
                         raise BadSTSPolicy()
                     if not is_plaintext(resp.headers.get('Content-Type', '')):
                         raise BadSTSPolicy()
-                    policy_text = await resp.text()
+                    if (int(resp.headers.get('Content-Length', '0')) > 
+                            HARD_RESP_LIMIT):
+                        raise BadSTSPolicy()
+                    policy_file = BytesIO()
+                    while policy_file.tell() <= HARD_RESP_LIMIT:
+                        chunk = await resp.content.read(CHUNK)
+                        if not chunk:
+                            break
+                        policy_file.write(chunk)
+                    else:
+                        raise BadSTSPolicy()
+                    charset = (resp.charset if resp.charset is not None
+                               else 'ascii')
+                    policy_text = policy_file.getvalue().decode(charset)
         except:
             return STSFetchResult.FETCH_ERROR, None
 
diff --git a/setup.py b/setup.py
index e3fe725..1e21f95 100644
--- a/setup.py
+++ b/setup.py
@@ -6,7 +6,7 @@
     long_description = f.read()
 
 setup(name='postfix_mta_sts_resolver',
-      version='0.2.4',
+      version='0.2.5',
       description='Daemon which provides TLS client policy for Postfix via socketmap, according to domain MTA-STS policy',
       url='https://github.com/Snawoot/postfix-mta-sts-resolver',
       author='Vladislav Yarmak',