Fix gateway_client cookie handling and test

Author: kevin-bates

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

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: <date>" (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()