From 7db8895750163f0c79b839367e9c61491c7b5083 Mon Sep 17 00:00:00 2001
From: n9k <originalcocaineninja@protonmail.com>
Date: Tue, 14 Jun 2022 03:33:09 +0000
Subject: [PATCH] Eyes: send Retry-After header during cooldown

---
 anonstream/helpers/user.py |  1 +
 anonstream/routes/core.py  |  8 +++++++-
 anonstream/user.py         | 13 ++++++-------
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/anonstream/helpers/user.py b/anonstream/helpers/user.py
index 5dbc64e..c5ed466 100644
--- a/anonstream/helpers/user.py
+++ b/anonstream/helpers/user.py
@@ -43,6 +43,7 @@ def generate_user(timestamp, token, broadcaster, presence):
         'last': {
             'seen': timestamp,
             'watching': -inf,
+            'eyes': -inf,
         },
         'presence': presence,
         'linespan': deque(),
diff --git a/anonstream/routes/core.py b/anonstream/routes/core.py
index c4642a8..c4be7de 100644
--- a/anonstream/routes/core.py
+++ b/anonstream/routes/core.py
@@ -1,12 +1,15 @@
 # SPDX-FileCopyrightText: 2022 n9k [https://git.076.ne.jp/ninya9k]
 # SPDX-License-Identifier: AGPL-3.0-or-later
 
+import math
+
 from quart import current_app, request, render_template, abort, make_response, redirect, url_for, abort
+from werkzeug.exceptions import TooManyRequests
 
 from anonstream.captcha import get_captcha_image
 from anonstream.segments import segments, StopSendingSegments
 from anonstream.stream import is_online, get_stream_uptime
-from anonstream.user import watched, create_eyes, renew_eyes, EyesException
+from anonstream.user import watched, create_eyes, renew_eyes, EyesException, RatelimitedEyes
 from anonstream.routes.wrappers import with_user_from, auth_required
 from anonstream.utils.security import generate_csp
 
@@ -27,6 +30,9 @@ async def stream(user):
 
     try:
         eyes_id = create_eyes(user, dict(request.headers))
+    except RatelimitedEyes as e:
+        retry_after, *_ = e.args
+        return TooManyRequests(), {'Retry-After': math.ceil(retry_after)}
     except EyesException:
         return abort(429)
 
diff --git a/anonstream/user.py b/anonstream/user.py
index 7519142..75c50eb 100644
--- a/anonstream/user.py
+++ b/anonstream/user.py
@@ -239,13 +239,11 @@ def get_users_by_presence(timestamp):
 @with_timestamp
 def create_eyes(timestamp, user, headers):
     # Enforce cooldown
-    last_created = max(
-        map(operator.itemgetter('created'), user['eyes']['current'].values()),
-        default=-inf,
-    )
-    last_created_ago = timestamp - last_created
-    if last_created_ago < CONFIG['FLOOD_VIDEO_COOLDOWN']:
-        raise RatelimitedEyes
+    last_created_ago = timestamp - user['last']['eyes']
+    cooldown_ended_ago = last_created_ago - CONFIG['FLOOD_VIDEO_COOLDOWN']
+    cooldown_remaining = -cooldown_ended_ago
+    if cooldown_remaining > 0:
+        raise RatelimitedEyes(cooldown_remaining)
 
     # Enforce max_eyes & overwrite
     if len(user['eyes']['current']) >= CONFIG['FLOOD_VIDEO_MAX_EYES']:
@@ -262,6 +260,7 @@ def create_eyes(timestamp, user, headers):
     # Create eyes
     eyes_id = user['eyes']['total']
     user['eyes']['total'] += 1
+    user['last']['eyes'] = timestamp
     user['eyes']['current'][eyes_id] = {
         'id': eyes_id,
         'token': user['token'],