From 54afac769c44054d7b9458c8fe0e42278ec70766 Mon Sep 17 00:00:00 2001 From: Dmitry Erlikh Date: Sat, 14 Nov 2020 14:11:28 +0100 Subject: [PATCH] Http exception cookies support (#5197) --- CHANGES/4277.feature | 1 + aiohttp/helpers.py | 92 ++++++++++++++++++++++++++++++++- aiohttp/web_exceptions.py | 4 +- aiohttp/web_protocol.py | 1 + aiohttp/web_response.py | 93 ++++----------------------------- docs/web_exceptions.rst | 80 +++++++++++++++++++++++++++++ tests/test_helpers.py | 99 +++++++++++++++++++++++++++++++++++- tests/test_web_functional.py | 13 +++++ tests/test_web_response.py | 87 ------------------------------- 9 files changed, 298 insertions(+), 172 deletions(-) create mode 100644 CHANGES/4277.feature diff --git a/CHANGES/4277.feature b/CHANGES/4277.feature new file mode 100644 index 00000000000..bc5b360bf73 --- /dev/null +++ b/CHANGES/4277.feature @@ -0,0 +1 @@ +Add set_cookie and del_cookie methods to HTTPException diff --git a/aiohttp/helpers.py b/aiohttp/helpers.py index 6ace5137e56..f197db4fba2 100644 --- a/aiohttp/helpers.py +++ b/aiohttp/helpers.py @@ -16,6 +16,7 @@ import weakref from collections import namedtuple from contextlib import suppress +from http.cookies import SimpleCookie from math import ceil from pathlib import Path from types import TracebackType @@ -44,7 +45,7 @@ import async_timeout import attr -from multidict import MultiDict, MultiDictProxy +from multidict import CIMultiDict, MultiDict, MultiDictProxy from typing_extensions import Protocol, final from yarl import URL @@ -705,6 +706,7 @@ class HeadersMixin: __slots__ = ("_content_type", "_content_dict", "_stored_content_type") def __init__(self) -> None: + super().__init__() self._content_type = None # type: Optional[str] self._content_dict = None # type: Optional[Dict[str, str]] self._stored_content_type = sentinel @@ -799,3 +801,91 @@ def __bool__(self) -> bool: def __repr__(self) -> str: content = ", ".join(map(repr, self._maps)) return f"ChainMapProxy({content})" + + +class CookieMixin: + def __init__(self) -> None: + super().__init__() + self._cookies = SimpleCookie() # type: SimpleCookie[str] + + @property + def cookies(self) -> "SimpleCookie[str]": + return self._cookies + + def set_cookie( + self, + name: str, + value: str, + *, + expires: Optional[str] = None, + domain: Optional[str] = None, + max_age: Optional[Union[int, str]] = None, + path: str = "/", + secure: Optional[bool] = None, + httponly: Optional[bool] = None, + version: Optional[str] = None, + samesite: Optional[str] = None, + ) -> None: + """Set or update response cookie. + + Sets new cookie or updates existent with new value. + Also updates only those params which are not None. + """ + + old = self._cookies.get(name) + if old is not None and old.coded_value == "": + # deleted cookie + self._cookies.pop(name, None) + + self._cookies[name] = value + c = self._cookies[name] + + if expires is not None: + c["expires"] = expires + elif c.get("expires") == "Thu, 01 Jan 1970 00:00:00 GMT": + del c["expires"] + + if domain is not None: + c["domain"] = domain + + if max_age is not None: + c["max-age"] = str(max_age) + elif "max-age" in c: + del c["max-age"] + + c["path"] = path + + if secure is not None: + c["secure"] = secure + if httponly is not None: + c["httponly"] = httponly + if version is not None: + c["version"] = version + if samesite is not None: + c["samesite"] = samesite + + def del_cookie( + self, name: str, *, domain: Optional[str] = None, path: str = "/" + ) -> None: + """Delete cookie. + + Creates new empty expired cookie. + """ + # TODO: do we need domain/path here? + self._cookies.pop(name, None) + self.set_cookie( + name, + "", + max_age=0, + expires="Thu, 01 Jan 1970 00:00:00 GMT", + domain=domain, + path=path, + ) + + +def populate_with_cookies( + headers: "CIMultiDict[str]", cookies: "SimpleCookie[str]" +) -> None: + for cookie in cookies.values(): + value = cookie.output(header="")[1:] + headers.add(hdrs.SET_COOKIE, value) diff --git a/aiohttp/web_exceptions.py b/aiohttp/web_exceptions.py index 573f61afc3e..28fe0aafcc5 100644 --- a/aiohttp/web_exceptions.py +++ b/aiohttp/web_exceptions.py @@ -6,6 +6,7 @@ from yarl import URL from . import hdrs +from .helpers import CookieMixin from .typedefs import LooseHeaders, StrOrURL __all__ = ( @@ -75,7 +76,7 @@ ############################################################ -class HTTPException(Exception): +class HTTPException(CookieMixin, Exception): # You should set in subclasses: # status = 200 @@ -92,6 +93,7 @@ def __init__( text: Optional[str] = None, content_type: Optional[str] = None, ) -> None: + super().__init__() if reason is None: reason = self.default_reason diff --git a/aiohttp/web_protocol.py b/aiohttp/web_protocol.py index 8a99aba71f9..8e54ed50758 100644 --- a/aiohttp/web_protocol.py +++ b/aiohttp/web_protocol.py @@ -460,6 +460,7 @@ async def _handle_request( resp = Response( status=exc.status, reason=exc.reason, text=exc.text, headers=exc.headers ) + resp._cookies = exc._cookies reset = await self.finish_response(request, resp, start_time) except asyncio.CancelledError: raise diff --git a/aiohttp/web_response.py b/aiohttp/web_response.py index db9c36243b1..ac3ee266243 100644 --- a/aiohttp/web_response.py +++ b/aiohttp/web_response.py @@ -9,7 +9,7 @@ import zlib from concurrent.futures import Executor from email.utils import parsedate -from http.cookies import Morsel, SimpleCookie +from http.cookies import Morsel from typing import ( TYPE_CHECKING, Any, @@ -27,7 +27,14 @@ from . import hdrs, payload from .abc import AbstractStreamWriter -from .helpers import PY_38, HeadersMixin, rfc822_formatted_time, sentinel +from .helpers import ( + PY_38, + CookieMixin, + HeadersMixin, + populate_with_cookies, + rfc822_formatted_time, + sentinel, +) from .http import RESPONSES, SERVER_SOFTWARE, HttpVersion10, HttpVersion11 from .payload import Payload from .typedefs import JSONEncoder, LooseHeaders @@ -64,7 +71,7 @@ class ContentCoding(enum.Enum): ############################################################ -class StreamResponse(BaseClass, HeadersMixin): +class StreamResponse(BaseClass, HeadersMixin, CookieMixin): __slots__ = ( "_length_check", @@ -73,7 +80,6 @@ class StreamResponse(BaseClass, HeadersMixin): "_chunked", "_compression", "_compression_force", - "_cookies", "_req", "_payload_writer", "_eof_sent", @@ -99,7 +105,6 @@ def __init__( self._chunked = False self._compression = False self._compression_force = None # type: Optional[ContentCoding] - self._cookies = SimpleCookie() # type: SimpleCookie[str] self._req = None # type: Optional[BaseRequest] self._payload_writer = None # type: Optional[AbstractStreamWriter] @@ -185,80 +190,6 @@ def enable_compression(self, force: Optional[ContentCoding] = None) -> None: def headers(self) -> "CIMultiDict[str]": return self._headers - @property - def cookies(self) -> "SimpleCookie[str]": - return self._cookies - - def set_cookie( - self, - name: str, - value: str, - *, - expires: Optional[str] = None, - domain: Optional[str] = None, - max_age: Optional[Union[int, str]] = None, - path: str = "/", - secure: Optional[bool] = None, - httponly: Optional[bool] = None, - version: Optional[str] = None, - samesite: Optional[str] = None, - ) -> None: - """Set or update response cookie. - - Sets new cookie or updates existent with new value. - Also updates only those params which are not None. - """ - - old = self._cookies.get(name) - if old is not None and old.coded_value == "": - # deleted cookie - self._cookies.pop(name, None) - - self._cookies[name] = value - c = self._cookies[name] - - if expires is not None: - c["expires"] = expires - elif c.get("expires") == "Thu, 01 Jan 1970 00:00:00 GMT": - del c["expires"] - - if domain is not None: - c["domain"] = domain - - if max_age is not None: - c["max-age"] = str(max_age) - elif "max-age" in c: - del c["max-age"] - - c["path"] = path - - if secure is not None: - c["secure"] = secure - if httponly is not None: - c["httponly"] = httponly - if version is not None: - c["version"] = version - if samesite is not None: - c["samesite"] = samesite - - def del_cookie( - self, name: str, *, domain: Optional[str] = None, path: str = "/" - ) -> None: - """Delete cookie. - - Creates new empty expired cookie. - """ - # TODO: do we need domain/path here? - self._cookies.pop(name, None) - self.set_cookie( - name, - "", - max_age=0, - expires="Thu, 01 Jan 1970 00:00:00 GMT", - domain=domain, - path=path, - ) - @property def content_length(self) -> Optional[int]: # Just a placeholder for adding setter @@ -399,9 +330,7 @@ async def _prepare_headers(self) -> None: version = request.version headers = self._headers - for cookie in self._cookies.values(): - value = cookie.output(header="")[1:] - headers.add(hdrs.SET_COOKIE, value) + populate_with_cookies(headers, self.cookies) if self._compression: await self._start_compression(request) diff --git a/docs/web_exceptions.rst b/docs/web_exceptions.rst index 9e5da48fcbd..14d522479db 100644 --- a/docs/web_exceptions.rst +++ b/docs/web_exceptions.rst @@ -142,6 +142,86 @@ Base HTTP Exception HTTP headers for the exception, :class:`multidict.CIMultiDict` + .. attribute:: cookies + + An instance of :class:`http.cookies.SimpleCookie` for *outgoing* cookies. + + .. versionadded:: 4.0 + + .. method:: set_cookie(name, value, *, path='/', expires=None, \ + domain=None, max_age=None, \ + secure=None, httponly=None, version=None, \ + samesite=None) + + Convenient way for setting :attr:`cookies`, allows to specify + some additional properties like *max_age* in a single call. + + .. versionadded:: 4.0 + + :param str name: cookie name + + :param str value: cookie value (will be converted to + :class:`str` if value has another type). + + :param expires: expiration date (optional) + + :param str domain: cookie domain (optional) + + :param int max_age: defines the lifetime of the cookie, in + seconds. The delta-seconds value is a + decimal non- negative integer. After + delta-seconds seconds elapse, the client + should discard the cookie. A value of zero + means the cookie should be discarded + immediately. (optional) + + :param str path: specifies the subset of URLs to + which this cookie applies. (optional, ``'/'`` by default) + + :param bool secure: attribute (with no value) directs + the user agent to use only (unspecified) + secure means to contact the origin server + whenever it sends back this cookie. + The user agent (possibly under the user's + control) may determine what level of + security it considers appropriate for + "secure" cookies. The *secure* should be + considered security advice from the server + to the user agent, indicating that it is in + the session's interest to protect the cookie + contents. (optional) + + :param bool httponly: ``True`` if the cookie HTTP only (optional) + + :param int version: a decimal integer, identifies to which + version of the state management + specification the cookie + conforms. (Optional, *version=1* by default) + + :param str samesite: Asserts that a cookie must not be sent with + cross-origin requests, providing some protection + against cross-site request forgery attacks. + Generally the value should be one of: ``None``, + ``Lax`` or ``Strict``. (optional) + + .. warning:: + + In HTTP version 1.1, ``expires`` was deprecated and replaced with + the easier-to-use ``max-age``, but Internet Explorer (IE6, IE7, + and IE8) **does not** support ``max-age``. + + .. method:: del_cookie(name, *, path='/', domain=None) + + Deletes cookie. + + .. versionadded:: 4.0 + + :param str name: cookie name + + :param str domain: optional cookie domain + + :param str path: optional cookie path, ``'/'`` by default + Successful Exceptions --------------------- diff --git a/tests/test_helpers.py b/tests/test_helpers.py index d8db499263c..ce243a00ef4 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -7,7 +7,8 @@ from unittest import mock import pytest -from multidict import MultiDict +from multidict import CIMultiDict, MultiDict +from re_assert import Matches from yarl import URL from aiohttp import helpers @@ -654,3 +655,99 @@ def test_is_expected_content_type_non_json_not_match(): assert not is_expected_content_type( response_content_type=response_ct, expected_content_type=expected_ct ) + + +def test_cookies_mixin(): + sut = helpers.CookieMixin() + + assert sut.cookies == {} + assert str(sut.cookies) == "" + + sut.set_cookie("name", "value") + assert str(sut.cookies) == "Set-Cookie: name=value; Path=/" + sut.set_cookie("name", "other_value") + assert str(sut.cookies) == "Set-Cookie: name=other_value; Path=/" + + sut.cookies["name"] = "another_other_value" + sut.cookies["name"]["max-age"] = 10 + assert ( + str(sut.cookies) == "Set-Cookie: name=another_other_value; Max-Age=10; Path=/" + ) + + sut.del_cookie("name") + expected = ( + 'Set-Cookie: name=("")?; ' + "expires=Thu, 01 Jan 1970 00:00:00 GMT; Max-Age=0; Path=/" + ) + assert Matches(expected) == str(sut.cookies) + + sut.set_cookie("name", "value", domain="local.host") + expected = "Set-Cookie: name=value; Domain=local.host; Path=/" + assert str(sut.cookies) == expected + + +def test_cookies_mixin_path(): + sut = helpers.CookieMixin() + + assert sut.cookies == {} + + sut.set_cookie("name", "value", path="/some/path") + assert str(sut.cookies) == "Set-Cookie: name=value; Path=/some/path" + sut.set_cookie("name", "value", expires="123") + assert str(sut.cookies) == "Set-Cookie: name=value; expires=123; Path=/" + sut.set_cookie( + "name", + "value", + domain="example.com", + path="/home", + expires="123", + max_age="10", + secure=True, + httponly=True, + version="2.0", + samesite="lax", + ) + assert ( + str(sut.cookies).lower() == "set-cookie: name=value; " + "domain=example.com; " + "expires=123; " + "httponly; " + "max-age=10; " + "path=/home; " + "samesite=lax; " + "secure; " + "version=2.0" + ) + + +def test_sutonse_cookie__issue_del_cookie(): + sut = helpers.CookieMixin() + + assert sut.cookies == {} + assert str(sut.cookies) == "" + + sut.del_cookie("name") + expected = ( + 'Set-Cookie: name=("")?; ' + "expires=Thu, 01 Jan 1970 00:00:00 GMT; Max-Age=0; Path=/" + ) + assert Matches(expected) == str(sut.cookies) + + +def test_cookie_set_after_del(): + sut = helpers.CookieMixin() + + sut.del_cookie("name") + sut.set_cookie("name", "val") + # check for Max-Age dropped + expected = "Set-Cookie: name=val; Path=/" + assert str(sut.cookies) == expected + + +def test_populate_with_cookies(): + cookies_mixin = helpers.CookieMixin() + cookies_mixin.set_cookie("name", "value") + headers = CIMultiDict() + + helpers.populate_with_cookies(headers, cookies_mixin.cookies) + assert headers == CIMultiDict({"Set-Cookie": "name=value; Path=/"}) diff --git a/tests/test_web_functional.py b/tests/test_web_functional.py index f81f1790851..754b5d31b49 100644 --- a/tests/test_web_functional.py +++ b/tests/test_web_functional.py @@ -1910,3 +1910,16 @@ async def handler(_): resp = await client.get("/") assert CONTENT_TYPE not in resp.headers assert TRANSFER_ENCODING not in resp.headers + + +async def test_httpfound_cookies_302(aiohttp_client): + async def handler(_): + resp = web.HTTPFound("/") + resp.set_cookie("my-cookie", "cookie-value") + raise resp + + app = web.Application() + app.router.add_get("/", handler) + client = await aiohttp_client(app) + resp = await client.get("/", allow_redirects=False) + assert "my-cookie" in resp.cookies diff --git a/tests/test_web_response.py b/tests/test_web_response.py index 9998063a369..74994e03605 100644 --- a/tests/test_web_response.py +++ b/tests/test_web_response.py @@ -618,93 +618,6 @@ def test_force_close() -> None: assert resp.keep_alive is False -def test_response_cookies() -> None: - resp = StreamResponse() - - assert resp.cookies == {} - assert str(resp.cookies) == "" - - resp.set_cookie("name", "value") - assert str(resp.cookies) == "Set-Cookie: name=value; Path=/" - resp.set_cookie("name", "other_value") - assert str(resp.cookies) == "Set-Cookie: name=other_value; Path=/" - - resp.cookies["name"] = "another_other_value" - resp.cookies["name"]["max-age"] = 10 - assert ( - str(resp.cookies) == "Set-Cookie: name=another_other_value; Max-Age=10; Path=/" - ) - - resp.del_cookie("name") - expected = ( - 'Set-Cookie: name=("")?; ' - "expires=Thu, 01 Jan 1970 00:00:00 GMT; Max-Age=0; Path=/" - ) - assert Matches(expected) == str(resp.cookies) - - resp.set_cookie("name", "value", domain="local.host") - expected = "Set-Cookie: name=value; Domain=local.host; Path=/" - assert str(resp.cookies) == expected - - -def test_response_cookie_path() -> None: - resp = StreamResponse() - - assert resp.cookies == {} - - resp.set_cookie("name", "value", path="/some/path") - assert str(resp.cookies) == "Set-Cookie: name=value; Path=/some/path" - resp.set_cookie("name", "value", expires="123") - assert str(resp.cookies) == "Set-Cookie: name=value; expires=123; Path=/" - resp.set_cookie( - "name", - "value", - domain="example.com", - path="/home", - expires="123", - max_age="10", - secure=True, - httponly=True, - version="2.0", - samesite="lax", - ) - assert ( - str(resp.cookies).lower() == "set-cookie: name=value; " - "domain=example.com; " - "expires=123; " - "httponly; " - "max-age=10; " - "path=/home; " - "samesite=lax; " - "secure; " - "version=2.0" - ) - - -def test_response_cookie__issue_del_cookie() -> None: - resp = StreamResponse() - - assert resp.cookies == {} - assert str(resp.cookies) == "" - - resp.del_cookie("name") - expected = ( - 'Set-Cookie: name=("")?; ' - "expires=Thu, 01 Jan 1970 00:00:00 GMT; Max-Age=0; Path=/" - ) - assert Matches(expected) == str(resp.cookies) - - -def test_cookie_set_after_del() -> None: - resp = StreamResponse() - - resp.del_cookie("name") - resp.set_cookie("name", "val") - # check for Max-Age dropped - expected = "Set-Cookie: name=val; Path=/" - assert str(resp.cookies) == expected - - def test_set_status_with_reason() -> None: resp = StreamResponse()