From 8a1a15d7250fde86f9bf3a296861d1dd79fd9eca Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 22 Jun 2019 22:07:41 -0700
Subject: [PATCH] Use aiofiles for static, refs #272

---
 datasette/app.py   | 50 +++++++++++++++++++++++++++++++++++-----------
 setup.py           |  1 +
 tests/fixtures.py  |  3 +++
 tests/test_html.py |  8 ++++++++
 4 files changed, 50 insertions(+), 12 deletions(-)

diff --git a/datasette/app.py b/datasette/app.py
index d5f1b43f..af89d7f6 100644
--- a/datasette/app.py
+++ b/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
diff --git a/setup.py b/setup.py
index 39ebd21d..3a8201cb 100644
--- a/setup.py
+++ b/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]
diff --git a/tests/fixtures.py b/tests/fixtures.py
index afee2f54..b1f54185 100644
--- a/tests/fixtures.py
+++ b/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"):
diff --git a/tests/test_html.py b/tests/test_html.py
index 6b673c13..f9ff393f 100644
--- a/tests/test_html.py
+++ b/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:")