refactor: improve protocol availability

Author: chadsec1

app/logic/data.py

@@ -7,24 +7,34 @@
         COLDWIRE_LEN_OFFSET,
         COLDWIRE_DATA_SEP
 )
+import secrets
+import base64
 
 
 redis_client = get_redis()
 
 
-def get_redis_list(client, key: str) -> list:
-    data = b""
-    while True:
-        raw = client.lpop(key)
-        if raw is None:
-            break
+def b64u_decode(data: str) -> bytes:
+    padding = 4 - (len(data) % 4)
+    if padding != 4:
+        data += "=" * padding
+    return base64.urlsafe_b64decode(data)
 
-        data += raw
 
-    return data
+def delete_data(user_id: str, acks: list[str]) -> None:
+    byte_acks = [b64u_decode(p) for p in acks]
+
+    values = redis_client.lrange(user_id, 0, -1)  
+    for v in values:
+        if any(v.startswith(pref) for pref in byte_acks):
+            res = redis_client.lrem(user_id, 0, v)
 
 def check_new_data(user_id: str) -> bytes:
-    return get_redis_list(redis_client, user_id)
+    data = redis_client.lrange(user_id, 0, -1)
+    if not data:
+        return b""
+
+    return b"".join(data)
 
 def data_processor(user_id: str, recipient: str, blob: bytes) -> None:
     if recipient.isdigit():
@@ -41,10 +51,10 @@ def data_processor(user_id: str, recipient: str, blob: bytes) -> None:
         if COLDWIRE_DATA_SEP in user_id:
             raise ValueError("User ID cannot have null byte!")
 
-        payload = user_id + COLDWIRE_DATA_SEP + blob
+        payload =  user_id + COLDWIRE_DATA_SEP + blob
         length_prefix = len(payload).to_bytes(COLDWIRE_LEN_OFFSET, "big")
 
-        payload = length_prefix + payload
+        payload = secrets.token_bytes(32) + length_prefix + payload
 
         redis_client.rpush(recipient, payload)
 

app/logic/federation_utils.py

@@ -21,6 +21,7 @@
 from base64 import b64encode, b64decode
 from datetime import datetime, timezone, timedelta
 import json
+import secrets
 
 redis_client = get_redis()
 
@@ -117,7 +118,7 @@ def federation_processor(url: str, sender: str, recipient: str, blob: bytes) ->
     payload = sender_with_url + COLDWIRE_DATA_SEP + blob
     length_prefix = len(payload).to_bytes(COLDWIRE_LEN_OFFSET, "big")
 
-    payload = length_prefix + payload
+    payload = secrets.token_bytes(32) + length_prefix + payload
 
     redis_client.rpush(recipient, payload)
 

app/routes/data.py

@@ -1,19 +1,22 @@
-from fastapi import APIRouter, Request, HTTPException, Response, Depends, Form, UploadFile, File
-from app.logic.data import check_new_data, data_processor
+from fastapi import APIRouter, Request, HTTPException, Response, Depends, Form, UploadFile, File, Query
+from app.logic.data import check_new_data, delete_data, data_processor
 from app.utils.jwt import verify_jwt_token
 from app.core.constants import LONGPOLL_MAX
+from typing import Optional
 import asyncio
 import json
 
 router = APIRouter()
 
 
 @router.get("/data/longpoll")
-async def get_data_longpoll(request: Request, response: Response, user=Depends(verify_jwt_token)):
+async def get_data_longpoll(request: Request, response: Response, acks: Optional[list[str]] = Query(None), user=Depends(verify_jwt_token)):
+    if acks:
+        await asyncio.to_thread(delete_data, user["id"], acks)
+        
     for _ in range(LONGPOLL_MAX):
         if await request.is_disconnected():
-            # Don't attempt to check for new data if client disconnects before LONGPOLL_MAX seconds
-            # This is crucial to perserve data as they get deleted after being read 
+            # Don't bother checking for new data if client disconnects before LONGPOLL_MAX seconds
             return Response(content=b'', media_type="application/octet-stream")
 
         data = await asyncio.to_thread(check_new_data, user["id"])