修复反向代理后客户共用限流与登录预算

tests/test_deployment.py

@@ -1,14 +1,19 @@
 """Exercise the server factory using disposable state/log paths and no model key."""
 
 import io
+import json
 import logging
 import os
 from pathlib import Path
+import shlex
 import tempfile
 import unittest
 from unittest.mock import patch
 
 from deploy import wsgi
+from logistics_chat import create_app
+from waitress.adjustments import Adjustments
+from waitress.proxy_headers import proxy_headers_middleware
 
 
 class DeploymentTests(unittest.TestCase):
@@ -44,5 +49,80 @@ class DeploymentTests(unittest.TestCase):
         self.assertEqual(app.config['HOST'], '127.0.0.1')
 
 
+class ProxyIdentityTests(unittest.TestCase):
+    def application(self, runtime):
+        root = Path(__file__).resolve().parents[1]
+        if runtime == 'docker':
+            command = next(line[4:] for line in (root/'Dockerfile').read_text(encoding='utf-8').splitlines()
+                           if line.startswith('CMD '))
+            args = json.loads(command)[3:]
+        else:
+            command = next(line[len('ExecStart='):] for line in
+                           (root/'deploy/baiyue-logistics-chat.service').read_text(encoding='utf-8').splitlines()
+                           if line.startswith('ExecStart='))
+            args = shlex.split(command)[3:]
+        options, _ = Adjustments.parse_args(args)
+        options.pop('call'); options.pop('help')
+        settings = Adjustments(**options)
+        folder = tempfile.TemporaryDirectory(); self.addCleanup(folder.cleanup)
+        data = Path(folder.name)
+        app = create_app({'TESTING':True,'DATABASE':data/'history.db','UPLOAD_DIR':data/'uploads',
+                          'BACKUP_DIR':data/'backups','LOG_DIR':data/'logs','RATE_LIMIT_PER_MIN':1000})
+        app.config['PROXY_TEST_PEER'] = '172.18.0.1' if runtime == 'docker' else '127.0.0.1'
+        app.wsgi_app = proxy_headers_middleware(app.wsgi_app, trusted_proxy=settings.trusted_proxy,
+            trusted_proxy_count=settings.trusted_proxy_count, trusted_proxy_headers=settings.trusted_proxy_headers)
+        return app
+
+    def clients(self, app):
+        clients = []
+        for address in ('198.51.100.10', '198.51.100.20'):
+            client = app.test_client()
+            client.environ_base['REMOTE_ADDR'] = app.config['PROXY_TEST_PEER']
+            client.environ_base['HTTP_X_FORWARDED_FOR'] = address
+            client.post('/api/visitor/session', json={'key':str(len(clients)+1)*64})
+            clients.append(client)
+        app.extensions['rate_hits'].clear()
+        return clients
+
+    def test_proxy_clients_do_not_share_request_rate_limits(self):
+        for runtime in ('docker', 'systemd'):
+            with self.subTest(runtime=runtime):
+                app = self.application(runtime)
+                a, b = self.clients(app)
+                app.config['RATE_LIMIT_PER_MIN'] = 2
+                self.assertEqual(a.get('/api/get_history').status_code, 200)
+                self.assertEqual(a.get('/api/get_history').status_code, 200)
+                self.assertEqual(a.get('/api/get_history').status_code, 429)
+                self.assertEqual(b.get('/api/get_history').status_code, 200)
+
+    def test_proxy_clients_do_not_share_failed_login_budgets(self):
+        for runtime in ('docker', 'systemd'):
+            with self.subTest(runtime=runtime):
+                app = self.application(runtime)
+                a, b = self.clients(app)
+                for _ in range(5):
+                    self.assertEqual(a.post('/api/customer/login',json={'username':'admin','password':'wrong'}).status_code,403)
+                self.assertEqual(a.post('/api/customer/login',json={'username':'admin','password':'123456'}).status_code,429)
+                self.assertTrue(b.post('/api/customer/login',json={'username':'admin','password':'123456'}).json.get('authenticated'))
+
+    def test_only_the_last_forwarded_address_is_used_and_other_proxy_claims_are_ignored(self):
+        from flask import request
+        for runtime in ('docker', 'systemd'):
+            with self.subTest(runtime=runtime):
+                app = self.application(runtime)
+                @app.get('/proxy-probe')
+                def probe():
+                    return {'ip':request.remote_addr,'host':request.host,'scheme':request.scheme}
+                response = app.test_client().get('/proxy-probe',base_url='https://chat.example.com',
+                    environ_overrides={'REMOTE_ADDR':app.config['PROXY_TEST_PEER']},headers={
+                        'X-Forwarded-For':'203.0.113.99, 198.51.100.10',
+                        'X-Forwarded-Host':'evil.example','X-Forwarded-Proto':'http'})
+                self.assertEqual(response.json,{'ip':'198.51.100.10','host':'chat.example.com','scheme':'https'})
+                if runtime == 'systemd':
+                    untrusted = app.test_client().get('/proxy-probe',environ_overrides={'REMOTE_ADDR':'203.0.113.1'},
+                        headers={'X-Forwarded-For':'198.51.100.10'})
+                    self.assertEqual(untrusted.json['ip'],'203.0.113.1')
+
+
 if __name__ == '__main__':
     unittest.main()