feat: add TrafficExtension for multi-tenancy

src/charm.py

@@ -13,7 +13,10 @@
 import botocore.exceptions
 import yaml
 from canonical_service_mesh.enums import Action
-from canonical_service_mesh.k8s.resource_manager import PolicyResourceManager
+from canonical_service_mesh.k8s.resource_manager import (
+    KubernetesResourceManager,
+    PolicyResourceManager,
+)
 from canonical_service_mesh.k8s.types.istio import AuthorizationPolicy
 from canonical_service_mesh.models.istio import (
     AuthorizationPolicySpec,
@@ -61,6 +64,7 @@
 )
 from jinja2 import Template
 from lightkube import Client
+from lightkube.generic_resource import GenericNamespacedResource, create_namespaced_resource
 from lightkube.models.core_v1 import ServicePort
 from lightkube.models.meta_v1 import ObjectMeta
 from object_storage import S3Requirer
@@ -100,6 +104,13 @@
     "The service-mesh relation is present but config 'istio_waypoint_principal' is not set. "
     "Set it to the SPIFFE principal of the platform namespace's waypoint proxy."
 )
+TRAFFIC_EXTENSION_TEMPLATE_PATH = "src/service_mesh/traffic-extension.yaml.j2"
+TrafficExtension = create_namespaced_resource(
+    "extensions.istio.io",
+    "v1alpha1",
+    "TrafficExtension",
+    "trafficextensions",
+)
 # path inside the workload container where the artifact store's TLS CA bundle is written to be then
 # referenced by the AWS_CA_BUNDLE environment variable, so that the tracking server can trust the
 # store's TLS certificate - NOTE: under Pebble's home directory, writable by the non-root user:
@@ -234,6 +245,13 @@ def __init__(self, *args):
             self._remove_authorization_policies,
         )
 
+        self.framework.observe(self.on.remove, self._remove_identity_traffic_extension)
+
+        self.framework.observe(
+            self.on[SERVICE_MESH_RELATION_NAME].relation_broken,
+            self._remove_identity_traffic_extension,
+        )
+
         # Log forwarding to Loki
         self._logging = LogForwarder(charm=self)
 
@@ -384,6 +402,20 @@ def _policy_resource_manager(self) -> PolicyResourceManager:
             logger=self.logger,
         )
 
+    @property
+    def _traffic_extension_resource_manager(self) -> KubernetesResourceManager:
+        """Create and return the manager of the charm's TrafficExtension resources."""
+        return KubernetesResourceManager(
+            labels={
+                "app.kubernetes.io/instance": f"{self.app.name}-{self.model.name}",
+                # The manager selects and cleans up resources using this ownership scope.
+                "kubernetes-resource-handler-scope": f"{self.app.name}-identity-traffic-extension",
+            },
+            resource_types={TrafficExtension},
+            lightkube_client=Client(field_manager=f"{self.app.name}-{self.model.name}"),
+            logger=self.logger,
+        )
+
     @property
     def secrets_manifests_wrapper(self):
         if not self._secrets_manifests_wrapper:
@@ -1403,13 +1435,47 @@ def _build_tracking_server_authorization_policy(self) -> Optional[AuthorizationP
             spec=spec.model_dump(by_alias=True, exclude_unset=True, exclude_none=True),
         )
 
+    def _reconcile_identity_traffic_extension(self) -> None:
+        """Reconcile the TrafficExtension identifying in-mesh callers by their namespace."""
+        if not self.unit.is_leader():
+            return
+        if not self.model.get_relation(SERVICE_MESH_RELATION_NAME):
+            # without a service mesh there is no waypoint to extend; any previously created
+            # resource is removed by the relation-broken handler.
+            return
+
+        traffic_extension = self._build_identity_traffic_extension()
+        self._traffic_extension_resource_manager.reconcile([traffic_extension])
+
+    def _build_identity_traffic_extension(self) -> GenericNamespacedResource:
+        """Build the TrafficExtension stamping in-mesh callers' namespace as their identity.
+
+        The waypoint proxy fronting the tracking server overwrites the trusted user-ID header with
+        the namespace of the calling workload, taken from its verified mTLS identity, so that
+        in-mesh clients are mapped to an MLflow user they cannot spoof. Targeting the tracking
+        server's Service (rather than the waypoint's Gateway) keeps the extension from running for
+        the other workloads that the same waypoint fronts.
+        """
+        rendered = Template(Path(TRAFFIC_EXTENSION_TEMPLATE_PATH).read_text()).render(
+            name=f"{self.app.name}-source-namespace-as-identity",
+            namespace=self.model.name,
+            service_name=self._service_name,
+            identity_header_name=self.model.config["identity_header_name"],
+        )
+        return TrafficExtension.from_dict(yaml.safe_load(rendered))
+
     def _remove_authorization_policies(self, _):
         if not self.unit.is_leader():
             return
         self._policy_resource_manager.reconcile(
             policies=[], mesh_type=MeshType.istio, raw_policies=[]
         )
 
