Defer loading cookies until needed in auth ext + return better error message

edb/server/protocol/auth_ext/http.py

@@ -140,6 +140,8 @@ async def handle_request(
         )
 
         try:
+            request.load_cookies()
+
             match args:
                 # PKCE token exchange route
                 case ("token",):

edb/server/protocol/http.py

@@ -0,0 +1,89 @@
+
+import logging
+import re
+
+from http.cookies import (  # type: ignore
+  SimpleCookie, Morsel, CookieError,
+  _CookiePattern, _unquote
+)
+
+
+logger = logging.getLogger('edb.server')
+
+
+class InvalidMorsel(Morsel):  # type: ignore
+    pass
+
+
+class SafeSimpleCookie(SimpleCookie):
+
+    def parse(self, str: str, patt: re.Pattern[str]=_CookiePattern) -> None:
+        # Copied from SimpleCookie.__parse_string
+
+        i = 0                 # Our starting point
+        n = len(str)          # Length of string
+        parsed_items = []     # Parsed (type, key, value) triples
+        morsel_seen = False   # A key=value pair was previously encountered
+
+        TYPE_ATTRIBUTE = 1
+        TYPE_KEYVALUE = 2
+
+        # We first parse the whole cookie string and reject it if it's
+        # syntactically invalid (this helps avoid some classes of injection
+        # attacks).
+        while 0 <= i < n:
+            # Start looking for a cookie
+            match = patt.match(str, i)
+            if not match:
+                # No more cookies
+                break
+
+            key, value = match.group("key"), match.group("val")
+            i = match.end(0)
+
+            if key[0] == "$":
+                if not morsel_seen:
+                    # We ignore attributes which pertain to the cookie
+                    # mechanism as a whole, such as "$Version".
+                    # See RFC 2965. (Does anyone care?)
+                    continue
+                parsed_items.append((TYPE_ATTRIBUTE, key[1:], value))
+            elif key.lower() in Morsel._reserved:  # type: ignore
+                if not morsel_seen:
+                    # Invalid cookie string
+                    return
+                if value is None:
+                    if key.lower() in Morsel._flags:  # type: ignore
+                        parsed_items.append((TYPE_ATTRIBUTE, key, True))
+                    else:
+                        # Invalid cookie string
+                        return
+                else:
+                    parsed_items.append((TYPE_ATTRIBUTE, key, _unquote(value)))
+            elif value is not None:
+                parsed_items.append(
+                    (TYPE_KEYVALUE, key, self.value_decode(value))
+                )
+                morsel_seen = True
+            else:
+                # Invalid cookie string
+                return
+
+        # The cookie string is valid, apply it.
+        M = None         # current morsel
+        for tp, key, value in parsed_items:
+            if tp == TYPE_ATTRIBUTE:
+                assert M is not None
+                M[key] = value
+            else:
+                assert tp == TYPE_KEYVALUE
+                rval, cval = value  # type: ignore
+                M = self.get(key, Morsel())
+                try:
+                    M.set(key, rval, cval)
+                    dict.__setitem__(self, key, M)
+                except CookieError as ex:
+                    logger.warning(
+                        f"ignoring invalid cookie: {key}={rval}: {ex}"
+                    )
+                    M = InvalidMorsel()

edb/server/protocol/protocol.pxd

@@ -35,6 +35,7 @@ cdef class HttpRequest:
         public bytes authorization
         public object params
         public object forwarded
+        public bytes cookie_header
         public object cookies
 
 

edb/server/protocol/protocol.pyi

@@ -40,6 +40,9 @@ class HttpRequest:
     forwarded: dict[bytes, bytes]
     cookies: http.cookies.SimpleCookie
 
+    def load_cookies(self) -> None:
+        ...
+
 class HttpResponse:
     status: http.HTTPStatus
     close_connection: bool

edb/server/protocol/protocol.pyx

@@ -50,6 +50,7 @@ from edb.server.dbview cimport dbview
 # can't cimport `protocol.binary` for some reason.
 from edb.server.pgproto.debug cimport PG_DEBUG
 
+from .http import SafeSimpleCookie
 from . import auth
 from . cimport auth_helpers
 from . import edgeql_ext
@@ -80,7 +81,11 @@ cdef class HttpRequest:
         self.authorization = b''
         self.content_type = b''
         self.forwarded = {}
-        self.cookies = http.cookies.SimpleCookie()
+        self.cookie_header = b''
+        self.cookies = SafeSimpleCookie()
+
+    def load_cookies(self):
+        self.cookies.parse(self.cookie_header.decode('ascii'))
 
 
 cdef class HttpResponse:
@@ -264,7 +269,7 @@ cdef class HttpProtocol:
             forwarded_key = name[len(b'x-forwarded-'):]
             self.current_request.forwarded[forwarded_key] = value
         elif name == b'cookie':
-            self.current_request.cookies.load(value.decode('ascii'))
+            self.current_request.cookie_header = value
 
     def on_body(self, body: bytes):
         self.current_request.body += body