From 29695bd24f35542204f9c3cf6784e127f0c7b2e2 Mon Sep 17 00:00:00 2001
From: Andrew Svetlov <andrew.svetlov@gmail.com>
Date: Wed, 29 Jul 2015 23:41:16 +0300
Subject: [PATCH] Return back to passing response parameter into remember and
 forget methods

---
 aiohttp_security/abc.py              | 10 +++++-----
 aiohttp_security/api.py              | 10 ++++------
 aiohttp_security/cookies_identity.py | 28 ++++++++++------------------
 3 files changed, 19 insertions(+), 29 deletions(-)

diff --git a/aiohttp_security/abc.py b/aiohttp_security/abc.py
index df064b8..dd038ab 100644
--- a/aiohttp_security/abc.py
+++ b/aiohttp_security/abc.py
@@ -9,16 +9,16 @@ class AbstractIdentityPolicy(metaclass=abc.ABCMeta):
     @asyncio.coroutine
     @abc.abstractmethod
     def identify(self, request):
-        """ Return the claimed identity of the user associated request or
+        """Return the claimed identity of the user associated request or
         ``None`` if no identity can be found associated with the request."""
         pass
 
     @asyncio.coroutine
     @abc.abstractmethod
-    def remember(self, request, identity, **kwargs):
+    def remember(self, request, response, identity, **kwargs):
         """Remember identity.
 
-        Return MultiDict with headers on this request's response.
+        Modify response object by filling it's headers with remembered user.
 
         An individual identity policy and its consumers can decide on
         the composition and meaning of **kw.
@@ -27,8 +27,8 @@ class AbstractIdentityPolicy(metaclass=abc.ABCMeta):
 
     @asyncio.coroutine
     @abc.abstractmethod
-    def forget(self, request):
-        """ Modify request.response which can be used to 'forget' the
+    def forget(self, request, response):
+        """ Modify response which can be used to 'forget' the
         current identity on subsequent requests."""
         pass
 
diff --git a/aiohttp_security/api.py b/aiohttp_security/api.py
index b7b8e49..cfbe35e 100644
--- a/aiohttp_security/api.py
+++ b/aiohttp_security/api.py
@@ -7,17 +7,15 @@ AUTZ_KEY = 'aiohttp_security_autz_policy'
 
 
 @asyncio.coroutine
-def remember(request, identity, **kwargs):
+def remember(request, response, identity, **kwargs):
     identity_policy = request.app[IDENTITY_KEY]
-    headers = yield from identity_policy.remember(request, identity, **kwargs)
-    return headers
+    yield from identity_policy.remember(request, response, identity, **kwargs)
 
 
 @asyncio.coroutine
-def forget(request):
+def forget(request, response):
     identity_policy = request.app[IDENTITY_KEY]
-    headers = yield from identity_policy.forget(request)
-    return headers
+    yield from identity_policy.forget(request, response)
 
 
 @asyncio.coroutine
diff --git a/aiohttp_security/cookies_identity.py b/aiohttp_security/cookies_identity.py
index cf829f6..ddfde0a 100644
--- a/aiohttp_security/cookies_identity.py
+++ b/aiohttp_security/cookies_identity.py
@@ -6,6 +6,9 @@ from aiohttp import hdrs, CIMultiDict
 from .abc import AbstractIdentityPolicy
 
 
+sentinel = object()
+
+
 class CookiesIdentityPolicy(AbstractIdentityPolicy):
 
     def __init__(self):
@@ -18,23 +21,12 @@ class CookiesIdentityPolicy(AbstractIdentityPolicy):
         return identity
 
     @asyncio.coroutine
-    def remember(self, request, identity, **kwargs):
-        cookies = http.cookies.SimpleCookie()
-        max_age = kwargs.pop('max_age', self._max_age)
-        cookies[self._cookie_name] = identity
-        cookie = cookies[self._cookie_name]
-        cookie['max-age'] = max_age
-        cookie.update(kwargs)
-
-        value = cookie.output(header='')[1:]
-        result = CIMultiDict({hdrs.SET_COOKIE: value})
-        return result
+    def remember(self, request, response, identity, max_age=sentinel,
+                 **kwargs):
+        if max_age is sentinel:
+            max_age = self._max_age
+        response.set_cookie(self._cookie_name, max_age=max_age, **kwargs)
 
     @asyncio.coroutine
-    def forget(self, request):
-        cookies = http.cookies.SimpleCookie()
-        cookies[self._cookie_name] = ''
-        cookie = cookies[self._cookie_name]
-        value = cookie.output(header='')[1:]
-        result = CIMultiDict({hdrs.SET_COOKIE: value})
-        return result
+    def forget(self, request, response):
+        response.del_cookie(self._cookie_name)