Some more project structure

anonstream/__init__.py

@@ -42,9 +42,10 @@ async def create_app():
     assert app.config['MAX_CHAT_MESSAGES'] >= app.config['MAX_CHAT_SCROLLBACK']
     assert app.config['THRESHOLD_ABSENT'] >= app.config['THRESHOLD_IDLE']
 
-    app.chat = {'messages': OrderedDict(), 'nonce_hashes': set()}
-    app.users = {}
-    app.websockets = set()
+    app.messages_by_id = OrderedDict()
+    app.users_by_token = {}
+    app.messages = app.messages_by_id.values()
+    app.users = app.users_by_token.values()
     app.segments_directory_cache = DirectoryCache(config['stream']['segments_dir'])
 
     async with app.app_context():

anonstream/chat.py

@@ -1,22 +1,37 @@
 import time
 from datetime import datetime
 
-from quart import escape
+from quart import current_app, escape
 
-from anonstream.user import users_for_websocket
 from anonstream.helpers.chat import generate_nonce_hash
+from anonstream.utils.chat import message_for_websocket
+
+MESSAGES_BY_ID = current_app.messages_by_id
+MESSAGES = current_app.messages
+USERS_BY_TOKEN = current_app.users_by_token
+USERS = current_app.users
 
 class Rejected(Exception):
     pass
 
-async def broadcast(websockets, payload):
-    for queue in websockets:
-        await queue.put(payload)
+async def broadcast(users, payload):
+    for user in users:
+        for queue in user['websockets']:
+            await queue.put(payload)
 
-async def add_chat_message(chat, users, websockets, user, nonce, comment):
+def messages_for_websocket():
+    return list(map(
+        lambda message: message_for_websocket(
+            user=USERS_BY_TOKEN[message['token']],
+            message=message,
+        ),
+        MESSAGES,
+    ))
+
+async def add_chat_message(user, nonce, comment):
     # check message
-    nonce_hash = generate_nonce_hash(nonce)
-    if nonce_hash in chat['nonce_hashes']:
+    message_id = generate_nonce_hash(nonce)
+    if message_id in MESSAGES_BY_ID:
         raise Rejected('Discarded suspected duplicate message')
     if len(comment) == 0:
         raise Rejected('Message was empty')
@@ -25,17 +40,19 @@ async def add_chat_message(chat, users, websockets, user, nonce, comment):
     timestamp_ms = time.time_ns() // 1_000_000
     timestamp = timestamp_ms // 1000
     try:
-        last_message = next(reversed(chat['messages'].values()))
+        last_message = next(reversed(MESSAGES))
     except StopIteration:
-        message_id = timestamp_ms
+        seq = timestamp_ms
     else:
-        if timestamp <= last_message['id']:
-            message_id = last_message['id'] + 1
+        if timestamp_ms > last_message['seq']:
+            seq = timestamp_ms
+        else:
+            seq = last_message['seq'] + 1
     dt = datetime.utcfromtimestamp(timestamp)
     markup = escape(comment)
