remote: Detach streams from unary RPCs

labgrid/remote/client.py

@@ -32,6 +32,11 @@
 # TODO: drop if Python >= 3.11 guaranteed
 from exceptiongroup import ExceptionGroup  # pylint: disable=redefined-builtin
 
+from labgrid.remote.grpc.interceptor.client import (
+    IdentityClientStreamStreamInterceptor,
+    IdentityClientUnaryUnaryInterceptor,
+)
+
 from .common import (
     ResourceEntry,
     ResourceMatch,
@@ -120,9 +125,19 @@ def __attrs_post_init__(self):
             ("grpc.http2.max_pings_without_data", 0),  # no limit
         ]
 
+        identity = {
+            "username": self.getuser(),
+            "hostname": self.gethostname(),
+            "user_agent": f"labgrid-client {labgrid_version()}",
+        }
+        interceptors = [
+            IdentityClientUnaryUnaryInterceptor(**identity),
+            IdentityClientStreamStreamInterceptor(**identity),
+        ]
         self.channel = grpc.aio.insecure_channel(
             target=self.address,
             options=channel_options,
+            interceptors=interceptors,
         )
         self.stub = labgrid_coordinator_pb2_grpc.CoordinatorStub(self.channel)
 
@@ -139,10 +154,6 @@ async def start(self):
 
         self.pump_task = self.loop.create_task(self.message_pump())
         msg = labgrid_coordinator_pb2.ClientInMessage()
-        msg.startup.version = labgrid_version()
-        msg.startup.name = f"{self.gethostname()}/{self.getuser()}"
-        self.out_queue.put_nowait(msg)
-        msg = labgrid_coordinator_pb2.ClientInMessage()
         msg.subscribe.all_places = True
         self.out_queue.put_nowait(msg)
         msg = labgrid_coordinator_pb2.ClientInMessage()

labgrid/remote/common.py

@@ -7,6 +7,8 @@
 import logging
 from datetime import datetime
 from fnmatch import fnmatchcase
+from typing import Optional
+import warnings
 
 import attr
 
@@ -481,6 +483,21 @@ def from_pb2(cls, pb2: labgrid_coordinator_pb2.Reservation):
         )
 
 
+def get_metadata_single_value_by_key(metadata, key: str) -> Optional[str]:
+    """Look up a single value by key in a metadata sequence of (key, value) pairs."""
+    values = [v for k, v in metadata or () if k == key]
+
+    if not values:
+        return None
+
+    if len(values) > 1:
+        warnings.warn(
+            "Multiple metadata KV pairs with the same key. The value of the first matching KV pair will be returned."
+        )
+
+    return values[0]
+
+
 async def queue_as_aiter(q):
     try:
         while True:

labgrid/remote/coordinator.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 import argparse
+import contextvars
 import logging
 import asyncio
 import traceback
@@ -10,11 +11,15 @@
 import copy
 import random
 import signal
+from typing import Optional
 
 import attr
 import grpc
 from grpc_reflection.v1alpha import reflection
 
