From c519a056c26c40fcea7efa73f1c923302b5a23d6 Mon Sep 17 00:00:00 2001
From: Kumi <git@kumi.email>
Date: Sat, 3 Aug 2024 16:33:32 +0200
Subject: [PATCH] feat(proxy): improve streaming and error handling

Enhanced proxy endpoint by adding streaming support to handle large files efficiently. Introduced error handling for various HTTP errors and Internal Server errors to provide more robust feedback. Updated content disposition to handle attachment downloads properly.
---
 app.py | 48 ++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 38 insertions(+), 10 deletions(-)

diff --git a/app.py b/app.py
index 5ff6eed..5a3e7c2 100644
--- a/app.py
+++ b/app.py
@@ -1,22 +1,50 @@
-from flask import Flask, render_template, request
+from flask import Flask, render_template, request, Response, abort
+from werkzeug.exceptions import InternalServerError
 
 from urllib.request import urlopen
+from urllib.error import HTTPError
+from urllib.parse import unquote
 
 app = Flask(__name__)
 
-@app.route('/')
+
+@app.route("/")
 def home():
     domain = request.url_root
-    return render_template('index.html', domain=domain)
+    return render_template("index.html", domain=domain)
 
-@app.route('/<path:url>')
+
+@app.route("/<path:url>")
 def proxy(url):
-    jsdelivr_url = f'https://cdnjs.cloudflare.com/{url}'
+    jsdelivr_url = f"https://cdnjs.cloudflare.com/{url}"
 
-    response = urlopen(jsdelivr_url)
-    content = response.read()
+    def generate():
+        # Subfunction to allow streaming the data instead of
+        # downloading all of it at once
+        try:
+            with urlopen(unquote(url)) as data:
+                while True:
+                    chunk = data.read(1024 * 1024)
+                    if not chunk:
+                        break
+                    yield chunk
+        except HTTPError as e:
+            abort(e.code)
 
-    return content
+    try:
+        with urlopen(unquote(jsdelivr_url)) as data:
+            content_type = data.headers["content-type"]
+    except HTTPError as e:
+        abort(e.code)
+    except KeyError:
+        raise InternalServerError()
 
-if __name__ == '__main__':
-    app.run()
\ No newline at end of file
+    headers = dict()
+    filename = url.split("/")[-1]
+    headers["Content-Disposition"] = f'attachment; filename="{filename}"'
+
+    return Response(generate(), content_type=content_type, headers=headers)
+
+
+if __name__ == "__main__":
+    app.run()