feat(event_handler): add cookies as 1st class citizen in v2 (aws-powertools#1487)

Co-authored-by: Heitor Lessa <[email protected]>
Co-authored-by: Heitor Lessa <[email protected]>

Author: rubenfonseca

aws_lambda_powertools/event_handler/api_gateway.py

@@ -15,6 +15,7 @@
 from aws_lambda_powertools.event_handler import content_types
 from aws_lambda_powertools.event_handler.exceptions import NotFoundError, ServiceError
 from aws_lambda_powertools.shared import constants
+from aws_lambda_powertools.shared.cookies import Cookie
 from aws_lambda_powertools.shared.functions import resolve_truthy_env_var_choice
 from aws_lambda_powertools.shared.json_encoder import Encoder
 from aws_lambda_powertools.utilities.data_classes import (
@@ -147,7 +148,7 @@ def __init__(
         content_type: Optional[str],
         body: Union[str, bytes, None],
         headers: Optional[Dict[str, Union[str, List[str]]]] = None,
-        cookies: Optional[List[str]] = None,
+        cookies: Optional[List[Cookie]] = None,
     ):
         """
 
@@ -162,7 +163,7 @@ def __init__(
             Optionally set the response body. Note: bytes body will be automatically base64 encoded
         headers: dict[str, Union[str, List[str]]]
             Optionally set specific http headers. Setting "Content-Type" here would override the `content_type` value.
-        cookies: list[str]
+        cookies: list[Cookie]
             Optionally set cookies.
         """
         self.status_code = status_code

aws_lambda_powertools/shared/cookies.py

@@ -0,0 +1,118 @@
+from datetime import datetime
+from enum import Enum
+from io import StringIO
+from typing import List, Optional
+
+
+class SameSite(Enum):
+    """
+    SameSite allows a server to define a cookie attribute making it impossible for
+    the browser to send this cookie along with cross-site requests. The main
+    goal is to mitigate the risk of cross-origin information leakage, and provide
+    some protection against cross-site request forgery attacks.
+
+    See https://tools.ietf.org/html/draft-ietf-httpbis-cookie-same-site-00 for details.
+    """
+
+    DEFAULT_MODE = ""
+    LAX_MODE = "Lax"
+    STRICT_MODE = "Strict"
+    NONE_MODE = "None"
+
+
+def _format_date(timestamp: datetime) -> str:
+    # Specification example: Wed, 21 Oct 2015 07:28:00 GMT
+    return timestamp.strftime("%a, %d %b %Y %H:%M:%S GMT")
+
+
+class Cookie:
+    """
+    A Cookie represents an HTTP cookie as sent in the Set-Cookie header of an
+    HTTP response or the Cookie header of an HTTP request.
+
+    See https://tools.ietf.org/html/rfc6265 for details.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        value: str,
+        path: str = "",
+        domain: str = "",
+        secure: bool = True,
+        http_only: bool = False,
+        max_age: Optional[int] = None,
+        expires: Optional[datetime] = None,
+        same_site: Optional[SameSite] = None,
+        custom_attributes: Optional[List[str]] = None,
+    ):
+        """
+
+        Parameters
+        ----------
+        name: str
+            The name of this cookie, for example session_id
+        value: str
+            The cookie value, for instance an uuid
+        path: str
+            The path for which this cookie is valid. Optional
+        domain: str
+            The domain for which this cookie is valid. Optional
+        secure: bool
+            Marks the cookie as secure, only sendable to the server with an encrypted request over the HTTPS protocol
+        http_only: bool
+            Enabling this attribute makes the cookie inaccessible to the JavaScript `Document.cookie` API
+        max_age: Optional[int]
+            Defines the period of time after which the cookie is invalid. Use negative values to force cookie deletion.
+        expires: Optional[datetime]
+            Defines a date where the permanent cookie expires.
+        same_site: Optional[SameSite]
+            Determines if the cookie should be sent to third party websites
+        custom_attributes: Optional[List[str]]
+            List of additional custom attributes to set on the cookie
+        """
+        self.name = name
+        self.value = value
+        self.path = path
+        self.domain = domain
+        self.secure = secure
+        self.expires = expires
+        self.max_age = max_age
+        self.http_only = http_only
+        self.same_site = same_site
+        self.custom_attributes = custom_attributes
+
+    def __str__(self) -> str:
+        payload = StringIO()
+        payload.write(f"{self.name}={self.value}")
+
+        if self.path:
+            payload.write(f"; Path={self.path}")
+
+        if self.domain:
+            payload.write(f"; Domain={self.domain}")
+
+        if self.expires:
+            payload.write(f"; Expires={_format_date(self.expires)}")
+
+        if self.max_age:
+            if self.max_age > 0:
+                payload.write(f"; MaxAge={self.max_age}")
+            else:
+                # negative or zero max-age should be set to 0
+                payload.write("; MaxAge=0")
+
+        if self.http_only:
+            payload.write("; HttpOnly")
+
+        if self.secure:
+            payload.write("; Secure")
+
+        if self.same_site:
+            payload.write(f"; SameSite={self.same_site.value}")
+
+        if self.custom_attributes:
+            for attr in self.custom_attributes:
+                payload.write(f"; {attr}")
+
+        return payload.getvalue()

aws_lambda_powertools/shared/headers_serializer.py

@@ -2,14 +2,16 @@
 from collections import defaultdict
 from typing import Any, Dict, List, Union
 
+from aws_lambda_powertools.shared.cookies import Cookie
+
 
 class BaseHeadersSerializer:
     """
     Helper class to correctly serialize headers and cookies for Amazon API Gateway,
     ALB and Lambda Function URL response payload.
     """
 
-    def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str]) -> Dict[str, Any]:
+    def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Cookie]) -> Dict[str, Any]:
         """
         Serializes headers and cookies according to the request type.
         Returns a dict that can be merged with the response payload.
@@ -25,7 +27,7 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str
 
 
 class HttpApiHeadersSerializer(BaseHeadersSerializer):
-    def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str]) -> Dict[str, Any]:
+    def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Cookie]) -> Dict[str, Any]:
         """
         When using HTTP APIs or LambdaFunctionURLs, everything is taken care automatically for us.
         We can directly assign a list of cookies and a dict of headers to the response payload, and the
@@ -44,11 +46,11 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str
             else:
                 combined_headers[key] = ", ".join(values)
 
-        return {"headers": combined_headers, "cookies": cookies}
+        return {"headers": combined_headers, "cookies": list(map(str, cookies))}
 
 
 class MultiValueHeadersSerializer(BaseHeadersSerializer):
-    def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str]) -> Dict[str, Any]:
+    def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Cookie]) -> Dict[str, Any]:
         """
         When using REST APIs, headers can be encoded using the `multiValueHeaders` key on the response.
         This is also the case when using an ALB integration with the `multiValueHeaders` option enabled.