+from labgrid.remote.grpc.interceptor.server import IdentityServerInterceptor
+from labgrid.remote.identity import ClientIdentity, infer_peer_identity
+
 from .common import (
     ResourceEntry,
     ResourceMatch,
@@ -30,6 +35,10 @@
 from .generated import labgrid_coordinator_pb2_grpc
 from ..util import atomic_replace, labgrid_version, yaml, Timeout
 
+client_identity_context: contextvars.ContextVar[Optional[ClientIdentity]] = contextvars.ContextVar(
+    "client_identity", default=None
+)
+
 
 @contextmanager
 def warn_if_slow(prefix, *, level=logging.WARNING, limit=0.1):
@@ -317,9 +326,17 @@ async def ClientStream(self, request_iterator, context):
         assert peer not in self.clients
         out_msg_queue = asyncio.Queue()
 
+        identity = client_identity_context.get()
+        if identity:
+            logging.debug("client identity provided in gRPC metadata")
+            logging.debug(identity)
+            self.clients[peer] = ClientSession(self, peer, identity.id, out_msg_queue, identity.user_agent)
+
         async def request_task():
             name = None
             version = None
+            if peer in self.clients:
+                session = self.clients[peer]
             try:
                 async for in_msg in request_iterator:
                     in_msg: labgrid_coordinator_pb2.ClientInMessage
@@ -330,6 +347,9 @@ async def request_task():
                         out_msg.sync.id = in_msg.sync.id
                         out_msg_queue.put_nowait(out_msg)
                     elif kind == "startup":
+                        if peer in self.clients:
+                            logging.debug("already setup, probably because identity was provided in metadata")
+                            continue
                         version = in_msg.startup.version
                         name = in_msg.startup.name
                         session = self.clients[peer] = ClientSession(self, peer, name, out_msg_queue, version)
@@ -409,9 +429,23 @@ async def ExporterStream(self, request_iterator, context):
         out_msg.hello.version = labgrid_version()
         yield out_msg
 
+        identity = client_identity_context.get()
+        if identity:
+            logging.debug("exporter identity provided in gRPC metadata")
+            logging.debug(identity)
+            if existing := self.get_exporter_by_name(identity.id):
+                await context.abort(
+                    grpc.StatusCode.ALREADY_EXISTS,
+                    f"startup failed: exporter with name '{identity.id}' is already connected from {existing.peer}",
+                )
+            self.exporters[peer] = ExporterSession(self, peer, identity.id, command_queue, identity.user_agent)
+            startup_done.set()
+
         async def request_task():
             name = None
             version = None
+            if peer in self.exporters:
+                session = self.exporters[peer]
             try:
                 async for in_msg in request_iterator:
                     in_msg: labgrid_coordinator_pb2.ExporterInMessage
@@ -422,6 +456,9 @@ async def request_task():
                         cmd.complete(in_msg.response)
                         logging.debug("Command %s is done", cmd)
                     elif kind == "startup":
+                        if peer in self.exporters:
+                            logging.debug("already setup, probably because identity was provided in metadata")
+                            continue
                         version = in_msg.startup.version
                         name = in_msg.startup.name
                         if existing := self.get_exporter_by_name(name):
@@ -852,7 +889,7 @@ async def AcquirePlace(self, request, context):
         peer = context.peer()
         name = request.placename
         try:
-            username = self.clients[peer].name
+            username = infer_peer_identity(self.clients, context, client_identity_context)
         except KeyError:
             await context.abort(grpc.StatusCode.FAILED_PRECONDITION, f"Peer {peer} does not have a valid session")
         print(request)
@@ -922,7 +959,7 @@ async def AllowPlace(self, request, context):
         user = request.user
         peer = context.peer()
         try:
-            username = self.clients[peer].name
+            username = infer_peer_identity(self.clients, context, client_identity_context)
         except KeyError:
             await context.abort(grpc.StatusCode.FAILED_PRECONDITION, f"Peer {peer} does not have a valid session")
         try:
@@ -1076,7 +1113,10 @@ async def CreateReservation(self, request: labgrid_coordinator_pb2.CreateReserva
                     await context.abort(grpc.StatusCode.INVALID_ARGUMENT, f"Value {v} is invalid")
                 fltr[k] = v
 
-        owner = self.clients[peer].name
+        try:
+            owner = infer_peer_identity(self.clients, context, client_identity_context)
+        except KeyError:
+            await context.abort(grpc.StatusCode.FAILED_PRECONDITION, f"Peer {peer} does not have a valid session")
         res = Reservation(owner=owner, prio=request.prio, filters=fltrs)
         self.reservations[res.token] = res
         self.schedule_reservations()
@@ -1108,7 +1148,6 @@ async def GetReservations(self, request: labgrid_coordinator_pb2.GetReservations
         reservations = [x.as_pb2() for x in self.reservations.values()]
         return labgrid_coordinator_pb2.GetReservationsResponse(reservations=reservations)
 
-
 async def serve(listen, cleanup) -> None:
     asyncio.current_task().set_name("coordinator-serve")
     # It seems since https://github.com/grpc/grpc/pull/34647, the
@@ -1126,6 +1165,7 @@ async def serve(listen, cleanup) -> None:
     ]
     server = grpc.aio.server(
         options=channel_options,
+        interceptors=[IdentityServerInterceptor(client_identity_context)],
     )
     coordinator = Coordinator()
     labgrid_coordinator_pb2_grpc.add_CoordinatorServicer_to_server(coordinator, server)

labgrid/remote/exporter.py

@@ -20,6 +20,11 @@
 import attr
 import grpc
 
+from labgrid.remote.grpc.interceptor.client import (
+    IdentityClientStreamStreamInterceptor,
+    IdentityClientUnaryUnaryInterceptor,
+)
+
 from .config import ResourceConfig
 from .common import ResourceEntry, queue_as_aiter
 from .generated import labgrid_coordinator_pb2, labgrid_coordinator_pb2_grpc
@@ -831,9 +836,14 @@ def __init__(self, config) -> None:
         if urlsplit(f"//{config['coordinator']}").port is None:
             config["coordinator"] += ":20408"
 
+        identity = (None, self.name, f"labgrid-exporter {labgrid_version()}")
         self.channel = grpc.aio.insecure_channel(
             target=config["coordinator"],
             options=channel_options,
+            interceptors=[
+                IdentityClientUnaryUnaryInterceptor(*identity),
+                IdentityClientStreamStreamInterceptor(*identity),
+            ],
         )
         self.stub = labgrid_coordinator_pb2_grpc.CoordinatorStub(self.channel)
         self.out_queue = asyncio.Queue()
@@ -845,7 +855,6 @@ def __init__(self, config) -> None:
 
     async def run(self) -> None:
         self.pump_task = self.loop.create_task(self.message_pump())
-        self.send_started()
 
         config_template_env = {
             "env": os.environ,
@@ -892,12 +901,6 @@ async def run(self) -> None:
         except asyncio.CancelledError:
             return
 
-    def send_started(self):
-        msg = labgrid_coordinator_pb2.ExporterInMessage()
-        msg.startup.version = labgrid_version()
-        msg.startup.name = self.name
-        self.out_queue.put_nowait(msg)
-
     async def message_pump(self):
         got_message = False
         try:

labgrid/remote/grpc/interceptor/client.py

@@ -0,0 +1,32 @@
+from typing import Optional
+
+from grpc.aio import ClientCallDetails, ClientInterceptor, StreamStreamClientInterceptor, UnaryUnaryClientInterceptor
+
+from labgrid.remote.identity import HOSTNAME_KEY, USER_AGENT_KEY, USERNAME_KEY
+
+
+class BaseIdentityClientInterceptor(ClientInterceptor):
+    def __init__(self, username: Optional[str], hostname: str, user_agent: Optional[str]):
+        super().__init__()
+        self.username = username
+        self.hostname = hostname
+        self.user_agent = user_agent
+
+    def _inject(self, client_call_details: ClientCallDetails):
+        if self.username:
+            client_call_details.metadata.add(USERNAME_KEY, self.username)
+        client_call_details.metadata.add(HOSTNAME_KEY, self.hostname)
+        if self.user_agent:
+            client_call_details.metadata.add(USER_AGENT_KEY, self.user_agent)
+
+
+class IdentityClientUnaryUnaryInterceptor(UnaryUnaryClientInterceptor, BaseIdentityClientInterceptor):
+    async def intercept_unary_unary(self, continuation, client_call_details, request):
+        self._inject(client_call_details)
+        return await continuation(client_call_details, request)
+
+
+class IdentityClientStreamStreamInterceptor(StreamStreamClientInterceptor, BaseIdentityClientInterceptor):
+    async def intercept_stream_stream(self, continuation, client_call_details, request_iterator):
+        self._inject(client_call_details)
+        return await continuation(client_call_details, request_iterator)

labgrid/remote/grpc/interceptor/server.py

@@ -0,0 +1,33 @@
+import contextvars
+import logging
+from asyncio import iscoroutine
+
+from grpc.aio import ServerInterceptor
+
+from labgrid.remote.identity import ClientIdentity, NoIdentityPresent
+
+
+class IdentityServerInterceptor(ServerInterceptor):
+    def __init__(self, client_identity_contextvar: contextvars.ContextVar):
+        super().__init__()
+        self.client_identity_contextvar = client_identity_contextvar
+
+    async def intercept_service(self, continuation, handler_call_details):
+        # continuation may return a handler
+        # OR an awaitable depending on grpcio build
+        maybe_handler = continuation(handler_call_details)
+        handler = await maybe_handler if iscoroutine(maybe_handler) else maybe_handler
+        if handler is None:
+            return None
+
+        metadata = handler_call_details.invocation_metadata
+        logging.debug(metadata)
+
+        try:
+            client_identity = ClientIdentity.from_metadata(metadata)
+            logging.debug(client_identity)
+            self.client_identity_contextvar.set(client_identity)
+        except NoIdentityPresent:
+            pass
+
+        return handler

labgrid/remote/identity.py

@@ -0,0 +1,65 @@
+import contextvars
+import logging
+from typing import Optional
+
+from labgrid.remote.common import get_metadata_single_value_by_key
+
+USERNAME_KEY = "x-lg-username"
+HOSTNAME_KEY = "x-lg-hostname"
+USER_AGENT_KEY = "x-lg-user-agent"
+
+
+class NoIdentityPresent(Exception):
+    """Raised when metadata-based identity information is missing from the request."""
+
+
+class ClientIdentity:
+    """Represents the identity of a connected client, derived from gRPC metadata."""
+
+    def __init__(self, identity_id: str, user_agent: Optional[str]):
+        self.id = identity_id
+        self.user_agent = user_agent
+
+    def __str__(self):
+        return f"ClientIdentity(id={self.id}, user_agent={self.user_agent})"
+
+    @classmethod
+    def from_metadata(cls, metadata: tuple):
+        """Construct a ClientIdentity from gRPC request metadata.
+
+        Args:
+            metadata: A sequence of (key, value) pairs from the gRPC context.
+
+        Returns:
+            A ClientIdentity with id set to ``hostname/username`` (or just
+            ``hostname`` if no username is present) and (optional) user_agent.
+
+        Raises:
+            NoIdentityPresent: If the hostname key is missing from metadata.
+        """
+        username = get_metadata_single_value_by_key(metadata, USERNAME_KEY)
+        hostname = get_metadata_single_value_by_key(metadata, HOSTNAME_KEY)
+        user_agent = get_metadata_single_value_by_key(metadata, USER_AGENT_KEY)
+
+        if not hostname:
+            raise NoIdentityPresent()
+
+        if username:
+            return cls(f"{hostname}/{username}", user_agent)
+
+        return cls(hostname, user_agent)
+
+
+def infer_peer_identity(clients, context, identity_contextvar: contextvars.ContextVar[Optional[ClientIdentity]]):
+    logger = logging.getLogger("infer_peer_identity")
+
+    user = identity_contextvar.get()
+    if user:
+        logger.debug("identity sourced from metadata")
+        return user.id
+
+    try:
+        logger.debug("identity sourced from self.clients")
+        return clients[context.peer()].name
+    except KeyError:
+        raise