Clean API to control mocked behavior on proxy (#5775)

Author: cbeauchesne

requirements.txt

@@ -9,7 +9,7 @@ filelock==3.12.2  # for parametric tests
 jsonschema==4.16.0
 kubernetes==29.0.0  #lib-injection kubernetes
 matplotlib  # for performance tests
-# mitmproxy==9.0.1  # sueful for proxy development, but failing on gitlab
+mitmproxy==9.0.1
 msgpack==1.0.4
 mypy==1.15.0
 opentelemetry-api==1.27.0 # for parametric tests

utils/_context/containers.py

@@ -541,8 +541,6 @@ def save_image_info(self, dir_path: str):
 
 
 class ProxyContainer(TestedContainer):
-    command_host_port = 11111  # Which port exposed to host to sent proxy commands
-
     def __init__(
         self,
         *,
@@ -574,12 +572,12 @@ def __init__(
                 "SYSTEM_TESTS_IPV6": str(enable_ipv6),
                 "SYSTEM_TEST_MOCKED_BACKEND": str(mocked_backend),
             },
-            working_dir="/app",
+            working_dir="/app/utils",
             volumes={
                 "./utils/": {"bind": "/app/utils/", "mode": "ro"},
             },
-            ports={f"{ProxyPorts.proxy_commands}/tcp": ("127.0.0.1", self.command_host_port)},
-            command="python utils/proxy/core.py",
+            ports={f"{ProxyPorts.proxy_commands}/tcp": ("127.0.0.1", ProxyPorts.proxy_commands)},
+            command="python -m proxy.core",
             healthcheck={
                 "test": f"python -c \"import socket; s=socket.socket({socket_family}); s.settimeout(2); s.connect(('{host_target}', {ProxyPorts.weblog})); s.close()\"",  # noqa: E501
                 "retries": 30,

utils/_remote_config.py

@@ -5,32 +5,17 @@
 import base64
 import hashlib
 import json
-import os
 import re
 import time
 import uuid
 from typing import Any
 from collections.abc import Mapping
 
-import requests
-
 from utils._context.core import context
 from utils.dd_constants import RemoteConfigApplyState as ApplyState
 from utils.interfaces import library
 from utils._logger import logger
-from utils._context.containers import ProxyContainer
-
-
-def _post(path: str, payload: list[dict] | dict) -> None:
-    if "SYSTEM_TESTS_PROXY_HOST" in os.environ:
-        domain = os.environ["SYSTEM_TESTS_PROXY_HOST"]
-    elif "DOCKER_HOST" in os.environ:
-        m = re.match(r"(?:ssh:|tcp:|fd:|)//(?:[^@]+@|)([^:]+)", os.environ["DOCKER_HOST"])
-        domain = m.group(1) if m is not None else "localhost"
-    else:
-        domain = "localhost"
-
-    requests.post(f"http://{domain}:{ProxyContainer.command_host_port}{path}", data=json.dumps(payload), timeout=30)
+from utils.proxy.mocked_response import StaticJsonMockedResponse, SequentialRemoteConfigJsonMockedResponse
 
 
 class RemoteConfigStateResults:
@@ -126,7 +111,8 @@ def remote_config_applied(data: dict) -> bool:
         current_states.state = ApplyState.ACKNOWLEDGED
         return True
 
-    _post("/unique_command", raw_payload)
+    StaticJsonMockedResponse(path="/v0.7/config", mocked_json=raw_payload).send()
+
     library.wait_for(remote_config_applied, timeout=30)
     # ensure the library has enough time to apply the config to all subprocesses
     time.sleep(2)
@@ -140,7 +126,7 @@ def send_sequential_commands(commands: list[dict], *, wait_for_all_command: bool
     if len(commands) == 0:
         raise ValueError("No commands to send")
 
-    _post("/sequential_commands", commands)
+    SequentialRemoteConfigJsonMockedResponse(mocked_json_sequence=commands).send()
 
     if not wait_for_all_command:
         return

utils/proxy/_deserializer.py

@@ -27,8 +27,8 @@
     ExportLogsServiceRequest,
     ExportLogsServiceResponse,
 )
-from _decoders.protobuf_schemas import MetricPayload, TracePayload, SketchPayload, BackendResponsePayload
-from traces.trace_v1 import deserialize_v1_trace, _uncompress_agent_v1_trace
+from ._decoders.protobuf_schemas import MetricPayload, TracePayload, SketchPayload, BackendResponsePayload
+from .traces.trace_v1 import deserialize_v1_trace, _uncompress_agent_v1_trace
 
 
 logger = logging.getLogger(__name__)

utils/proxy/core.py

@@ -1,5 +1,5 @@
 # keep this import in first
-import scrubber  # noqa: F401
+from . import scrubber  # noqa: F401
 
 import asyncio
 from collections import defaultdict
@@ -11,11 +11,19 @@
 
 from mitmproxy import master, options, http
 from mitmproxy.addons import errorcheck, default_addons
-from mitmproxy.flow import Error as FlowError, Flow
+from mitmproxy.flow import Error as FlowError
 from mitmproxy.http import HTTPFlow, Request
 
-from _deserializer import deserialize
-from ports import ProxyPorts
+from ._deserializer import deserialize
+from .ports import ProxyPorts
+from .mocked_response import (
+    MOCKED_RESPONSE_PATH,
+    MockedResponse,
+    AddRemoteConfigEndpoint,
+    RemoveMetaStructsSupport,
+    SetSpanEventFlags,
+    StaticJsonMockedResponse,
+)
 
 # prevent permission issues on file created by the proxy when the host is linux
 os.umask(0)
@@ -43,24 +51,29 @@ def __init__(self) -> None:
         logger.addHandler(handler)
         logger.setLevel(logging.DEBUG)
 
-        self.host_log_folder = "logs"
+        self.host_log_folder = "/app/logs"
 
         self.rc_api_enabled = os.environ.get("SYSTEM_TESTS_RC_API_ENABLED") == "True"
         self.mocked_backend = os.environ.get("SYSTEM_TEST_MOCKED_BACKEND") == "True"
-
         self.span_meta_structs_disabled = os.environ.get("SYSTEM_TESTS_AGENT_SPAN_META_STRUCTS_DISABLED") == "True"
+        self.span_events = os.environ.get("SYSTEM_TESTS_AGENT_SPAN_EVENTS") != "False"
 
         self.tracing_agent_target_host = os.environ.get("PROXY_TRACING_AGENT_TARGET_HOST", "agent")
         self.tracing_agent_target_port = int(os.environ.get("PROXY_TRACING_AGENT_TARGET_PORT", "8127"))
 
-        span_events = os.environ.get("SYSTEM_TESTS_AGENT_SPAN_EVENTS")
-        self.span_events = span_events != "False"
+        self.mocked_response: MockedResponse | None = None
+        self.internal_mocked_responses: list[MockedResponse] = [
+            SetSpanEventFlags(span_events=self.span_events),
+        ]
 
-        self.rc_api_command = None
+        if self.rc_api_enabled:
+            # add the remote config endpoint on available agent endpoints
+            self.internal_mocked_responses.append(AddRemoteConfigEndpoint())
+            # mimic the default response from the agent
+            self.internal_mocked_responses.append(StaticJsonMockedResponse(path="/v0.7/config", mocked_json={}))
 
-        # mimic the old API
-        self.rc_api_sequential_commands = None
-        self.rc_api_runtime_ids_request_count: dict = {}
+        if self.span_meta_structs_disabled:
+            self.internal_mocked_responses.append(RemoveMetaStructsSupport())
 
         logger.info(f"rc_api_enabled: {self.rc_api_enabled}")
         logger.info(f"mocked_backend: {self.mocked_backend}")
@@ -75,20 +88,19 @@ def request(self, flow: HTTPFlow):
         # sockname is the local address (host, port) we received this connection on.
         port = flow.client_conn.sockname[1]
 
-        logger.info(f"{flow.request.method} {flow.request.pretty_url}, using proxy port {port}")
+        logger.info(f"{flow.request.method} {flow.request.pretty_url} using proxy port {port}")
 
         if port == ProxyPorts.proxy_commands:
-            if not self.rc_api_enabled:
-                flow.response = self.get_error_response(b"RC API is not enabled")
-            elif flow.request.path == "/unique_command":
-                logger.info("Store RC command to mock")
-                self.rc_api_command = flow.request.content
-                flow.response = http.Response.make(200, b"Ok")
-            elif flow.request.path == "/sequential_commands":
-                logger.info("Reset mocked RC sequential commands")
-                self.rc_api_sequential_commands = json.loads(flow.request.content)
-                self.rc_api_runtime_ids_request_count = defaultdict(int)
-                flow.response = http.Response.make(200, b"Ok")
+            if flow.request.path == MOCKED_RESPONSE_PATH and flow.request.method == "PUT":
+                source = json.loads(flow.request.content)
+                try:
+                    self.mocked_response = MockedResponse.build_from_json(source) if source else None
+                except Exception as e:
+                    logger.exception(f"Failed to build mocked response from {source}")
+                    flow.response = self.get_error_response(f"Invalid mocked response definition: {e}".encode())
+                else:
+                    logger.info(f"Store mocked response :{self.mocked_response}")
+                    flow.response = http.Response.make(200, b"Ok")
             else:
                 flow.response = http.Response.make(404, b"Not found")
 
@@ -239,70 +251,17 @@ def response(self, flow: HTTPFlow):
         except:
             logger.exception("Unexpected error")
 
-    def _modify_response(self, flow: Flow):
+    def _modify_response(self, flow: HTTPFlow):
         if self.request_is_from_tracer(flow.request):
-            if self.rc_api_enabled:
-                self._add_rc_capabilities_in_info_request(flow)
-
-                if flow.request.path == "/v0.7/config":
-                    # mimic the default response from the agent
-                    flow.response.status_code = 200
-                    flow.response.content = b"{}"
-                    flow.response.headers["Content-Type"] = "application/json"
-
-                    if self.rc_api_command is not None:
-                        request_content = json.loads(flow.request.content)
-                        logger.info("    => modifying rc response")
-                        flow.response.content = self.rc_api_command
-
-                    elif self.rc_api_sequential_commands is not None:
-                        request_content = json.loads(flow.request.content)
-                        runtime_id = request_content["client"]["client_tracer"]["runtime_id"]
-                        nth_api_command = self.rc_api_runtime_ids_request_count[runtime_id]
-                        response = self.rc_api_sequential_commands[nth_api_command]
-
-                        logger.info(f"    => Modifying RC response for runtime ID {runtime_id}")
-                        logger.info(f"    => Overwriting /v0.7/config response #{nth_api_command}")
-
-                        flow.response.content = json.dumps(response).encode()
-                        flow.response.headers["st-proxy-overwrite-rc-response"] = f"{nth_api_command}"
-
-                        if nth_api_command + 1 < len(self.rc_api_sequential_commands):
-                            self.rc_api_runtime_ids_request_count[runtime_id] = nth_api_command + 1
-
-            if self.span_meta_structs_disabled:
-                self._remove_meta_structs_support(flow)
-
-            self._modify_span_events_flag(flow)
-
-    def _remove_meta_structs_support(self, flow: Flow):
-        if flow.request.path == "/info" and str(flow.response.status_code) == "200":
-            c = json.loads(flow.response.content)
-            if "span_meta_structs" in c:
-                logger.info("    => Overwriting /info response to remove span_meta_structs field")
-                c.pop("span_meta_structs")
-                flow.response.content = json.dumps(c).encode()
-
-    def _add_rc_capabilities_in_info_request(self, flow: Flow):
-        if flow.request.path == "/info" and str(flow.response.status_code) == "200":
-            c = json.loads(flow.response.content)
-
-            if "/v0.7/config" not in c["endpoints"]:
-                logger.info("    => Overwriting /info response to include /v0.7/config")
-                c["endpoints"].append("/v0.7/config")
-                flow.response.content = json.dumps(c).encode()
-
-    def _modify_span_events_flag(self, flow: Flow):
-        """Modify the agent flag that signals support for native span event serialization.
-        There are three possible cases:
-        - Not configured: agent's response is not modified, the real agent behavior is preserved
-        - `true`: agent advertises support for native span events serialization
-        - `false`: agent advertises that it does not support native span events serialization
-        """
-        if flow.request.path == "/info" and str(flow.response.status_code) == "200":
-            c = json.loads(flow.response.content)
-            c["span_events"] = self.span_events
-            flow.response.content = json.dumps(c).encode()
+            for internal_mocked_response in self.internal_mocked_responses:
+                if internal_mocked_response.path == flow.request.path:
+                    logger.info(f"    => applying {internal_mocked_response}")
+                    internal_mocked_response.execute(flow)
+
+            if self.mocked_response is not None:
+                if self.mocked_response.path == flow.request.path:
+                    logger.info(f"    => applying {self.mocked_response}")
+                    self.mocked_response.execute(flow)
 
 
 def start_proxy() -> None:
@@ -324,7 +283,7 @@ def start_proxy() -> None:
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
     listen_host = "::" if os.environ.get("SYSTEM_TESTS_IPV6") == "True" else "0.0.0.0"  # noqa: S104
-    opts = options.Options(mode=modes, listen_host=listen_host, confdir="utils/proxy/.mitmproxy")
+    opts = options.Options(mode=modes, listen_host=listen_host, confdir="/app/utils/proxy/.mitmproxy")
     proxy = master.Master(opts, event_loop=loop)
     proxy.addons.add(*default_addons())
     proxy.addons.add(errorcheck.ErrorCheck())