Allow serving alternative endpoints on exporter

Author: NickCao

packages/jumpstarter/jumpstarter/client/client.py

@@ -6,8 +6,8 @@
 import grpc
 from anyio.from_thread import BlockingPortal
 from google.protobuf import empty_pb2
-from jumpstarter_protocol import jumpstarter_pb2_grpc
 
+from .grpc import SmartExporterStub
 from jumpstarter.client import DriverClient
 from jumpstarter.common.importlib import import_class
 
@@ -26,13 +26,31 @@ async def client_from_channel(
     stack: ExitStack,
     allow: list[str],
     unsafe: bool,
+    use_alternative_endpoints: bool = True,
 ) -> DriverClient:
     topo = defaultdict(list)
     last_seen = {}
     reports = {}
     clients = OrderedDict()
 
-    response = await jumpstarter_pb2_grpc.ExporterServiceStub(channel).GetReport(empty_pb2.Empty())
+    response = await SmartExporterStub([channel]).GetReport(empty_pb2.Empty())
+
+    channels = [channel]
+    if use_alternative_endpoints:
+        for endpoint in response.alternative_endpoints:
+            if endpoint.certificate:
+                channels.append(
+                    grpc.aio.secure_channel(
+                        endpoint.endpoint,
+                        grpc.ssl_channel_credentials(
+                            root_certificates=endpoint.certificate.encode(),
+                            private_key=endpoint.client_private_key.encode(),
+                            certificate_chain=endpoint.client_certificate.encode(),
+                        ),
+                    )
+                )
+
+    stub = SmartExporterStub(list(reversed(channels)))
 
     for index, report in enumerate(response.reports):
         topo[index] = []
