From cafd263d0c22dfded3bfba36f644ec02ba55cec9 Mon Sep 17 00:00:00 2001 From: Kevin Bates Date: Wed, 27 Aug 2025 10:42:32 -0700 Subject: [PATCH 1/7] Refactor cookie handling to facilitate improved testing --- jupyter_server/gateway/gateway_client.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/jupyter_server/gateway/gateway_client.py b/jupyter_server/gateway/gateway_client.py index 966cf03f35..49e81356ac 100644 --- a/jupyter_server/gateway/gateway_client.py +++ b/jupyter_server/gateway/gateway_client.py @@ -630,11 +630,17 @@ def load_connection_args(self, **kwargs): return kwargs - def update_cookies(self, cookie: SimpleCookie) -> None: - """Update cookies from existing requests for load balancers""" + def update_cookies(self, headers: dict) -> None: + """Update cookies from response headers for load balancers""" if not self.accept_cookies: return + # Create a SimpleCookie from the Set-Cookie header + cookie = SimpleCookie() + set_cookie = headers.get('Set-Cookie') + if set_cookie: + cookie.load(set_cookie) + store_time = datetime.now(tz=timezone.utc) for key, item in cookie.items(): # Convert "expires" arg into "max-age" to facilitate expiration management. @@ -820,11 +826,7 @@ async def gateway_request(endpoint: str, **kwargs: ty.Any) -> HTTPResponse: ) raise e - if gateway_client.accept_cookies: - # Update cookies on GatewayClient from server if configured. - cookie_values = response.headers.get("Set-Cookie") - if cookie_values: - cookie: SimpleCookie = SimpleCookie() - cookie.load(cookie_values) - gateway_client.update_cookies(cookie) + + gateway_client.update_cookies(response.headers) + return response From 27cd252e76e47c4a9bd5c906fc2bd5ca21b49d54 Mon Sep 17 00:00:00 2001 From: Kevin-Bates Date: Wed, 27 Aug 2025 11:20:28 -0700 Subject: [PATCH 2/7] Update gateway cookie test to reproduce #1557 --- tests/test_gateway.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 00aa64f111..540afc5e98 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -397,16 +397,20 @@ def test_gateway_request_with_expiring_cookies( GatewayClient.clear_instance() _ = jp_configurable_serverapp(argv=argv) - cookie: SimpleCookie = SimpleCookie() - cookie.load("SERVERID=1234567; Path=/") + headers = {} + # Use comma-separated Set-Cookie headers to demonstrate the issue + headers["Set-Cookie"] = "SERVERID=016723f5|12345678; Max-Age=172800; Path=/; HttpOnly," \ + "username-my-server=2|1:0|10:1756250589|35:username-my-server|144:123abc789|87654321; " \ + "Max-Age=172800; Path=/; HttpOnly" if expire_arg == "Expires": expire_param = format_datetime( datetime.now(tz=timezone.utc) + timedelta(seconds=expire_param) ) if expire_arg: - cookie["SERVERID"][expire_arg] = expire_param + # Replace the first SERVERID cookie's Max-Age with the new expire parameter + headers["Cookie"] = headers["Set-Cookie"].replace("Max-Age=172800", f"{expire_arg}={expire_param}") - GatewayClient.instance().update_cookies(cookie) + GatewayClient.instance().update_cookies(headers) args = {} if existing_cookies: From 188b5caa6b69843e2854a6821075a5b9171162c6 Mon Sep 17 00:00:00 2001 From: Kevin-Bates Date: Fri, 29 Aug 2025 10:19:27 -0700 Subject: [PATCH 3/7] Fix gateway_client cookie handling and test --- jupyter_server/gateway/gateway_client.py | 46 ++++++++++++------ tests/test_gateway.py | 59 +++++++++++++----------- 2 files changed, 64 insertions(+), 41 deletions(-) diff --git a/jupyter_server/gateway/gateway_client.py b/jupyter_server/gateway/gateway_client.py index 49e81356ac..921ff1424e 100644 --- a/jupyter_server/gateway/gateway_client.py +++ b/jupyter_server/gateway/gateway_client.py @@ -12,12 +12,13 @@ from abc import ABC, ABCMeta, abstractmethod from datetime import datetime, timezone from email.utils import parsedate_to_datetime -from http.cookies import SimpleCookie +from http.cookies import Morsel, SimpleCookie from socket import gaierror from jupyter_events import EventLogger from tornado import web from tornado.httpclient import AsyncHTTPClient, HTTPClientError, HTTPResponse +from tornado.httputil import HTTPHeaders from traitlets import ( Bool, Float, @@ -630,26 +631,41 @@ def load_connection_args(self, **kwargs): return kwargs - def update_cookies(self, headers: dict) -> None: - """Update cookies from response headers for load balancers""" + def update_cookies(self, headers: HTTPHeaders) -> None: + """Update cookies from response headers""" + if not self.accept_cookies: return - # Create a SimpleCookie from the Set-Cookie header - cookie = SimpleCookie() - set_cookie = headers.get('Set-Cookie') - if set_cookie: - cookie.load(set_cookie) + # Get individual Set-Cookie headers in list form. This handles multiple cookies + # that are otherwise comma-separated in the header and will break the parsing logic + # if only headers.get() is used. + cookie_headers = headers.get_list('Set-Cookie') + if not cookie_headers: + return store_time = datetime.now(tz=timezone.utc) - for key, item in cookie.items(): + for header in cookie_headers: + cookie = SimpleCookie() + try: + cookie.load(header) + except Exception as e: + self.log.warning("Failed to parse cookie header %s: %s", header, e) + continue + + if not cookie: + self.log.warning("No cookies found in header: %s", header) + continue + name, morsel = next(iter(cookie.items())) + # Convert "expires" arg into "max-age" to facilitate expiration management. # As "max-age" has precedence, ignore "expires" when "max-age" exists. - if item.get("expires") and not item.get("max-age"): - expire_timedelta = parsedate_to_datetime(item["expires"]) - store_time - item["max-age"] = str(expire_timedelta.total_seconds()) + if morsel.get("expires") and not morsel.get("max-age"): + expire_time = parsedate_to_datetime(morsel["expires"]) + expire_timedelta = expire_time - store_time + morsel["max-age"] = str(expire_timedelta.total_seconds()) - self._cookies[key] = (item, store_time) + self._cookies[name] = (morsel, store_time) def _clear_expired_cookies(self) -> None: """Clear expired cookies.""" @@ -826,7 +842,7 @@ async def gateway_request(endpoint: str, **kwargs: ty.Any) -> HTTPResponse: ) raise e - - gateway_client.update_cookies(response.headers) + if gateway_client.accept_cookies: + gateway_client.update_cookies(response.headers) return response diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 540afc5e98..031a6995a9 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -6,8 +6,6 @@ import os import uuid from datetime import datetime, timedelta, timezone -from email.utils import format_datetime -from http.cookies import SimpleCookie from io import BytesIO from queue import Empty from typing import Any, Union @@ -18,7 +16,7 @@ from jupyter_core.utils import ensure_async from tornado.concurrent import Future from tornado.httpclient import HTTPRequest, HTTPResponse -from tornado.httputil import HTTPServerRequest +from tornado.httputil import HTTPServerRequest, HTTPHeaders from tornado.queues import Queue from tornado.web import HTTPError from traitlets import Int, Unicode @@ -374,14 +372,13 @@ def test_gateway_request_timeout_pad_option( GatewayClient.clear_instance() - @pytest.mark.parametrize( - "accept_cookies,expire_arg,expire_param,existing_cookies,cookie_exists", + "accept_cookies,expire_arg,expire_param,existing_cookies", [ - (False, None, None, "EXISTING=1", False), - (True, None, None, "EXISTING=1", True), - (True, "Expires", 180, None, True), - (True, "Max-Age", "-360", "EXISTING=1", False), + (False, None, 0, "EXISTING=1"), + (True, None, 0, "EXISTING=1"), + (True, "expires", 180, None), + (True, "Max-Age", -360, "EXISTING=1"), ], ) def test_gateway_request_with_expiring_cookies( @@ -390,39 +387,49 @@ def test_gateway_request_with_expiring_cookies( expire_arg, expire_param, existing_cookies, - cookie_exists, ): argv = [f"--GatewayClient.accept_cookies={accept_cookies}"] GatewayClient.clear_instance() _ = jp_configurable_serverapp(argv=argv) - headers = {} - # Use comma-separated Set-Cookie headers to demonstrate the issue - headers["Set-Cookie"] = "SERVERID=016723f5|12345678; Max-Age=172800; Path=/; HttpOnly," \ - "username-my-server=2|1:0|10:1756250589|35:username-my-server|144:123abc789|87654321; " \ - "Max-Age=172800; Path=/; HttpOnly" - if expire_arg == "Expires": - expire_param = format_datetime( - datetime.now(tz=timezone.utc) + timedelta(seconds=expire_param) - ) - if expire_arg: - # Replace the first SERVERID cookie's Max-Age with the new expire parameter - headers["Cookie"] = headers["Set-Cookie"].replace("Max-Age=172800", f"{expire_arg}={expire_param}") + test_expiration = bool(expire_param < 0) + # Create mock headers with Set-Cookie values + headers = HTTPHeaders() + cookie_value = "SERVERID=1234567; Path=/; HttpOnly" + if expire_arg == "expires": + # Convert expire_param to a string in the format of "Expires: " (RFC 7231) + expire_param = (datetime.now(tz=timezone.utc) + timedelta(seconds=expire_param))\ + .strftime("%a, %d %b %Y %H:%M:%S GMT") + cookie_value = f"SERVERID=1234567; Path=/; expires={expire_param}; HttpOnly" + elif expire_arg == "Max-Age": + cookie_value = f"SERVERID=1234567; Path=/; Max-Age={expire_param}; HttpOnly" + + # Add a second cookie to test comma-separated cookies + headers.add("Set-Cookie", cookie_value) + headers.add("Set-Cookie", "ADDITIONAL_COOKIE=8901234; Path=/; HttpOnly") + + headers_list = headers.get_list("Set-Cookie") + print(headers_list) GatewayClient.instance().update_cookies(headers) args = {} if existing_cookies: args["headers"] = {"Cookie": existing_cookies} + connection_args = GatewayClient.instance().load_connection_args(**args) - if not cookie_exists: - assert "SERVERID" not in (connection_args["headers"].get("Cookie") or "") + if not accept_cookies or test_expiration: + # The first condition ensure the response cookie is not accepted, + # the second condition ensures that the cookie is not accepted if it is expired. + assert "SERVERID" not in connection_args["headers"].get("Cookie") else: - assert "SERVERID" in connection_args["headers"].get("Cookie") + # The cookie is accepted if it is not expired and accept_cookies is True. + assert "SERVERID" in connection_args["headers"].get("Cookie", "") + if existing_cookies: - assert "EXISTING" in connection_args["headers"].get("Cookie") + assert "EXISTING" in connection_args["headers"].get("Cookie", "") GatewayClient.clear_instance() From bf15cefef991abe7d1068bf0e87fa2f4b0db1844 Mon Sep 17 00:00:00 2001 From: Kevin Bates Date: Mon, 1 Sep 2025 08:36:52 -0700 Subject: [PATCH 4/7] Address linting issues --- jupyter_server/gateway/gateway_client.py | 5 ++--- tests/test_gateway.py | 10 ++++++---- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/jupyter_server/gateway/gateway_client.py b/jupyter_server/gateway/gateway_client.py index 921ff1424e..e5ad96daf3 100644 --- a/jupyter_server/gateway/gateway_client.py +++ b/jupyter_server/gateway/gateway_client.py @@ -18,7 +18,6 @@ from jupyter_events import EventLogger from tornado import web from tornado.httpclient import AsyncHTTPClient, HTTPClientError, HTTPResponse -from tornado.httputil import HTTPHeaders from traitlets import ( Bool, Float, @@ -42,7 +41,7 @@ MESSAGE_KEY = "msg" if ty.TYPE_CHECKING: - from http.cookies import Morsel + from tornado.httputil import HTTPHeaders class GatewayTokenRenewerMeta(ABCMeta, type(LoggingConfigurable)): # type: ignore[misc] @@ -640,7 +639,7 @@ def update_cookies(self, headers: HTTPHeaders) -> None: # Get individual Set-Cookie headers in list form. This handles multiple cookies # that are otherwise comma-separated in the header and will break the parsing logic # if only headers.get() is used. - cookie_headers = headers.get_list('Set-Cookie') + cookie_headers = headers.get_list("Set-Cookie") if not cookie_headers: return diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 031a6995a9..bf774bb431 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -16,7 +16,7 @@ from jupyter_core.utils import ensure_async from tornado.concurrent import Future from tornado.httpclient import HTTPRequest, HTTPResponse -from tornado.httputil import HTTPServerRequest, HTTPHeaders +from tornado.httputil import HTTPHeaders, HTTPServerRequest from tornado.queues import Queue from tornado.web import HTTPError from traitlets import Int, Unicode @@ -372,6 +372,7 @@ def test_gateway_request_timeout_pad_option( GatewayClient.clear_instance() + @pytest.mark.parametrize( "accept_cookies,expire_arg,expire_param,existing_cookies", [ @@ -399,8 +400,9 @@ def test_gateway_request_with_expiring_cookies( cookie_value = "SERVERID=1234567; Path=/; HttpOnly" if expire_arg == "expires": # Convert expire_param to a string in the format of "Expires: " (RFC 7231) - expire_param = (datetime.now(tz=timezone.utc) + timedelta(seconds=expire_param))\ - .strftime("%a, %d %b %Y %H:%M:%S GMT") + expire_param = (datetime.now(tz=timezone.utc) + timedelta(seconds=expire_param)).strftime( + "%a, %d %b %Y %H:%M:%S GMT" + ) cookie_value = f"SERVERID=1234567; Path=/; expires={expire_param}; HttpOnly" elif expire_arg == "Max-Age": cookie_value = f"SERVERID=1234567; Path=/; Max-Age={expire_param}; HttpOnly" @@ -421,7 +423,7 @@ def test_gateway_request_with_expiring_cookies( connection_args = GatewayClient.instance().load_connection_args(**args) if not accept_cookies or test_expiration: - # The first condition ensure the response cookie is not accepted, + # The first condition ensure the response cookie is not accepted, # the second condition ensures that the cookie is not accepted if it is expired. assert "SERVERID" not in connection_args["headers"].get("Cookie") else: From 6ced05133a447ad8cf53ae69417de65ff6513c31 Mon Sep 17 00:00:00 2001 From: Kevin Bates Date: Mon, 1 Sep 2025 12:19:57 -0700 Subject: [PATCH 5/7] Fix docs build --- jupyter_server/gateway/gateway_client.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/jupyter_server/gateway/gateway_client.py b/jupyter_server/gateway/gateway_client.py index e5ad96daf3..533836730f 100644 --- a/jupyter_server/gateway/gateway_client.py +++ b/jupyter_server/gateway/gateway_client.py @@ -18,6 +18,7 @@ from jupyter_events import EventLogger from tornado import web from tornado.httpclient import AsyncHTTPClient, HTTPClientError, HTTPResponse +from tornado.httputil import HTTPHeaders from traitlets import ( Bool, Float, @@ -40,9 +41,6 @@ STATUS_CODE_KEY = "status_code" MESSAGE_KEY = "msg" -if ty.TYPE_CHECKING: - from tornado.httputil import HTTPHeaders - class GatewayTokenRenewerMeta(ABCMeta, type(LoggingConfigurable)): # type: ignore[misc] """The metaclass necessary for proper ABC behavior in a Configurable.""" From 5523269f63f9dbda8e3221a5db50ec7755f547b3 Mon Sep 17 00:00:00 2001 From: Kevin Bates Date: Mon, 1 Sep 2025 14:16:03 -0700 Subject: [PATCH 6/7] One more attempt at pre-commit checks --- jupyter_server/gateway/gateway_client.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/jupyter_server/gateway/gateway_client.py b/jupyter_server/gateway/gateway_client.py index 533836730f..78bcb8682b 100644 --- a/jupyter_server/gateway/gateway_client.py +++ b/jupyter_server/gateway/gateway_client.py @@ -18,7 +18,6 @@ from jupyter_events import EventLogger from tornado import web from tornado.httpclient import AsyncHTTPClient, HTTPClientError, HTTPResponse -from tornado.httputil import HTTPHeaders from traitlets import ( Bool, Float, @@ -35,6 +34,9 @@ from jupyter_server import DEFAULT_EVENTS_SCHEMA_PATH, JUPYTER_SERVER_EVENTS_URI +if ty.TYPE_CHECKING: + from tornado.httputil import HTTPHeaders + ERROR_STATUS = "error" SUCCESS_STATUS = "success" STATUS_KEY = "status" From 4c4a229c5b5fa19c91906ab96466be591a188222 Mon Sep 17 00:00:00 2001 From: Kevin Bates Date: Mon, 1 Sep 2025 16:51:44 -0700 Subject: [PATCH 7/7] Remove TYPE_CHECKING, add ignore of TCH on gateway files --- jupyter_server/gateway/gateway_client.py | 4 +--- pyproject.toml | 1 + 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/jupyter_server/gateway/gateway_client.py b/jupyter_server/gateway/gateway_client.py index 78bcb8682b..533836730f 100644 --- a/jupyter_server/gateway/gateway_client.py +++ b/jupyter_server/gateway/gateway_client.py @@ -18,6 +18,7 @@ from jupyter_events import EventLogger from tornado import web from tornado.httpclient import AsyncHTTPClient, HTTPClientError, HTTPResponse +from tornado.httputil import HTTPHeaders from traitlets import ( Bool, Float, @@ -34,9 +35,6 @@ from jupyter_server import DEFAULT_EVENTS_SCHEMA_PATH, JUPYTER_SERVER_EVENTS_URI -if ty.TYPE_CHECKING: - from tornado.httputil import HTTPHeaders - ERROR_STATUS = "error" SUCCESS_STATUS = "success" STATUS_KEY = "status" diff --git a/pyproject.toml b/pyproject.toml index eb2a0ed58c..81fbb5bb7f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -156,6 +156,7 @@ unfixable = [ [tool.ruff.lint.extend-per-file-ignores] "jupyter_server/*" = ["S101", "RET", "S110", "UP031", "FBT", "FA100", "SLF001", "A002", "SIM105", "A001", "UP007", "PLR2004", "T201", "N818", "F403"] +"jupyter_server/gateway/*" = ["TCH" ] "tests/*" = ["UP031", "PT", 'EM', "TRY", "RET", "SLF", "C408", "F841", "FBT", "A002", "FLY", "N", "PERF", "ASYNC", "T201", "FA100", "E711", "INP", "TCH", "SIM105", "A001", "PLW0603"] "examples/*_config.py" = ["F821"]