Use aiofiles for static, refs #272

datasette/app.py

@@ -1,4 +1,5 @@
 import asyncio
+import aiofiles
 from mimetypes import guess_type
 import collections
 import hashlib
@@ -733,7 +734,13 @@ async def asgi_send_html(send, html, status=200, headers=None):
 
 
 async def asgi_send(send, content, status, headers, content_type="text/plain"):
-    # TODO: watch out for Content-Type due to mixed case:
+    await asgi_start(send, status, headers, content_type)
+    await send({"type": "http.response.body", "body": content.encode("utf8")})
+
+
+async def asgi_start(send, status, headers, content_type="text/plain"):
+    # Remove any existing content-type header
+    headers = dict([(k, v) for k, v in headers.items() if k.lower() != "content-type"])
     headers["content-type"] = content_type
     await send(
         {
@@ -745,20 +752,39 @@ async def asgi_send(send, content, status, headers, content_type="text/plain"):
             ],
         }
     )
-    await send({"type": "http.response.body", "body": content.encode("utf8")})
 
 
-def asgi_static(root_path):
+def asgi_static(root_path, chunk_size=4096):
     async def inner_static(scope, receive, send):
         path = scope["url_route"]["kwargs"]["path"]
-        # TODO: prevent ../../ style paths
-        full_path = Path(root_path) / path
-        await asgi_send(
-            send,
-            full_path.open().read(),
-            200,
-            {},
-            content_type=guess_type(str(full_path))[0] or "text/plain",
-        )
+        full_path = (Path(root_path) / path).absolute()
+        # Ensure full_path is within root_path to avoid weird "../" tricks
+        try:
+            full_path.relative_to(root_path)
+        except ValueError:
+            await asgi_send_html(send, "404", 404)
+            return
+        first = True
+        try:
+            async with aiofiles.open(full_path, mode="rb") as fp:
+                if first:
+                    await asgi_start(
+                        send, 200, {}, guess_type(str(full_path))[0] or "text/plain"
+                    )
+                    first = False
+                more_body = True
+                while more_body:
+                    chunk = await fp.read(chunk_size)
+                    more_body = len(chunk) == chunk_size
+                    await send(
+                        {
+                            "type": "http.response.body",
+                            "body": chunk,
+                            "more_body": more_body,
+                        }
+                    )
+        except FileNotFoundError:
+            await asgi_send_html(send, "404", 404)
+            return
 
     return inner_static

setup.py

@@ -49,6 +49,7 @@ setup(
         "pint==0.8.1",
         "pluggy>=0.12.0",
         "uvicorn>=0.8.1",
+        "aiofiles==0.4.0",
     ],
     entry_points="""
         [console_scripts]

tests/fixtures.py

@@ -52,7 +52,9 @@ class TestClient:
         )
         await instance.send_input({"type": "http.request"})
         # First message back should be response.start with headers and status
+        messages = []
         start = await instance.receive_output(2)
+        messages.append(start)
         assert start["type"] == "http.response.start"
         headers = dict(
             [(k.decode("utf8"), v.decode("utf8")) for k, v in start["headers"]]
@@ -62,6 +64,7 @@ class TestClient:
         body = b""
         while True:
             message = await instance.receive_output(2)
+            messages.append(message)
             assert message["type"] == "http.response.body"
             body += message["body"]
             if not message.get("more_body"):

tests/test_html.py

@@ -44,6 +44,14 @@ def test_homepage(app_client_two_attached_databases):
     ] == table_links
 
 
+def test_static(app_client):
+    response = app_client.get("/-/static/app2.css")
+    assert response.status == 404
+    response = app_client.get("/-/static/app.css")
+    assert response.status == 200
+    assert "text/css" == response.headers["content-type"]
+
+
 def test_memory_database_page():
     for client in make_app_client(memory=True):
         response = client.get("/:memory:")