From 581c6f1714b64e0c6e24bf47c238821803638966 Mon Sep 17 00:00:00 2001
From: Zed <zedeus@pm.me>
Date: Fri, 19 Jun 2026 22:00:28 +0200
Subject: [PATCH] Add retry logic and timeout to media proxy

---
 src/routes/media.nim | 98 +++++++++++++++++++++++++++-----------------
 1 file changed, 61 insertions(+), 37 deletions(-)

diff --git a/src/routes/media.nim b/src/routes/media.nim
index 40f5a6c..0fa02aa 100644
--- a/src/routes/media.nim
+++ b/src/routes/media.nim
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 import uri, strutils, httpclient, os, hashes, base64, re
 import asynchttpserver, asyncstreams, asyncfile, asyncnet
+import asyncdispatch
 
 import jester
 
@@ -32,47 +33,70 @@ template respond*(req: asynchttpserver.Request; headers) =
 
 proc proxyMedia*(req: jester.Request; url: string): Future[HttpCode] {.async.} =
   result = Http200
-  let
-    request = req.getNativeReq()
-    client = newAsyncHttpClient(maxRedirects = 0)
+  let request = req.getNativeReq()
 
-  try:
-    let res = await client.get(url)
-    if res.status != "200 OK":
-      if res.status != "404 Not Found":
-        echo "[media] Proxying failed, status: $1, url: $2" % [res.status, url]
-      return Http404
-
-    let hashed = $hash(url)
-    if request.headers.getOrDefault("If-None-Match") == hashed:
-      return Http304
-
-    let contentLength =
-      if res.headers.hasKey("content-length"):
-        res.headers["content-length", 0]
+  for attempt in 0 .. 2:
+    let client = newAsyncHttpClient(maxRedirects = 0)
+    var shouldRetry = false
+    try:
+      let resFut = client.get(url)
+      let completed = await withTimeout(resFut, 5000)
+      if not completed:
+        if attempt < 2:
+          echo "[media] Retry $1/2, timeout after 5s, url: $2" % [$(attempt + 1), url]
+          shouldRetry = true
+        else:
+          echo "[media] Proxying timeout after 5s, url: $1" % [url]
+          return Http504
       else:
-        ""
+        let res = resFut.read()
+        if res.status != "200 OK":
+          if res.status == "404 Not Found":
+            return Http404
+          if attempt < 2:
+            echo "[media] Retry $1/2, status: $2, url: $3" % [$(attempt + 1), res.status, url]
+            shouldRetry = true
+          else:
+            echo "[media] Proxying failed, status: $1, url: $2" % [res.status, url]
+            return Http404
+        else:
+          let hashed = $hash(url)
+          if request.headers.getOrDefault("If-None-Match") == hashed:
+            return Http304
 
-    let headers = newHttpHeaders({
-      "content-type": res.headers["content-type", 0],
-      "content-length": contentLength,
-      "cache-control": maxAge,
-      "etag": hashed
-    })
+          let contentLength =
+            if res.headers.hasKey("content-length"):
+              res.headers["content-length", 0]
+            else:
+              ""
 
-    respond(request, headers)
+          let headers = newHttpHeaders({
+            "content-type": res.headers["content-type", 0],
+            "content-length": contentLength,
+            "cache-control": maxAge,
+            "etag": hashed
+          })
 
-    var (hasValue, data) = (true, "")
-    while hasValue:
-      (hasValue, data) = await res.bodyStream.read()
-      if hasValue:
-        await request.client.send(data)
-    data.setLen 0
-  except HttpRequestError, ProtocolError, OSError:
-    echo "[media] Proxying exception, error: $1, url: $2" % [getCurrentExceptionMsg(), url]
-    result = Http404
-  finally:
-    client.close()
+          respond(request, headers)
+
+          var (hasValue, data) = (true, "")
+          while hasValue:
+            (hasValue, data) = await res.bodyStream.read()
+            if hasValue:
+              await request.client.send(data)
+          data.setLen 0
+          return Http200
+    except CatchableError:
+      if attempt < 2:
+        echo "[media] Retry $1/2, error: $2, url: $3" % [$(attempt + 1), getCurrentExceptionMsg(), url]
+        shouldRetry = true
+      else:
+        echo "[media] Proxying exception, error: $1, url: $2" % [getCurrentExceptionMsg(), url]
+        result = Http404
+    finally:
+      client.close()
+    if not shouldRetry:
+      break
 
 template check*(code): untyped =
   if code != Http200:
@@ -129,7 +153,7 @@ proc createMediaRouter*(cfg: Config) =
       if getHmac(url) != request.matches[1]:
         resp Http403, showError("Failed to verify signature", cfg)
 
-      if ".mp4" in url or ".ts" in url or ".m4s" in url:
+      if ".mp4" in url or ".ts" in url or ".m4s" in url or ".aac" in url:
         let code = await proxyMedia(request, url)
         check code