Allow serving alternative endpoints on exporter

Author: NickCao

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,74 @@
+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())
+
+    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())
+    )
+
+    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(),
+    )
+
+    server_credentials = grpc.ssl_server_credentials([(pem_key, pem_crt)])
+
+    endpoints_pb = []
+    for endpoint in endpoints:
+        server.add_secure_port(endpoint, server_credentials)
+        # FIXME: generate and check token
+        endpoints_pb.append(
+            jumpstarter_pb2.Endpoint(endpoint=endpoint, token="", certificate=pem_crt),
+        )
+
+    return endpoints_pb