From a7c5a41396752c39a5a9b688e2dccfaca152a62f Mon Sep 17 00:00:00 2001
From: Alex Oleshkevich <alex.oleshkevich@gmail.com>
Date: Wed, 12 Jan 2022 12:57:47 +0300
Subject: [PATCH] Allow Session scoped cookies. (#1387)

* Allow Session scoped cookies.

* Update docs/middleware.md

Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>

* Improve typing.

Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>
---
 docs/middleware.md               |  2 +-
 starlette/middleware/sessions.py | 14 +++++++-------
 tests/middleware/test_session.py | 17 +++++++++++++++++
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/docs/middleware.md b/docs/middleware.md
index 6d4d1a611..5fe7ce516 100644
--- a/docs/middleware.md
+++ b/docs/middleware.md
@@ -98,7 +98,7 @@ The following arguments are supported:
 
 * `secret_key` - Should be a random string.
 * `session_cookie` - Defaults to "session".
-* `max_age` - Session expiry time in seconds. Defaults to 2 weeks.
+* `max_age` - Session expiry time in seconds. Defaults to 2 weeks. If set to `None` then the cookie will last as long as the browser session.
 * `same_site` - SameSite flag prevents the browser from sending session cookie along with cross-site requests. Defaults to `'lax'`.
 * `https_only` - Indicate that Secure flag should be set (can be used with HTTPS only). Defaults to `False`.
 
diff --git a/starlette/middleware/sessions.py b/starlette/middleware/sessions.py
index ad7a6ee89..3ff1e3de1 100644
--- a/starlette/middleware/sessions.py
+++ b/starlette/middleware/sessions.py
@@ -16,7 +16,7 @@ def __init__(
         app: ASGIApp,
         secret_key: typing.Union[str, Secret],
         session_cookie: str = "session",
-        max_age: int = 14 * 24 * 60 * 60,  # 14 days, in seconds
+        max_age: typing.Optional[int] = 14 * 24 * 60 * 60,  # 14 days, in seconds
         same_site: str = "lax",
         https_only: bool = False,
     ) -> None:
@@ -55,12 +55,12 @@ async def send_wrapper(message: Message) -> None:
                     data = b64encode(json.dumps(scope["session"]).encode("utf-8"))
                     data = self.signer.sign(data)
                     headers = MutableHeaders(scope=message)
-                    header_value = "%s=%s; path=%s; Max-Age=%d; %s" % (
-                        self.session_cookie,
-                        data.decode("utf-8"),
-                        path,
-                        self.max_age,
-                        self.security_flags,
+                    header_value = "{session_cookie}={data}; path={path}; {max_age}{security_flags}".format(  # noqa E501
+                        session_cookie=self.session_cookie,
+                        data=data.decode("utf-8"),
+                        path=path,
+                        max_age=f"Max-Age={self.max_age}; " if self.max_age else "",
+                        security_flags=self.security_flags,
                     )
                     headers.append("Set-Cookie", header_value)
                 elif not initial_session_was_empty:
diff --git a/tests/middleware/test_session.py b/tests/middleware/test_session.py
index 07296bcbb..867a96735 100644
--- a/tests/middleware/test_session.py
+++ b/tests/middleware/test_session.py
@@ -129,3 +129,20 @@ def test_invalid_session_cookie(test_client_factory):
     # we expect it to not raise an exception if we provide a bogus session cookie
     response = client.get("/view_session", cookies={"session": "invalid"})
     assert response.json() == {"session": {}}
+
+
+def test_session_cookie(test_client_factory):
+    app = create_app()
+    app.add_middleware(SessionMiddleware, secret_key="example", max_age=None)
+    client = test_client_factory(app)
+
+    response = client.post("/update_session", json={"some": "data"})
+    assert response.json() == {"session": {"some": "data"}}
+
+    # check cookie max-age
+    set_cookie = response.headers["set-cookie"]
+    assert "Max-Age" not in set_cookie
+
+    client.cookies.clear_session_cookies()
+    response = client.get("/view_session")
+    assert response.json() == {"session": {}}