[Chief] Follow up fix for session escaping (mlrun#2950)

mlrun/api/utils/clients/chief.py

@@ -230,9 +230,14 @@ def _resolve_request_kwargs_from_request(
             # we will url-encode them (aka quote), so the value would be safe against such escaping.
             # e.g.: instead of having "x":"y" being escaped to "\"x\":\"y\"", it will be escaped to "%22x%22:%22y%22"
             elif cookie_name == "session" and mlrun.mlconf.is_running_on_iguazio():
-                request_kwargs["cookies"][cookie_name] = urllib.parse.quote(
+
+                # unquote first, to avoid double quoting ourselves, in case the cookie is already quoted
+                unquoted_session = urllib.parse.unquote(
                     request_kwargs["cookies"][cookie_name]
                 )
+                request_kwargs["cookies"][cookie_name] = urllib.parse.quote(
+                    unquoted_session
+                )
 
         request_kwargs.update(**kwargs)
         return request_kwargs

mlrun/api/utils/clients/iguazio.py

@@ -453,7 +453,13 @@ def _send_request_to_api(
         self._prepare_request_kwargs(session, path, kwargs=kwargs)
         response = self._session.request(method, url, verify=False, **kwargs)
         if not response.ok:
-            self._handle_error_response(method, path, response, error_message, kwargs)
+            try:
+                response_body = response.json()
+            except Exception:
+                response_body = {}
+            self._handle_error_response(
+                method, path, response, response_body, error_message, kwargs
+            )
         return response
 
     def _generate_auth_info_from_session_verification_response(
@@ -686,19 +692,22 @@ def _prepare_request_kwargs(self, session, path, *, kwargs):
                 if isinstance(dict_[key], enum.Enum):
                     dict_[key] = dict_[key].value
 
-    def _handle_error_response(self, method, path, response, error_message, kwargs):
+    def _handle_error_response(
+        self, method, path, response, response_body, error_message, kwargs
+    ):
         log_kwargs = copy.deepcopy(kwargs)
         log_kwargs.update({"method": method, "path": path})
-        if response.content:
-            try:
-                data = response.json()
-                ctx = data.get("meta", {}).get("ctx")
-                errors = data.get("errors", [])
-            except Exception:
-                pass
-            else:
+        try:
+            ctx = response_body.get("meta", {}).get("ctx")
+            errors = response_body.get("errors", [])
+        except Exception:
+            pass
+        else:
+            if errors:
                 error_message = f"{error_message}: {str(errors)}"
+            if errors or ctx:
                 log_kwargs.update({"ctx": ctx, "errors": errors})
+
         logger.warning("Request to iguazio failed", **log_kwargs)
         mlrun.errors.raise_for_status(response, error_message)
 
@@ -787,8 +796,12 @@ async def _send_request_to_api_async(
                 method, url, verify_ssl=False, **kwargs
             )
             if not response.ok:
+                try:
+                    response_body = await response.json()
+                except Exception:
+                    response_body = {}
                 self._handle_error_response(
-                    method, path, response, error_message, kwargs
+                    method, path, response, response_body, error_message, kwargs
                 )
             yield response
         finally:

tests/api/utils/clients/test_chief.py

@@ -18,9 +18,12 @@
 import unittest.mock
 
 import aiohttp
+import aiohttp.web
 import fastapi.encoders
 import pytest
+import starlette.datastructures
 from aiohttp import ClientConnectorError
+from aiohttp.test_utils import TestClient, TestServer
 
 import mlrun.api.schemas
 import mlrun.api.utils.clients.chief
@@ -264,3 +267,46 @@ def _generate_background_task(
         status=mlrun.api.schemas.BackgroundTaskStatus(state=state.value),
         spec=mlrun.api.schemas.BackgroundTaskSpec(),
     )
+
+
+@pytest.mark.parametrize(
+    "session_cookie, expected_cookie_header",
+    [
+        # no escape is needed, sanity
+        ("j", "j"),
+        # escape is needed
+        ('j:{"', "j%3A%7B%22"),
+        # do not double escape
+        ("j%3A%7B%22", "j%3A%7B%22"),
+    ],
+)
+@pytest.mark.asyncio
+async def test_do_not_escape_cookie(
+    chief_client, session_cookie, expected_cookie_header
+):
+    async def handler(request):
+        assert (
+            request.headers["cookie"] == f"session={expected_cookie_header}"
+        ), "Cookie header escaping is malfunctioning"
+        assert (
+            request.cookies["session"] == expected_cookie_header
+        ), "Cookie session escaping is malfunctioning"
+        return aiohttp.web.Response(status=200)
+
+    mock_request = fastapi.Request({"type": "http"})
+    mock_request._headers = starlette.datastructures.Headers()
+    mock_request._cookies = {"session": session_cookie}
+    mock_request._query_params = starlette.datastructures.QueryParams()
+
+    app = aiohttp.web.Application()
+    app.router.add_post("/api/v1/operations/migrations", handler)
+    async with TestClient(TestServer(app)) as client:
+        chief_client._api_url = ""
+        await chief_client._ensure_session()
+        chief_client._session._client = client
+
+        # set that to make sure session escaping is on
+        # coupled with chief_client._resolve_request_kwargs_from_request logic.
+        mlrun.mlconf.igz_version = "0.5.0"
+        response = await chief_client.trigger_migrations(mock_request)
+        assert response.status_code == 200