-    chat['messages'][message_id] = {
+    MESSAGES_BY_ID[message_id] = {
         'id': message_id,
-        'nonce_hash': nonce_hash,
+        'seq': seq,
         'token': user['token'],
         'timestamp': timestamp,
         'date': dt.strftime('%Y-%m-%d'),
@@ -45,15 +62,12 @@ async def add_chat_message(chat, users, websockets, user, nonce, comment):
         'markup': markup,
     }
 
-    # collect nonce hash
-    chat['nonce_hashes'].add(nonce_hash)
-
     # broadcast message to websockets
     await broadcast(
-        websockets,
+        USERS,
         payload={
             'type': 'chat',
-            'id': message_id,
+            'seq': seq,
             'token_hash': user['token_hash'],
             'markup': markup,
         }

anonstream/helpers/user.py

@@ -14,7 +14,7 @@ def generate_token_hash(token):
     digest = hashlib.sha256(parts).digest()
     return base64.b32encode(digest)[:26].lower().decode()
 
-def generate_user(token, broadcaster, timestamp):
+def generate_user(timestamp, token, broadcaster):
     colour = generate_colour(
         seed='name\0' + token,
         bg=CONFIG['CHAT_BACKGROUND_COLOUR'],
@@ -23,6 +23,7 @@ def generate_user(token, broadcaster, timestamp):
     return {
         'token': token,
         'token_hash': generate_token_hash(token),
+        'websockets': set(),
         'broadcaster': broadcaster,
         'name': None,
         'color': colour_to_color(colour),
@@ -49,7 +50,10 @@ def is_idle(timestamp, user):
     return is_present(timestamp, user) and not is_watching(timestamp, user)
 
 def is_present(timestamp, user):
-    return user['seen']['last'] >= timestamp - CONFIG['THRESHOLD_ABSENT']
+    return (
+        user['seen']['last'] >= timestamp - CONFIG['THRESHOLD_ABSENT']
+        or len(user['websockets']) > 0
+    )
 
 def is_absent(timestamp, user):
     return not is_present(timestamp, user)

anonstream/routes/nojs.py

@@ -24,8 +24,8 @@ async def nojs_chat(user):
     return await render_template(
         'nojs_chat.html',
         user=user,
-        users=current_app.users,
-        messages=current_app.chat['messages'].values(),
+        users_by_token=current_app.users_by_token,
+        messages=current_app.messages,
         get_default_name=get_default_name,
     )
 
@@ -53,14 +53,7 @@ async def nojs_submit_message(user):
     nonce = form.get('nonce', '')
 
     try:
-        await add_chat_message(
-            chat=current_app.chat,
-            users=current_app.users,
-            websockets=current_app.websockets,
-            user=user,
-            nonce=nonce,
-            comment=comment,
-        )
+        await add_chat_message(user, nonce, comment)
     except Rejected as e:
         notice, *_ = e.args
         notice_id = add_notice(user, notice)

anonstream/routes/websocket.py

@@ -9,21 +9,11 @@ from anonstream.routes.wrappers import with_user_from
 @with_user_from(websocket)
 async def live(user):
     queue = asyncio.Queue()
-    current_app.websockets.add(queue)
+    user['websockets'].add(queue)
 
-    producer = websocket_outbound(
-        queue=queue,
-        messages=current_app.chat['messages'].values(),
-        users=current_app.users,
-    )
-    consumer = websocket_inbound(
-        queue=queue,
-        chat=current_app.chat,
-        users=current_app.users,
-        connected_websockets=current_app.websockets,
-        user=user,
-    )
+    producer = websocket_outbound(queue)
+    consumer = websocket_inbound(queue, user)
     try:
         await asyncio.gather(producer, consumer)
     finally:
-        current_app.websockets.remove(queue)
+        user['websockets'].remove(queue)

anonstream/routes/wrappers.py

@@ -5,17 +5,22 @@ from quart import current_app, request, abort, make_response
 from werkzeug.security import check_password_hash
 
 from anonstream.user import sunset, user_for_websocket
-from anonstream.websocket import broadcast
+from anonstream.chat import broadcast
 from anonstream.helpers.user import generate_user
 from anonstream.utils.user import generate_token
 
+CONFIG = current_app.config
+MESSAGES = current_app.messages
+USERS_BY_TOKEN = current_app.users_by_token
+USERS = current_app.users
+
 def check_auth(context):
     auth = context.authorization
     return (
         auth is not None
         and auth.type == "basic"
-        and auth.username == current_app.config["AUTH_USERNAME"]
-        and check_password_hash(current_app.config["AUTH_PWHASH"], auth.password)
+        and auth.username == CONFIG["AUTH_USERNAME"]
+        and check_password_hash(CONFIG["AUTH_PWHASH"], auth.password)
     )
 
 def auth_required(f):
@@ -44,40 +49,40 @@ def with_user_from(context):
             # Check if broadcaster
             broadcaster = check_auth(context)
             if broadcaster:
-                token = current_app.config['AUTH_TOKEN']
+                token = CONFIG['AUTH_TOKEN']
             else:
                 token = context.args.get('token') or context.cookies.get('token') or generate_token()
 
             # Remove non-visible absent users
-            token_hashes = sunset(
-                messages=current_app.chat['messages'].values(),
-                users=current_app.users,
+            sunsetted_token_hashes = sunset(
+                messages=MESSAGES,
+                users_by_token=USERS_BY_TOKEN,
             )
-            if len(token_hashes) > 0:
+            if sunsetted_token_hashes:
                 await broadcast(
-                    current_app.websockets,
+                    users=USERS,
                     payload={
                         'type': 'rem-users',
-                        'token_hashes': token_hashes,
+                        'token_hashes': sunsetted_token_hashes,
                     }
                 )
 
             # Update / create user
-            user = current_app.users.get(token)
+            user = USERS_BY_TOKEN.get(token)
             if user is not None:
                 user['seen']['last'] = timestamp
             else:
                 user = generate_user(
-                    secret=current_app.config['SECRET_KEY'],
+                    timestamp=timestamp,
                     token=token,
                     broadcaster=broadcaster,
-                    timestamp=timestamp,
                 )
-                current_app.users[token] = user
+                USERS_BY_TOKEN[token] = user
                 await broadcast(
-                    current_app.websockets,
+                    USERS,
                     payload={
                         'type': 'add-user',
+                        'token_hash': user['token_hash'],
                         'user': user_for_websocket(user),
                     }
                 )

anonstream/static/anonstream.js

@@ -52,7 +52,7 @@ const create_chat_message = (object) => {
 
   const chat_message = document.createElement("li");
   chat_message.classList.add("chat-message");
-  chat_message.dataset.id = object.id;
+  chat_message.dataset.seq = object.seq;
   chat_message.dataset.tokenHash = object.token_hash;
 
   const chat_message_name = document.createElement("span");
@@ -126,11 +126,11 @@ const on_websocket_message = (event) => {
       users = receipt.users;
       update_user_styles();
 
-      const ids = new Set(receipt.chat.map((message) => {return message.id;}));
+      const seqs = new Set(receipt.messages.map((message) => {return message.seq;}));
       const to_delete = [];
       for (const chat_message of chat_messages.children) {
-        const chat_message_id = parseInt(chat_message.dataset.id);
-        if (!ids.has(chat_message_id)) {
+        const chat_message_seq = parseInt(chat_message.dataset.seq);
+        if (!seqs.has(chat_message_seq)) {
           to_delete.push(chat_message);
         }
       }
@@ -138,9 +138,10 @@ const on_websocket_message = (event) => {
         chat_message.remove();
       }
 
-      const last_id = Math.max(...[...chat_messages.children].map((element) => parseInt(element.dataset.id)));
-      for (const message of receipt.chat) {
-        if (message.id > last_id) {
+      const last = chat_messages.children.length == 0 ? null : chat_messages.children[chat_messages.children.length - 1];
+      const last_seq = last === null ? null : parseInt(last.dataset.seq);
+      for (const message of receipt.messages) {
+        if (message.seq > last_seq) {
           const chat_message = create_chat_message(message);
           chat_messages.insertAdjacentElement("beforeend", chat_message);
         }
@@ -184,7 +185,7 @@ const on_websocket_message = (event) => {
 
     case "add-user":
         console.log("ws add-user", receipt);
-        users[receipt.user.token_hash] = receipt.user;
+        users[receipt.token_hash] = receipt.user;
         update_user_styles();
         break;
 

anonstream/templates/nojs_chat.html

@@ -48,7 +48,7 @@
   <ul id="chat-messages">
     {% for message in messages | reverse %}
       <li class="chat-message">
-        {% with user = users[message.token] %}
+        {% with user = users_by_token[message.token] %}
           <span class="chat-message__name" style="color:{{ user.color }};" data-id="{{ message.id }}">{{ user.name or get_default_name(user) }}</span>:&nbsp;<span class="chat-message__markup">{{ message.markup }}</span>
         {% endwith %}
       </li>

anonstream/user.py

@@ -10,6 +10,8 @@ from anonstream.utils.colour import color_to_colour, get_contrast, NotAColor
 from anonstream.utils.user import user_for_websocket
 
 CONFIG = current_app.config
+MESSAGES = current_app.messages
+USERS = current_app.users
 
 class BadAppearance(Exception):
     pass
@@ -71,19 +73,19 @@ def see(user):
     user['seen']['last'] = int(time.time())
 
 @with_timestamp
-def users_for_websocket(timestamp, messages, users):
+def users_for_websocket(timestamp):
     visible_users = filter(
-        lambda user: is_visible(timestamp, messages, user),
-        users.values(),
+        lambda user: is_visible(timestamp, MESSAGES, user),
+        USERS,
     )
     return {
-        user['token_hash']: user_for_websocket(user, include_token_hash=False)
+        user['token_hash']: user_for_websocket(user)
         for user in visible_users
     }
 
 last_checkup = -inf
 
-def sunset(messages, users):
+def sunset(messages, users_by_token):
     global last_checkup
 
     timestamp = int(time.time())
@@ -91,13 +93,13 @@ def sunset(messages, users):
         return []
 
     to_delete = []
-    for token in users:
-        user = users[token]
+    for token in users_by_token:
+        user = users_by_token[token]
         if not is_visible(timestamp, messages, user):
             to_delete.append(token)
 
     for index, token in enumerate(to_delete):
-        to_delete[index] = users.pop(token)['token_hash']
+        to_delete[index] = users_by_token.pop(token)['token_hash']
 
     last_checkup = timestamp
     return to_delete

anonstream/utils/chat.py

@@ -8,11 +8,9 @@ class NonceReuse(Exception):
 def generate_nonce():
     return secrets.token_urlsafe(16)
 
-def message_for_websocket(users, message):
-    message_keys = ('id', 'date', 'time_minutes', 'time_seconds', 'markup')
+def message_for_websocket(user, message):
+    message_keys = ('seq', 'date', 'time_minutes', 'time_seconds', 'markup')
     user_keys = ('token_hash',)
-
-    user = users[message['token']]
     return {
         **{key: message[key] for key in message_keys},
         **{key: user[key] for key in user_keys},

anonstream/utils/user.py

@@ -9,10 +9,8 @@ from quart import escape, Markup
 def generate_token():
     return secrets.token_hex(16)
 
-def user_for_websocket(user, include_token_hash=True):
+def user_for_websocket(user):
     keys = ['broadcaster', 'name', 'color', 'tripcode']
-    if include_token_hash:
-        keys.append('token_hash')
     return {key: user[key] for key in keys}
 
 def concatenate_for_notice(string, *tuples):

anonstream/websocket.py

@@ -3,26 +3,23 @@ import asyncio
 from quart import current_app, websocket
 
 from anonstream.stream import get_stream_title, get_stream_uptime
-from anonstream.chat import broadcast, add_chat_message, Rejected
+from anonstream.chat import messages_for_websocket, add_chat_message, Rejected
 from anonstream.user import users_for_websocket, see
 from anonstream.wrappers import with_first_argument
 from anonstream.helpers.user import is_present
-from anonstream.utils.chat import generate_nonce, message_for_websocket
+from anonstream.utils.chat import generate_nonce
 from anonstream.utils.websocket import parse_websocket_data, Malformed
 
 CONFIG = current_app.config
 
-async def websocket_outbound(queue, messages, users):
+async def websocket_outbound(queue):
     payload = {
         'type': 'init',
         'nonce': generate_nonce(),
         'title': get_stream_title(),
         'uptime': get_stream_uptime(),
-        'chat': list(map(
-            with_first_argument(users)(message_for_websocket),
-            messages,
-        )),
-        'users': users_for_websocket(messages, users),
+        'messages': messages_for_websocket(),
+        'users': users_for_websocket(),
         'default': {
             True: CONFIG['DEFAULT_HOST_NAME'],
             False: CONFIG['DEFAULT_ANON_NAME'],
@@ -33,7 +30,7 @@ async def websocket_outbound(queue, messages, users):
         payload = await queue.get()
         await websocket.send_json(payload)
 
-async def websocket_inbound(queue, chat, users, connected_websockets, user):
+async def websocket_inbound(queue, user):
     while True:
         receipt = await websocket.receive_json()
         see(user)
@@ -47,14 +44,7 @@ async def websocket_inbound(queue, chat, users, connected_websockets, user):
             }
         else:
             try:
-                markup = await add_chat_message(
-                    chat,
-                    users,
-                    connected_websockets,
-                    user,
-                    nonce,
-                    comment,
-                )
+                markup = await add_chat_message(user, nonce, comment)
             except Rejected as e:
                 notice, *_ = e.args
                 payload = {