From f5e7751835764a819678f58be0098cd7a62cb691 Mon Sep 17 00:00:00 2001
From: Michal Prihoda <michal@prihoda.net>
Date: Fri, 22 Jun 2018 11:47:18 +0200
Subject: [PATCH] Return correct status codes from auth rate limiter failure.

---
 core/admin/mailu/internal/__init__.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/core/admin/mailu/internal/__init__.py b/core/admin/mailu/internal/__init__.py
index 45084fe5..6419ad10 100644
--- a/core/admin/mailu/internal/__init__.py
+++ b/core/admin/mailu/internal/__init__.py
@@ -1,3 +1,5 @@
+from flask_limiter import RateLimitExceeded
+
 from mailu import limiter
 
 import socket
@@ -6,6 +8,14 @@ import flask
 
 internal = flask.Blueprint('internal', __name__)
 
+@internal.app_errorhandler(RateLimitExceeded)
+def rate_limit_handler(e):
+    response = flask.Response()
+    response.headers['Auth-Status'] = 'Authentication rate limit from one source exceeded'
+    response.headers['Auth-Error-Code'] = '451 4.3.2'
+    if int(flask.request.headers['Auth-Login-Attempt']) < 10:
+        response.headers['Auth-Wait'] = '3'
+    return response
 
 @limiter.request_filter
 def whitelist_webmail():