@@ -52,7 +70,7 @@ async def client_from_channel(
         client = client_class(
             uuid=UUID(report.uuid),
             labels=report.labels,
-            channel=channel,
+            stub=stub,
             portal=portal,
             stack=stack.enter_context(ExitStack()),
             children={reports[k].labels["jumpstarter.dev/name"]: clients[k] for k in topo[index]},

packages/jumpstarter/jumpstarter/client/core.py

@@ -5,11 +5,12 @@
 import logging
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
+from typing import Any
 
 from anyio import create_task_group
 from google.protobuf import empty_pb2
 from grpc import StatusCode
-from grpc.aio import AioRpcError, Channel
+from grpc.aio import AioRpcError
 from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc, router_pb2_grpc
 
 from jumpstarter.common import Metadata
@@ -60,16 +61,14 @@ class AsyncDriverClient(
     Backing implementation of blocking driver client.
     """
 
-    channel: Channel
+    stub: Any
 
     log_level: str = "INFO"
     logger: logging.Logger = field(init=False)
 
     def __post_init__(self):
         if hasattr(super(), "__post_init__"):
             super().__post_init__()
-        jumpstarter_pb2_grpc.ExporterServiceStub.__init__(self, self.channel)
-        router_pb2_grpc.RouterServiceStub.__init__(self, self.channel)
         self.logger = logging.getLogger(self.__class__.__name__)
         self.logger.setLevel(self.log_level)
 
@@ -89,7 +88,7 @@ async def call_async(self, method, *args):
         )
 
         try:
-            response = await self.DriverCall(request)
+            response = await self.stub.DriverCall(request)
         except AioRpcError as e:
             match e.code():
                 case StatusCode.UNIMPLEMENTED:
@@ -113,7 +112,7 @@ async def streamingcall_async(self, method, *args):
         )
 
         try:
-            async for response in self.StreamingDriverCall(request):
+            async for response in self.stub.StreamingDriverCall(request):
                 yield decode_value(response.result)
         except AioRpcError as e:
             match e.code():
@@ -128,7 +127,7 @@ async def streamingcall_async(self, method, *args):
 
     @asynccontextmanager
     async def stream_async(self, method):
-        context = self.Stream(
+        context = self.stub.Stream(
             metadata=StreamRequestMetadata.model_construct(request=DriverStreamRequest(uuid=self.uuid, method=method))
             .model_dump(mode="json", round_trip=True)
             .items(),
@@ -142,7 +141,7 @@ async def resource_async(
         self,
         stream,
     ):
-        context = self.Stream(
+        context = self.stub.Stream(
             metadata=StreamRequestMetadata.model_construct(request=ResourceStreamRequest(uuid=self.uuid))
             .model_dump(mode="json", round_trip=True)
             .items(),
@@ -160,7 +159,7 @@ def __log(self, level: int, msg: str):
     @asynccontextmanager
     async def log_stream_async(self):
         async def log_stream():
-            async for response in self.LogStream(empty_pb2.Empty()):
+            async for response in self.stub.LogStream(empty_pb2.Empty()):
                 self.__log(logging.getLevelName(response.severity), response.message)
 
         async with create_task_group() as tg:

packages/jumpstarter/jumpstarter/client/grpc.py

@@ -1,12 +1,16 @@
 from __future__ import annotations
 
-from dataclasses import dataclass, field
+from collections import OrderedDict
+from dataclasses import InitVar, dataclass, field
 from datetime import datetime, timedelta
+from types import SimpleNamespace
+from typing import Any
 
 import yaml
 from google.protobuf import duration_pb2, field_mask_pb2, json_format
+from grpc import ChannelConnectivity
 from grpc.aio import Channel
-from jumpstarter_protocol import client_pb2, client_pb2_grpc, kubernetes_pb2
+from jumpstarter_protocol import client_pb2, client_pb2_grpc, jumpstarter_pb2_grpc, kubernetes_pb2, router_pb2_grpc
 from pydantic import BaseModel, ConfigDict, Field, field_serializer
 
 from jumpstarter.common.grpc import translate_grpc_exceptions
@@ -250,3 +254,25 @@ async def DeleteLease(self, *, name: str):
                     name="namespaces/{}/leases/{}".format(self.namespace, name),
                 )
             )
+
+
+@dataclass(frozen=True, slots=True)
+class SmartExporterStub:
+    channels: InitVar[list[Channel]]
+
+    __stubs: dict[Channel, Any] = field(init=False, default_factory=OrderedDict)
+
+    def __post_init__(self, channels):
+        for channel in channels:
+            stub = SimpleNamespace()
+            jumpstarter_pb2_grpc.ExporterServiceStub.__init__(stub, channel)
+            router_pb2_grpc.RouterServiceStub.__init__(stub, channel)
+            self.__stubs[channel] = stub
+
+    def __getattr__(self, name):
+        for channel, stub in self.__stubs.items():
+            # find the first channel that's ready
+            if channel.get_state(try_to_connect=True) == ChannelConnectivity.READY:
+                return getattr(stub, name)
+        # or fallback to the last channel (via router)
+        return getattr(next(reversed(self.__stubs.values())), name)

packages/jumpstarter/jumpstarter/config/exporter.py

@@ -2,7 +2,7 @@
 
 from contextlib import asynccontextmanager, contextmanager, suppress
 from pathlib import Path
-from typing import Any, ClassVar, Literal, Optional, Self
+from typing import Any, ClassVar, List, Literal, Optional, Self
 
 import grpc
 import yaml
@@ -83,6 +83,8 @@ class ExporterConfigV1Alpha1(BaseModel):
     token: str
     grpcOptions: dict[str, str | int] | None = Field(default_factory=dict)
 
+    alternative_endpoints: List[str] = Field(default_factory=list)
+
     export: dict[str, ExporterConfigV1Alpha1DriverInstance] = Field(default_factory=dict)
 
     path: Path | None = Field(default=None)
@@ -171,6 +173,7 @@ def channel_factory():
             device_factory=ExporterConfigV1Alpha1DriverInstance(children=self.export).instantiate,
             tls=self.tls,
             grpc_options=self.grpcOptions,
+            alternative_endpoints=self.alternative_endpoints,
         ) as exporter:
             await exporter.serve()
 

packages/jumpstarter/jumpstarter/exporter/exporter.py

@@ -25,6 +25,7 @@ class Exporter(AbstractAsyncContextManager, Metadata):
     channel_factory: Callable[[], grpc.aio.Channel]
     device_factory: Callable[[], Driver]
     lease_name: str = field(init=False, default="")
+    alternative_endpoints: list[str] = field(default_factory=list)
     tls: TLSConfigV1Alpha1 = field(default_factory=TLSConfigV1Alpha1)
     grpc_options: dict[str, str] = field(default_factory=dict)
 
@@ -50,7 +51,7 @@ async def session(self):
             labels=self.labels,
             root_device=self.device_factory(),
         ) as session:
-            async with session.serve_unix_async() as path:
+            async with session.serve_unix_async(alternative_endpoints=self.alternative_endpoints) as path:
                 async with grpc.aio.secure_channel(
                     f"unix://{path}", grpc.local_channel_credentials(grpc.LocalConnectionType.UDS)
                 ) as channel:

packages/jumpstarter/jumpstarter/exporter/session.py

@@ -15,6 +15,7 @@
 )
 
 from .logging import LogHandler
+from .tls import with_alternative_endpoints
 from jumpstarter.common import Metadata, TemporarySocket
 from jumpstarter.common.streams import StreamRequestMetadata
 from jumpstarter.driver import Driver
@@ -53,12 +54,16 @@ def __init__(self, *args, root_device, **kwargs):
 
         self._logging_queue = deque(maxlen=32)
         self._logging_handler = LogHandler(self._logging_queue)
+        self._alternative_endpoints = []
 
     @asynccontextmanager
-    async def serve_port_async(self, port):
+    async def serve_ports_async(self, port, alternative_endpoints: list[str] | None = None):
         server = grpc.aio.server()
         server.add_insecure_port(port)
 
+        if alternative_endpoints is not None:
+            self._alternative_endpoints = with_alternative_endpoints(server, alternative_endpoints)
+
         jumpstarter_pb2_grpc.add_ExporterServiceServicer_to_server(self, server)
         router_pb2_grpc.add_RouterServiceServicer_to_server(self, server)
 
@@ -69,15 +74,15 @@ async def serve_port_async(self, port):
             await server.stop(grace=None)
 
     @asynccontextmanager
-    async def serve_unix_async(self):
+    async def serve_unix_async(self, alternative_endpoints: list[str] | None = None):
         with TemporarySocket() as path:
-            async with self.serve_port_async(f"unix://{path}"):
+            async with self.serve_ports_async(f"unix://{path}", alternative_endpoints):
                 yield path
 
     @contextmanager
-    def serve_unix(self):
+    def serve_unix(self, alternative_endpoints: list[str] | None = None):
         with start_blocking_portal() as portal:
-            with portal.wrap_async_context_manager(self.serve_unix_async()) as path:
+            with portal.wrap_async_context_manager(self.serve_unix_async(alternative_endpoints)) as path:
                 yield path
 
     def __getitem__(self, key: UUID):
@@ -92,6 +97,7 @@ async def GetReport(self, request, context):
                 instance.report(parent=parent, name=name)
                 for (_, parent, name, instance) in self.root_device.enumerate()
             ],
+            alternative_endpoints=self._alternative_endpoints,
         )
 
     async def DriverCall(self, request, context):

packages/jumpstarter/jumpstarter/exporter/tls.py

@@ -0,0 +1,98 @@
+from datetime import datetime, timedelta
+from ipaddress import IPv4Address, IPv6Address, ip_address
+
+import grpc
+from cryptography import x509
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from jumpstarter_protocol import jumpstarter_pb2
+
+
+def parse_endpoint(endpoint):
+    host, sep, port = endpoint.rpartition(":")
+
+    if sep == "":
+        raise ValueError("port not specified in endpoint {}".format(endpoint))
+
+    host = host.strip("[]")  # strip brackets from ipv6 addresses
+
+    try:
+        port = int(port)
+        if port < 0 or port > 65535:
+            raise ValueError("port number {} out of range".format(port))
+    except ValueError as e:
+        raise ValueError("invalid port {} in endpoint {}".format(port, endpoint)) from e
+
+    try:
+        return ip_address(host), port
+    except ValueError:
+        return host, port
+
+
+def with_alternative_endpoints(server, endpoints: list[str]):
+    sans = []
+    for endpoint in endpoints:
+        host, port = parse_endpoint(endpoint)
+        match host:
+            case str():
+                sans.append(x509.DNSName(host))
+            case IPv4Address() | IPv6Address():
+                sans.append(x509.IPAddress(host))
+
+    key = rsa.generate_private_key(public_exponent=65537, key_size=2048, backend=default_backend())
+    client_key = rsa.generate_private_key(public_exponent=65537, key_size=2048, backend=default_backend())
+
+    crt = (
+        x509.CertificateBuilder()
+        .subject_name(x509.Name([]))
+        .issuer_name(x509.Name([]))
+        .public_key(key.public_key())
+        .serial_number(x509.random_serial_number())
+        .not_valid_before(datetime.now())
+        .not_valid_after(datetime.now() + timedelta(days=365))
+        .add_extension(x509.SubjectAlternativeName(sans), critical=False)
+        .sign(private_key=key, algorithm=hashes.SHA256(), backend=default_backend())
+    )
+    client_crt = (
+        x509.CertificateBuilder()
+        .subject_name(x509.Name([]))
+        .issuer_name(x509.Name([]))
+        .public_key(client_key.public_key())
+        .serial_number(x509.random_serial_number())
+        .not_valid_before(datetime.now())
+        .not_valid_after(datetime.now() + timedelta(days=365))
+        .sign(private_key=client_key, algorithm=hashes.SHA256(), backend=default_backend())
+    )
+
+    pem_crt = crt.public_bytes(serialization.Encoding.PEM)
+    pem_key = key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.TraditionalOpenSSL,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+    pem_client_crt = client_crt.public_bytes(serialization.Encoding.PEM)
+    pem_client_key = client_key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.TraditionalOpenSSL,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+    server_credentials = grpc.ssl_server_credentials(
+        [(pem_key, pem_crt)], root_certificates=pem_client_crt, require_client_auth=True
+    )
+
+    endpoints_pb = []
+    for endpoint in endpoints:
+        server.add_secure_port(endpoint, server_credentials)
+        endpoints_pb.append(
+            jumpstarter_pb2.Endpoint(
+                endpoint=endpoint,
+                certificate=pem_crt,
+                client_certificate=pem_client_crt,
+                client_private_key=pem_client_key,
+            ),
+        )
+
+    return endpoints_pb

packages/jumpstarter/pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
     "anyio>=4.4.0,!=4.6.2",
     "aiohttp>=3.10.5",
     "tqdm>=4.66.5",
+    "cryptography>=43.0.3",
     "pydantic>=2.8.2"
 ]
 
@@ -25,7 +26,6 @@ dev = [
     "pytest-cov>=6.0.0",
     "pytest-anyio>=0.0.0",
     "pytest-asyncio>=0.0.0",
-    "cryptography>=43.0.3",
     "jumpstarter-driver-power",
     "jumpstarter-driver-network",
     "jumpstarter-driver-composite"