From 4aff3e163409d54387c45084ed5f69cce1497a62 Mon Sep 17 00:00:00 2001
From: Minecon724 <minecon724@noreply.git.m724.eu>
Date: Wed, 2 Apr 2025 16:47:29 +0200
Subject: [PATCH] Refactor limiter a bit

---
 fly.toml                 |  2 +-
 src/anonchat/__init__.py | 23 +++++++----------------
 2 files changed, 8 insertions(+), 17 deletions(-)

diff --git a/fly.toml b/fly.toml
index e463334..de4ab97 100644
--- a/fly.toml
+++ b/fly.toml
@@ -16,7 +16,7 @@ primary_region = 'ams'
   force_https = true
   auto_stop_machines = 'stop'
   auto_start_machines = true
-  min_machines_running = 0
+  min_machines_running = 1
   processes = ['app']
 
 [[vm]]
diff --git a/src/anonchat/__init__.py b/src/anonchat/__init__.py
index 7a00076..25f0216 100644
--- a/src/anonchat/__init__.py
+++ b/src/anonchat/__init__.py
@@ -8,6 +8,7 @@ from flask import render_template, request, jsonify
 from flask_wtf.csrf import CSRFProtect
 from flask_session import Session
 import redis
+from werkzeug.middleware.proxy_fix import ProxyFix
 from argon2 import PasswordHasher
 
 # Load environment variables from .env file
@@ -27,7 +28,7 @@ app.config['ADMIN_USERNAME'] = os.environ.get('ADMIN_USERNAME', 'admin')
 app.config['ADMIN_PASSWORD'] = os.environ.get('ADMIN_PASSWORD', None)
 app.config['ADMIN_FORCE_RESET'] = os.environ.get('ADMIN_FORCE_RESET', 'false').lower() == 'true'
 # Rate limit configurations
-app.config['RATELIMIT_STORAGE_URL'] = os.environ.get('RATELIMIT_STORAGE_URL', os.environ.get('REDIS_URL'))
+app.config['RATELIMIT_STORAGE_URI'] = os.environ.get('RATELIMIT_STORAGE_URI', os.environ.get('REDIS_URL'))
 app.config['RATELIMIT_HEADERS_ENABLED'] = True
 app.config['RATELIMIT_KEY_PREFIX'] = 'anonchat_rate_limit'
 # Whether app is behind a proxy (get from env, default to False)
@@ -45,6 +46,9 @@ elif app.config['SESSION_TYPE'] == 'filesystem':
     app.config['SESSION_FILE_MODE'] = os.environ.get('SESSION_FILE_MODE', 384)
 app.config['SESSION_KEY_PREFIX'] = 'anonchat_session:'
 
+if app.config['BEHIND_PROXY']:
+    app.wsgi_app = ProxyFix(app.wsgi_app)
+
 # Initialize password hasher
 # Parameters source: https://cheatsheetseries.owasp.org/cheatsheets/Password_Storage_Cheat_Sheet.html#argon2id
 password_hasher = PasswordHasher(
@@ -59,23 +63,10 @@ Session(app)
 # Initialize CSRF protection
 csrf = CSRFProtect(app)
 
-# Function to get client IP address, respecting X-Forwarded-For when behind a proxy
-def get_client_ip():
-    if app.config['BEHIND_PROXY']:
-        # Get the first IP in X-Forwarded-For, which should be the client
-        forwarded_for = request.headers.get('X-Forwarded-For')
-        if forwarded_for:
-            return forwarded_for.split(',')[0].strip()
-    # Fall back to remote_addr if not behind proxy or X-Forwarded-For not found
-    return request.remote_addr
-
 # Initialize limiter with custom key_func
-limiter = Limiter(
-    get_client_ip,  # Use our custom function instead of get_remote_address
-    app=app,
-    storage_uri=app.config['RATELIMIT_STORAGE_URL']
-)
+limiter = Limiter(get_remote_address, app=app)
 
+# Initialize database
 db = SQLAlchemy(app)
 
 # Import models