@@ -69,13 +71,13 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str
         if cookies:
             payload.setdefault("Set-Cookie", [])
             for cookie in cookies:
-                payload["Set-Cookie"].append(cookie)
+                payload["Set-Cookie"].append(str(cookie))
 
         return {"multiValueHeaders": payload}
 
 
 class SingleValueHeadersSerializer(BaseHeadersSerializer):
-    def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str]) -> Dict[str, Any]:
+    def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Cookie]) -> Dict[str, Any]:
         """
         The ALB integration has `multiValueHeaders` disabled by default.
         If we try to set multiple headers with the same key, or more than one cookie, print a warning.
@@ -93,7 +95,7 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str
                 )
 
             # We can only send one cookie, send the last one
-            payload["headers"]["Set-Cookie"] = cookies[-1]
+            payload["headers"]["Set-Cookie"] = str(cookies[-1])
 
         for key, values in headers.items():
             if isinstance(values, str):

docs/core/event_handler/api_gateway.md

@@ -323,7 +323,7 @@ You can use the `Response` class to have full control over the response. For exa
 
 === "fine_grained_responses.py"
 
-    ```python hl_lines="7 24-29"
+    ```python hl_lines="7 25-30"
     --8<-- "examples/event_handler_rest/src/fine_grained_responses.py"
     ```
 

docs/upgrade.md

@@ -53,7 +53,7 @@ def get_todos():
     return Response(
         # ...
         headers={"Content-Type": ["text/plain"]},
-        cookies=["CookieName=CookieValue"]
+        cookies=[Cookie(name="session_id", value="12345", secure=True, http_only=True)],
     )
 ```
 

examples/event_handler_rest/src/fine_grained_responses.py

@@ -6,6 +6,7 @@
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.event_handler import APIGatewayRestResolver, Response, content_types
 from aws_lambda_powertools.logging import correlation_paths
+from aws_lambda_powertools.shared.cookies import Cookie
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 tracer = Tracer()
@@ -26,7 +27,7 @@ def get_todos():
         content_type=content_types.APPLICATION_JSON,
         body=todos.json()[:10],
         headers=custom_headers,
-        cookies=["<cookie-name>=<cookie-value>; Secure; Expires=<date>"],
+        cookies=[Cookie(name="session_id", value="12345")],
     )
 
 

examples/event_handler_rest/src/fine_grained_responses_output.json

@@ -3,7 +3,7 @@
     "multiValueHeaders": {
         "Content-Type": ["application/json"],
         "X-Transaction-Id": ["3490eea9-791b-47a0-91a4-326317db61a9"],
-        "Set-Cookie": ["<cookie-name>=<cookie-value>; Secure; Expires=<date>"]
+        "Set-Cookie": ["session_id=12345; Secure"]
     },
     "body": "{\"todos\":[{\"userId\":1,\"id\":1,\"title\":\"delectus aut autem\",\"completed\":false},{\"userId\":1,\"id\":2,\"title\":\"quis ut nam facilis et officia qui\",\"completed\":false},{\"userId\":1,\"id\":3,\"title\":\"fugiat veniam minus\",\"completed\":false},{\"userId\":1,\"id\":4,\"title\":\"et porro tempora\",\"completed\":true},{\"userId\":1,\"id\":5,\"title\":\"laboriosam mollitia et enim quasi adipisci quia provident illum\",\"completed\":false},{\"userId\":1,\"id\":6,\"title\":\"qui ullam ratione quibusdam voluptatem quia omnis\",\"completed\":false},{\"userId\":1,\"id\":7,\"title\":\"illo expedita consequatur quia in\",\"completed\":false},{\"userId\":1,\"id\":8,\"title\":\"quo adipisci enim quam ut ab\",\"completed\":true},{\"userId\":1,\"id\":9,\"title\":\"molestiae perspiciatis ipsa\",\"completed\":false},{\"userId\":1,\"id\":10,\"title\":\"illo est ratione doloremque quia maiores aut\",\"completed\":true}]}",
     "isBase64Encoded": false

tests/e2e/event_handler/handlers/alb_handler.py

@@ -3,14 +3,22 @@
 app = ALBResolver()
 
 
-@app.get("/todos")
+@app.post("/todos")
 def hello():
+    payload = app.current_event.json_body
+
+    body = payload.get("body", "Hello World")
+    status_code = payload.get("status_code", 200)
+    headers = payload.get("headers", {})
+    cookies = payload.get("cookies", [])
+    content_type = headers.get("Content-Type", content_types.TEXT_PLAIN)
+
     return Response(
-        status_code=200,
-        content_type=content_types.TEXT_PLAIN,
-        body="Hello world",
-        cookies=["CookieMonster", "MonsterCookie"],
-        headers={"Foo": ["bar", "zbr"]},
+        status_code=status_code,
+        content_type=content_type,
+        body=body,
+        cookies=cookies,
+        headers=headers,
     )
 
 

tests/e2e/event_handler/handlers/api_gateway_http_handler.py

@@ -3,14 +3,22 @@
 app = APIGatewayHttpResolver()
 
 
-@app.get("/todos")
+@app.post("/todos")
 def hello():
+    payload = app.current_event.json_body
+
+    body = payload.get("body", "Hello World")
+    status_code = payload.get("status_code", 200)
+    headers = payload.get("headers", {})
+    cookies = payload.get("cookies", [])
+    content_type = headers.get("Content-Type", content_types.TEXT_PLAIN)
+
     return Response(
-        status_code=200,
-        content_type=content_types.TEXT_PLAIN,
-        body="Hello world",
-        cookies=["CookieMonster", "MonsterCookie"],
-        headers={"Foo": ["bar", "zbr"]},
+        status_code=status_code,
+        content_type=content_type,
+        body=body,
+        cookies=cookies,
+        headers=headers,
     )
 
 

tests/e2e/event_handler/handlers/api_gateway_rest_handler.py

@@ -3,14 +3,22 @@
 app = APIGatewayRestResolver()
 
 
-@app.get("/todos")
+@app.post("/todos")
 def hello():
+    payload = app.current_event.json_body
+
+    body = payload.get("body", "Hello World")
+    status_code = payload.get("status_code", 200)
+    headers = payload.get("headers", {})
+    cookies = payload.get("cookies", [])
+    content_type = headers.get("Content-Type", content_types.TEXT_PLAIN)
+
     return Response(
-        status_code=200,
-        content_type=content_types.TEXT_PLAIN,
-        body="Hello world",
-        cookies=["CookieMonster", "MonsterCookie"],
-        headers={"Foo": ["bar", "zbr"]},
+        status_code=status_code,
+        content_type=content_type,
+        body=body,
+        cookies=cookies,
+        headers=headers,
     )
 
 

tests/e2e/event_handler/handlers/lambda_function_url_handler.py

@@ -3,14 +3,22 @@
 app = LambdaFunctionUrlResolver()
 
 
-@app.get("/todos")
+@app.post("/todos")
 def hello():
+    payload = app.current_event.json_body
+
+    body = payload.get("body", "Hello World")
+    status_code = payload.get("status_code", 200)
+    headers = payload.get("headers", {})
+    cookies = payload.get("cookies", [])
+    content_type = headers.get("Content-Type", content_types.TEXT_PLAIN)
+
     return Response(
-        status_code=200,
-        content_type=content_types.TEXT_PLAIN,
-        body="Hello world",
-        cookies=["CookieMonster", "MonsterCookie"],
-        headers={"Foo": ["bar", "zbr"]},
+        status_code=status_code,
+        content_type=content_type,
+        body=body,
+        cookies=cookies,
+        headers=headers,
     )
 
 

tests/e2e/event_handler/infrastructure.py

@@ -61,7 +61,7 @@ def _create_api_gateway_http(self, function: Function):
         apigw = apigwv2.HttpApi(self.stack, "APIGatewayHTTP", create_default_stage=True)
         apigw.add_routes(
             path="/todos",
-            methods=[apigwv2.HttpMethod.GET],
+            methods=[apigwv2.HttpMethod.POST],
             integration=apigwv2integrations.HttpLambdaIntegration("TodosIntegration", function),
         )
 
@@ -71,7 +71,7 @@ def _create_api_gateway_rest(self, function: Function):
         apigw = apigwv1.RestApi(self.stack, "APIGatewayRest", deploy_options=apigwv1.StageOptions(stage_name="dev"))
 
         todos = apigw.root.add_resource("todos")
-        todos.add_method("GET", apigwv1.LambdaIntegration(function, proxy=True))
+        todos.add_method("POST", apigwv1.LambdaIntegration(function, proxy=True))
 
         CfnOutput(self.stack, "APIGatewayRestUrl", value=apigw.url)