response.set_cookie(), closes #795

datasette/actor_auth_cookie.py

@@ -1,6 +1,5 @@
 from datasette import hookimpl
 from itsdangerous import BadSignature
-from http.cookies import SimpleCookie
 
 
 @hookimpl

datasette/app.py

@@ -3,7 +3,6 @@ import asgi_csrf
 import collections
 import datetime
 import hashlib
-from http.cookies import SimpleCookie
 import itertools
 import json
 import os
@@ -442,19 +441,9 @@ class Datasette:
     def _write_messages_to_response(self, request, response):
         if getattr(request, "_messages", None):
             # Set those messages
-            cookie = SimpleCookie()
+            response.set_cookie("ds_messages", self.sign(request._messages, "messages"))
-            cookie["ds_messages"] = self.sign(request._messages, "messages")
-            cookie["ds_messages"]["path"] = "/"
-            # TODO: Co-exist with existing set-cookie headers
-            assert "set-cookie" not in response.headers
-            response.headers["set-cookie"] = cookie.output(header="").lstrip()
         elif getattr(request, "_messages_should_clear", False):
-            cookie = SimpleCookie()
+            response.set_cookie("ds_messages", "", expires=0, max_age=0)
-            cookie["ds_messages"] = ""
-            cookie["ds_messages"]["path"] = "/"
-            # TODO: Co-exist with existing set-cookie headers
-            assert "set-cookie" not in response.headers
-            response.headers["set-cookie"] = cookie.output(header="").lstrip()
 
     def _show_messages(self, request):
         if getattr(request, "_messages", None):

datasette/utils/asgi.py

@@ -4,10 +4,15 @@ from mimetypes import guess_type
 from urllib.parse import parse_qs, urlunparse, parse_qsl
 from pathlib import Path
 from html import escape
-from http.cookies import SimpleCookie
+from http.cookies import SimpleCookie, Morsel
 import re
 import aiofiles
 
+# Workaround for adding samesite support to pre 3.8 python
+Morsel._reserved["samesite"] = "SameSite"
+# Thanks, Starlette:
+# https://github.com/encode/starlette/blob/519f575/starlette/responses.py#L17
+
 
 class NotFound(Exception):
     pass
@@ -17,6 +22,9 @@ class Forbidden(Exception):
     pass
 
 
+SAMESITE_VALUES = ("strict", "lax", "none")
+
+
 class Request:
     def __init__(self, scope, receive):
         self.scope = scope
@@ -370,20 +378,24 @@ class Response:
         self.body = body
         self.status = status
         self.headers = headers or {}
+        self._set_cookie_headers = []
         self.content_type = content_type
 
     async def asgi_send(self, send):
         headers = {}
         headers.update(self.headers)
         headers["content-type"] = self.content_type
+        raw_headers = [
+            [key.encode("utf-8"), value.encode("utf-8")]
+            for key, value in headers.items()
+        ]
+        for set_cookie in self._set_cookie_headers:
+            raw_headers.append([b"set-cookie", set_cookie.encode("utf-8")])
         await send(
             {
                 "type": "http.response.start",
                 "status": self.status,
-                "headers": [
+                "headers": raw_headers,
-                    [key.encode("utf-8"), value.encode("utf-8")]
-                    for key, value in headers.items()
-                ],
             }
         )
         body = self.body
@@ -391,6 +403,37 @@ class Response:
             body = body.encode("utf-8")
         await send({"type": "http.response.body", "body": body})
 
+    def set_cookie(
+        self,
+        key,
+        value="",
+        max_age=None,
+        expires=None,
+        path="/",
+        domain=None,
+        secure=False,
+        httponly=False,
+        samesite="lax",
+    ):
+        assert samesite in SAMESITE_VALUES, "samesite should be one of {}".format(
+            SAMESITE_VALUES
+        )
+        cookie = SimpleCookie()
+        cookie[key] = value
+        for prop_name, prop_value in (
+            ("max_age", max_age),
+            ("expires", expires),
+            ("path", path),
+            ("domain", domain),
+            ("samesite", samesite),
+        ):
+            if prop_value is not None:
+                cookie[key][prop_name.replace("_", "-")] = prop_value
+        for prop_name, prop_value in (("secure", secure), ("httponly", httponly)):
+            if prop_value:
+                cookie[key][prop_name] = True
+        self._set_cookie_headers.append(cookie.output(header="").strip())
+
     @classmethod
     def html(cls, body, status=200, headers=None):
         return cls(

datasette/views/special.py

@@ -1,7 +1,6 @@
 import json
 from datasette.utils.asgi import Response
 from .base import BaseView
-from http.cookies import SimpleCookie
 import secrets
 
 
@@ -62,17 +61,8 @@ class AuthTokenView(BaseView):
             return Response("Root token has already been used", status=403)
         if secrets.compare_digest(token, self.ds._root_token):
             self.ds._root_token = None
-            cookie = SimpleCookie()
+            response = Response.redirect("/")
-            cookie["ds_actor"] = self.ds.sign({"id": "root"}, "actor")
+            response.set_cookie("ds_actor", self.ds.sign({"id": "root"}, "actor"))
-            cookie["ds_actor"]["path"] = "/"
-            response = Response(
-                body="",
-                status=302,
-                headers={
-                    "Location": "/",
-                    "set-cookie": cookie.output(header="").lstrip(),
-                },
-            )
             return response
         else:
             return Response("Invalid token", status=403)

docs/internals.rst

@@ -131,6 +131,36 @@ Each of these responses will use the correct corresponding content-type - ``text
 
 Each of the helper methods take optional ``status=`` and ``headers=`` arguments, documented above.
 
+.. _internals_response_set_cookie:
+
+Setting cookies with response.set_cookie()
+------------------------------------------
+
+To set cookies on the response, use the ``response.set_cookie(...)`` method. The method signature looks like this:
+
+.. code-block:: python
+
+    def set_cookie(
+        self,
+        key,
+        value="",
+        max_age=None,
+        expires=None,
+        path="/",
+        domain=None,
+        secure=False,
+        httponly=False,
+        samesite="lax",
+    ):
+
+You can use this with :ref:`datasette.sign() <datasette_sign>` to set signed cookies. Here's how you would set the ``ds_actor`` cookie for use with Datasette :ref:`authentication <authentication>`:
+
+.. code-block:: python
+
+    response = Response.redirect("/")
+    response.set_cookie("ds_actor", datasette.sign({"id": "cleopaws"}, "actor"))
+    return response
+
 .. _internals_datasette:
 
 Datasette class

tests/test_internals_response.py

@@ -1,4 +1,5 @@
 from datasette.utils.asgi import Response
+import pytest
 
 
 def test_response_html():
@@ -26,3 +27,28 @@ def test_response_redirect():
     response = Response.redirect("/foo")
     assert 302 == response.status
     assert "/foo" == response.headers["Location"]
+
+
+@pytest.mark.asyncio
+async def test_response_set_cookie():
+    events = []
+
+    async def send(event):
+        events.append(event)
+
+    response = Response.redirect("/foo")
+    response.set_cookie("foo", "bar", max_age=10, httponly=True)
+    await response.asgi_send(send)
+
+    assert [
+        {
+            "type": "http.response.start",
+            "status": 302,
+            "headers": [
+                [b"Location", b"/foo"],
+                [b"content-type", b"text/plain"],
+                [b"set-cookie", b"foo=bar; HttpOnly; Max-Age=10; Path=/; SameSite=lax"],
+            ],
+        },
+        {"type": "http.response.body", "body": b""},
+    ] == events