+    def _remove_identity_traffic_extension(self, _):
+        if not self.unit.is_leader():
+            return
+        self._traffic_extension_resource_manager.reconcile([])
+
     def _on_upgrade_charm(self, event) -> None:
         """Handle the upgrade-charm event by running migrations to possibly newer database schemas.
 
@@ -1537,6 +1603,8 @@ def _on_event(self, event) -> None:
 
             self._reconcile_policy_resource_manager()
 
+            self._reconcile_identity_traffic_extension()
+
             if not self.exporter_container.can_connect():
                 raise ErrorWithStatus(
                     f"Container {self._exporter_container_name} is not ready", WaitingStatus

src/service_mesh/traffic-extension.yaml.j2

@@ -0,0 +1,43 @@
+apiVersion: extensions.istio.io/v1alpha1
+kind: TrafficExtension
+metadata:
+  name: {{ name }}
+  namespace: {{ namespace }}
+spec:
+  # a waypoint only honours targetRefs, and a Service targetRef is waypoint-only: the extension
+  # therefore runs solely for traffic the waypoint handles on behalf of the tracking server.
+  targetRefs:
+    - kind: Service
+      group: ""
+      name: {{ service_name }}
+  # before the authorization filters, so that policies see the stamped identity
+  phase: AUTHN
+  lua:
+    inlineCode: |
+      local IDENTITY_HEADER = "{{ identity_header_name }}"
+
+      local function source_namespace(request_handle)
+        -- mTLS terminates on the outer HBONE listener; Istio propagates its verified peer URI SAN.
+        local principal = request_handle:streamInfo():filterState():get("io.istio.peer_principal")
+        if principal == nil then
+          return nil
+        end
+        -- spiffe://<trust-domain>/ns/<namespace>/sa/<service-account>
+        return string.match(principal, "^spiffe://[^/]+/ns/([^/]+)/sa/")
+      end
+
+      function envoy_on_request(request_handle)
+        -- any client-supplied user-ID header is dropped first, so that it cannot be spoofed:
+        request_handle:headers():remove(IDENTITY_HEADER)
+
+        local namespace = source_namespace(request_handle)
+        if namespace == nil then
+          -- no verified mTLS peer (plaintext or out of the mesh): fail closed
+          request_handle:respond(
+            {[":status"] = "401"}, "Unauthenticated: no verified source namespace"
+          )
+          return
+        end
+
+        request_handle:headers():add(IDENTITY_HEADER, namespace)
+      end

tests/integration/helpers.py

@@ -0,0 +1,88 @@
+# Copyright 2026 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+"""Shared helpers for the integration tests."""
+
+import logging
+import subprocess
+from random import choices
+from string import ascii_lowercase
+
+from tenacity import Retrying, retry_if_exception_type, stop_after_delay, wait_fixed
+
+logger = logging.getLogger(__name__)
+
+
+class CurlPod:
+    """A long-running pod issuing curl requests from within a given namespace.
+
+    Use it as a context manager, so the pod is created on entry and deleted on exit.
+    """
+
+    IMAGE = "curlimages/curl:8.8.0"
+
+    def __init__(self, namespace: str):
+        self.namespace = namespace
+        self.name = f"curl-{''.join(choices(ascii_lowercase, k=6))}"
+
+    def __enter__(self) -> "CurlPod":
+        self.create()
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self.delete()
+
+    def create(self) -> None:
+        """Create the pod and wait until it is ready, deleting it if it never gets ready."""
+        # a fresh namespace's default service account may not exist yet, rejecting the pod:
+        for attempt in Retrying(
+            stop=stop_after_delay(60),
+            wait=wait_fixed(2),
+            retry=retry_if_exception_type(subprocess.CalledProcessError),
+            reraise=True,
+        ):
+            with attempt:
+                self._kubectl(
+                    "run",
+                    self.name,
+                    f"--image={self.IMAGE}",
+                    "--restart=Never",
+                    "--command",
+                    "--",
+                    "sleep",
+                    "86400",
+                )
+        try:
+            self._kubectl("wait", f"pod/{self.name}", "--for=condition=Ready", "--timeout=180s")
+        except subprocess.CalledProcessError:
+            self.delete()
+            raise
+
+    def delete(self) -> None:
+        """Delete the pod, tolerating it being already gone."""
+        self._kubectl(
+            "delete", "pod", self.name, "--ignore-not-found", "--wait=false", check=False
+        )
+
+    def curl(self, *args: str, check: bool = True) -> subprocess.CompletedProcess:
+        """Run curl in the pod; `kubectl exec` propagates curl's own exit code."""
+        result = self._kubectl("exec", self.name, "--", "curl", *args, check=False)
+        logger.info(
+            "curl from pod %s/%s: exit_code=%s stdout=%r stderr=%r",
+            self.namespace,
+            self.name,
+            result.returncode,
+            result.stdout,
+            result.stderr,
+        )
+        if check:
+            result.check_returncode()
+        return result
+
+    def _kubectl(self, *args: str, check: bool = True) -> subprocess.CompletedProcess:
+        return subprocess.run(
+            ["kubectl", "-n", self.namespace, *args],
+            check=check,
+            capture_output=True,
+            text=True,
+        )