Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Sep 2, 2022
5 changes: 3 additions & 2 deletions aws_lambda_powertools/event_handler/api_gateway.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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,
):
"""

Expand All @@ -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
Expand Down
117 changes: 117 additions & 0 deletions aws_lambda_powertools/shared/cookies.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,117 @@
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:
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,
expires: Optional[datetime] = None,
max_age: Optional[int] = 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
expires: Optional[datetime]
Defines a date where the permanent cookie expires.
max_age: Optional[int]
Defines the period of time after which the cookie is invalid. Use negative values to force cookie deletion.
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()
16 changes: 9 additions & 7 deletions aws_lambda_powertools/shared/headers_serializer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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.
Expand All @@ -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.
Expand All @@ -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):
Expand Down
2 changes: 1 addition & 1 deletion docs/core/event_handler/api_gateway.md
Original file line number Diff line number Diff line change
Expand Up @@ -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"
```

Expand Down
2 changes: 1 addition & 1 deletion docs/upgrade.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)],
)
```

Expand Down
3 changes: 2 additions & 1 deletion examples/event_handler_rest/src/fine_grained_responses.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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", secure=True)],
)


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
20 changes: 14 additions & 6 deletions tests/e2e/event_handler/handlers/alb_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)


Expand Down
20 changes: 14 additions & 6 deletions tests/e2e/event_handler/handlers/api_gateway_http_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)


Expand Down
20 changes: 14 additions & 6 deletions tests/e2e/event_handler/handlers/api_gateway_rest_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)


Expand Down
20 changes: 14 additions & 6 deletions tests/e2e/event_handler/handlers/lambda_function_url_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)


Expand Down
4 changes: 2 additions & 2 deletions tests/e2e/event_handler/infrastructure.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
)

Expand All @@ -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)

Expand Down
Loading