Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 34 additions & 4 deletions labgrid/remote/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)

Expand All @@ -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()
Expand Down Expand Up @@ -836,6 +847,21 @@ async def allow(self):

print(f"allowed {self.args.user} for place {place.name}")

async def unshare(self):
"""Remove another user's access to a previously acquired place"""
place = self.get_acquired_place()
if "/" not in self.args.user:
raise UserError(f"user {self.args.user} must be in <host>/<username> format")
request = labgrid_coordinator_pb2.UnsharePlaceRequest(name=place.name, user=self.args.user)

try:
await self.stub.UnsharePlace(request)
await self.sync_with_coordinator()
except grpc.aio.AioRpcError as e:
raise ServerError(e.details())

print(f"unshared place {place.name} with {self.args.user}")

def get_target_resources(self, place):
self._check_allowed(place)
resources = {}
Expand Down Expand Up @@ -1978,6 +2004,10 @@ def get_parser(auto_doc_mode=False) -> "argparse.ArgumentParser | AutoProgramArg
subparser.add_argument("user", help="<host>/<username>")
subparser.set_defaults(func=ClientSession.allow)

subparser = subparsers.add_parser("unshare", help="remove another user's access to a place")
subparser.add_argument("user", help="<host>/<username>")
subparser.set_defaults(func=ClientSession.unshare)

subparser = subparsers.add_parser("env", help="generate a labgrid environment file for a place")
subparser.set_defaults(func=ClientSession.print_env)

Expand Down
17 changes: 17 additions & 0 deletions labgrid/remote/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
import logging
from datetime import datetime
from fnmatch import fnmatchcase
from typing import Optional
import warnings

import attr

Expand Down Expand Up @@ -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:
Expand Down
94 changes: 90 additions & 4 deletions labgrid/remote/coordinator.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#!/usr/bin/env python3
import argparse
import contextvars
import logging
import asyncio
import traceback
Expand All @@ -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,
Expand All @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand All @@ -941,9 +978,55 @@ async def AllowPlace(self, request, context):
self.save_later()
return labgrid_coordinator_pb2.AllowPlaceResponse()

@locked
async def UnsharePlace(self, request, context):
logging.debug("UnsharePlace name=%s user=%s", request.name, request.user)

placename = request.name
user = request.user
peer = context.peer()
try:
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:
place = self.places[placename]
except KeyError:
await context.abort(grpc.StatusCode.INVALID_ARGUMENT, f"Place {placename} does not exist")
if not place.acquired:
await context.abort(grpc.StatusCode.FAILED_PRECONDITION, f"Place {placename} is not acquired")
if not place.acquired == username:
await context.abort(
grpc.StatusCode.FAILED_PRECONDITION, f"Place {placename} is not acquired by {username}"
)
try:
place.allowed.remove(user)
place.touch()
self._publish_place(place)
self.save_later()
except KeyError:
await context.abort(grpc.StatusCode.FAILED_PRECONDITION, f"Place {placename} is not shared with {user}")

return labgrid_coordinator_pb2.UnsharePlaceResponse()

def _get_places(self):
return {k: v.asdict() for k, v in self.places.items()}

@locked
async def GetPlace(self, request, context):
name = request.name
logging.debug("GetPlace name=%s", name)
if not name or not isinstance(name, str):
await context.abort(grpc.StatusCode.INVALID_ARGUMENT, "name was not a string")
try:
place = self.places[name]
return labgrid_coordinator_pb2.GetPlaceResponse(place=place.as_pb2())
except KeyError:
await context.abort(grpc.StatusCode.INVALID_ARGUMENT, f"Place {name} does not exist")
except Exception:
logging.exception("error during get place")
await context.abort(grpc.StatusCode.INTERNAL, "internal error")

@locked
async def GetPlaces(self, unused_request, unused_context):
logging.debug("GetPlaces")
Expand Down Expand Up @@ -1076,7 +1159,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()
Expand Down Expand Up @@ -1108,7 +1194,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
Expand All @@ -1126,6 +1211,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)
Expand Down
17 changes: 10 additions & 7 deletions labgrid/remote/exporter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand All @@ -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,
Expand Down Expand Up @@ -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:
Expand Down
Loading