From 008e2f63c217aa066027a872ee706b07bd084857 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Tue, 9 Jun 2020 15:19:37 -0700 Subject: [PATCH] response.set_cookie(), closes #795 --- datasette/actor_auth_cookie.py | 1 - datasette/app.py | 15 ++------- datasette/utils/asgi.py | 53 +++++++++++++++++++++++++++++--- datasette/views/special.py | 14 ++------- docs/internals.rst | 30 ++++++++++++++++++ tests/test_internals_response.py | 26 ++++++++++++++++ 6 files changed, 108 insertions(+), 31 deletions(-) diff --git a/datasette/actor_auth_cookie.py b/datasette/actor_auth_cookie.py index f3a0f306..a2aa6889 100644 --- a/datasette/actor_auth_cookie.py +++ b/datasette/actor_auth_cookie.py @@ -1,6 +1,5 @@ from datasette import hookimpl from itsdangerous import BadSignature -from http.cookies import SimpleCookie @hookimpl diff --git a/datasette/app.py b/datasette/app.py index 633ca4fe..71fa9afb 100644 --- a/datasette/app.py +++ b/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() - 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() + response.set_cookie("ds_messages", self.sign(request._messages, "messages")) elif getattr(request, "_messages_should_clear", False): - cookie = SimpleCookie() - 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() + response.set_cookie("ds_messages", "", expires=0, max_age=0) def _show_messages(self, request): if getattr(request, "_messages", None): diff --git a/datasette/utils/asgi.py b/datasette/utils/asgi.py index cdd6b148..5a152570 100644 --- a/datasette/utils/asgi.py +++ b/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": [ - [key.encode("utf-8"), value.encode("utf-8")] - for key, value in headers.items() - ], + "headers": raw_headers, } ) 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( diff --git a/datasette/views/special.py b/datasette/views/special.py index 7a5fbe21..7f4284a1 100644 --- a/datasette/views/special.py +++ b/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() - 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(), - }, - ) + response = Response.redirect("/") + response.set_cookie("ds_actor", self.ds.sign({"id": "root"}, "actor")) return response else: return Response("Invalid token", status=403) diff --git a/docs/internals.rst b/docs/internals.rst index d92c985f..7978e3d7 100644 --- a/docs/internals.rst +++ b/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() ` to set signed cookies. Here's how you would set the ``ds_actor`` cookie for use with Datasette :ref:`authentication `: + +.. code-block:: python + + response = Response.redirect("/") + response.set_cookie("ds_actor", datasette.sign({"id": "cleopaws"}, "actor")) + return response + .. _internals_datasette: Datasette class diff --git a/tests/test_internals_response.py b/tests/test_internals_response.py index 7c11f858..820b20b2 100644 --- a/tests/test_internals_response.py +++ b/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