fix(workloadapi): harden client and sources for safe streaming and errors

Signed-off-by: Max Lambrecht <maxlambrecht@gmail.com>

Author: maxlambrecht

spiffe/src/spiffe/workloadapi/handle_error.py

@@ -36,11 +36,18 @@ def wrapper(*args, **kw):
             except ArgumentError as ae:
                 raise ae
             except PySpiffeError as pe:
-                raise error_cls(str(pe))
+                # Avoid double-wrapping if it's already the expected error type
+                if isinstance(pe, error_cls):
+                    raise pe
+                raise error_cls(str(pe)) from pe
             except grpc.RpcError as rpc_error:
                 if isinstance(rpc_error, grpc.Call):
-                    raise error_cls(str(rpc_error.details()))
-                raise error_cls(DEFAULT_WL_API_ERROR_MESSAGE)
+                    details = rpc_error.details()
+                    code = rpc_error.code()
+                    raise error_cls(
+                        f'{DEFAULT_WL_API_ERROR_MESSAGE}: {details} ({code})'
+                    ) from rpc_error
+                raise error_cls(DEFAULT_WL_API_ERROR_MESSAGE) from rpc_error
             except Exception as e:
                 raise error_cls(str(e))
 

spiffe/src/spiffe/workloadapi/jwt_source.py

@@ -16,14 +16,14 @@
 
 import logging
 import threading
-from typing import Optional, Set, Callable, List
+from typing import Optional, Set, Callable, List, FrozenSet
 
 from spiffe.spiffe_id.spiffe_id import SpiffeId
 from spiffe.bundle.jwt_bundle.jwt_bundle import JwtBundle
 from spiffe.bundle.jwt_bundle.jwt_bundle_set import JwtBundleSet
 from spiffe.spiffe_id.spiffe_id import TrustDomain
 from spiffe.svid.jwt_svid import JwtSvid
-from spiffe.workloadapi.workload_api_client import WorkloadApiClient
+from spiffe.workloadapi.workload_api_client import WorkloadApiClient, StreamCancelHandler
 from spiffe.workloadapi.errors import JwtSourceError
 from spiffe.errors import ArgumentError
 
@@ -79,9 +79,12 @@ def __init__(
         self._subscribers: List[Callable] = []
         self._subscribers_lock = threading.Lock()
 
+        # Track ownership: if we create the client, we own it
+        self._owns_client = workload_api_client is None
         self._workload_api_client = (
             workload_api_client if workload_api_client else WorkloadApiClient(socket_path)
         )
+        self._client_cancel_handler: Optional[StreamCancelHandler] = None
 
         # Start the watcher in a separate thread
         threading.Thread(target=self._start_watcher, daemon=True).start()
@@ -100,12 +103,16 @@ def __init__(
             raise JwtSourceError(f"Failed to create JwtSource: {self._error}") from self._error
 
     @property
-    def bundles(self) -> Set[JwtBundle]:
+    def bundles(self) -> FrozenSet[JwtBundle]:
         """Returns the set of all JwtBundles."""
         with self._lock:
+            if self._error is not None:
+                raise JwtSourceError(
+                    f'Cannot get Jwt Bundles: source has error: {self._error}'
+                )
             if self._closed:
                 raise JwtSourceError('Cannot get Jwt Bundles: source is closed')
-            return self._jwt_bundle_set.bundles
+            return frozenset(self._jwt_bundle_set.bundles)
 
     def fetch_svid(self, audience: Set[str], subject: Optional[SpiffeId] = None) -> JwtSvid:
         """Fetches an JWT-SVID from the source.
@@ -124,7 +131,9 @@ def fetch_svid(self, audience: Set[str], subject: Optional[SpiffeId] = None) ->
         jwt_svid = self._workload_api_client.fetch_jwt_svid(audience, subject)
         return jwt_svid
 
-    def fetch_svids(self, audiences: Set[str], subject: Optional[SpiffeId] = None) -> JwtSvid:
+    def fetch_svids(
+        self, audiences: Set[str], subject: Optional[SpiffeId] = None
+    ) -> List[JwtSvid]:
         """Fetches all JWT-SVIDs from the source.
 
         Args:
@@ -148,6 +157,8 @@ def get_bundle_for_trust_domain(self, trust_domain: TrustDomain) -> Optional[Jwt
             JwtSourceError: In case this JWT Source is closed.
         """
         with self._lock:
+            if self._error is not None:
+                raise JwtSourceError(f'Cannot get JWT Bundle: source has error: {self._error}')
             if self._closed:
                 raise JwtSourceError('Cannot get JWT Bundle: source is closed')
             return self._jwt_bundle_set.get_bundle_for_trust_domain(trust_domain)
@@ -161,9 +172,11 @@ def close(self) -> None:
         """
         _logger.info("Closing JWT Source")
         with self._lock:
-            # the cancel method throws a grpc exception, that can be discarded
+            if self._closed:
+                return
             try:
-                self._client_cancel_handler.cancel()
+                if self._client_cancel_handler:
+                    self._client_cancel_handler.cancel()
             except Exception as err:
                 _logger.exception(
                     'JWT Source: Exception canceling the Workload API client connection: {}'.format(
@@ -172,6 +185,14 @@ def close(self) -> None:
                 )
             self._closed = True
 
+        if self._owns_client:
+            try:
+                self._workload_api_client.close()
+            except Exception as err:
+                _logger.exception(
+                    'Exception closing owned Workload API client: {}'.format(str(err))
+                )
+
     def is_closed(self) -> bool:
         """Checks if the source has been closed, disallowing further operations."""
         with self._lock:
@@ -195,7 +216,10 @@ def unsubscribe_for_updates(self, callback: Callable[[], None]) -> None:
             callback (Callable[[], None]): The callback function to unregister.
         """
         with self._subscribers_lock:
-            self._subscribers.remove(callback)
+            try:
+                self._subscribers.remove(callback)
+            except ValueError:
+                pass
 
     def _start_watcher(self) -> None:
         self._client_cancel_handler = self._workload_api_client.stream_jwt_bundles(
@@ -213,16 +237,23 @@ def _set_jwt_bundle_set(self, jwt_bundle_set: JwtBundleSet) -> None:
 
     def _notify_subscribers(self) -> None:
         with self._subscribers_lock:
-            for callback in self._subscribers:
-                try:
-                    callback()
-                except Exception as err:
-                    _logger.exception(f"An error occurred while notifying a subscriber: {err}")
+            subscribers = list(self._subscribers)
+        for callback in subscribers:
+            try:
+                callback()
+            except Exception as err:
+                _logger.exception(f"An error occurred while notifying a subscriber: {err}")
 
     def _on_error(self, error: Exception) -> None:
         self._log_error(error)
         with self._lock:
             self._error = error
+            self._closed = True
+            try:
+                if self._client_cancel_handler:
+                    self._client_cancel_handler.cancel()
+            except Exception as err:
+                _logger.exception(f"Exception canceling stream on error: {err}")
         self._initialization_event.set()
 
     @staticmethod

spiffe/src/spiffe/workloadapi/workload_api_client.py

@@ -17,7 +17,6 @@
 import logging
 import os
 import threading
-import time
 from typing import Optional, List, Mapping, Callable, Dict, Set
 
 import grpc
@@ -117,13 +116,31 @@ def reset(self):
 class StreamCancelHandler:
     def __init__(self):
         self.response_iterator = None
+        self._cancel_event = threading.Event()
+        self._lock = threading.Lock()
 
     def set_iterator(self, iterator):
-        self.response_iterator = iterator
+        with self._lock:
+            self.response_iterator = iterator
+            # If already cancelled, cancel the iterator immediately to avoid race
+            if self._cancel_event.is_set() and hasattr(iterator, "cancel"):
+                try:
+                    iterator.cancel()
+                except Exception:
+                    pass
 
     def cancel(self):
-        if self.response_iterator:
-            self.response_iterator.cancel()
+        self._cancel_event.set()
+        with self._lock:
+            if self.response_iterator:
+                self.response_iterator.cancel()
+
+    def is_cancelled(self) -> bool:
+        return self._cancel_event.is_set()
+
+    def wait_cancelled(self, timeout: float) -> bool:
+        """Waits until the handler is cancelled or timeout is reached."""
+        return self._cancel_event.wait(timeout)
 
 
 class WorkloadApiClient:
@@ -298,15 +315,8 @@ def fetch_jwt_bundles(self) -> JwtBundleSet:
             FetchJwtBundleError: In case there is an error in fetching the JWT-Bundle from the Workload API or
                                 in case the set of jwt_authorities cannot be parsed from the Workload API Response.
         """
-
-        responses = self._spiffe_workload_api_stub.FetchJWTBundles(
-            workload_pb2.JWTBundlesRequest(), timeout=10
-        )
-        res = next(responses)
-        jwt_bundles: Dict[TrustDomain, JwtBundle] = self._create_td_jwt_bundle_dict(res)
-        if not jwt_bundles:
-            raise FetchJwtBundleError('JWT Bundles response is empty')
-
+        response = self._call_fetch_jwt_bundles()
+        jwt_bundles: Dict[TrustDomain, JwtBundle] = self._create_td_jwt_bundle_dict(response)
         return JwtBundleSet(jwt_bundles)
 
     @handle_error(error_cls=ValidateJwtSvidError)
@@ -442,13 +452,17 @@ def _watch_x509_context_updates(
         on_error: Callable[[Exception], None],
     ):
         while True:
+            if cancel_handler.is_cancelled():
+                break
             try:
                 response_iterator = self._spiffe_workload_api_stub.FetchX509SVID(
                     workload_pb2.X509SVIDRequest()
                 )
                 cancel_handler.set_iterator(response_iterator)
 
                 for item in response_iterator:
+                    if cancel_handler.is_cancelled():
+                        break
                     x509_context = self._process_x509_context(item)
                     on_success(x509_context)
 
@@ -461,7 +475,9 @@ def _watch_x509_context_updates(
                     on_error(WorkloadApiError(f"gRPC error: {str(grpc_err.code())}"))
                     break
 
-                time.sleep(retry_handler.get_backoff())
+                backoff = retry_handler.get_backoff()
+                if cancel_handler.wait_cancelled(backoff):
+                    break
 
             except Exception as err:
                 on_error(WorkloadApiError(str(err)))
@@ -475,13 +491,17 @@ def _watch_jwt_bundles_updates(
         on_error: Callable[[Exception], None],
     ):
         while True:
+            if cancel_handler.is_cancelled():
+                break
             try:
                 response_iterator = self._spiffe_workload_api_stub.FetchJWTBundles(
                     workload_pb2.JWTBundlesRequest()
                 )
                 cancel_handler.set_iterator(response_iterator)
 
                 for item in response_iterator:
+                    if cancel_handler.is_cancelled():
+                        break
                     jwt_bundles = self._process_jwt_bundles(item)
                     on_success(jwt_bundles)
 
@@ -494,7 +514,9 @@ def _watch_jwt_bundles_updates(
                     on_error(WorkloadApiError(f"gRPC error: {str(grpc_err.code())}"))
                     break
 
-                time.sleep(retry_handler.get_backoff())
+                backoff = retry_handler.get_backoff()
+                if cancel_handler.wait_cancelled(backoff):
+                    break
 
             except Exception as err:
                 on_error(WorkloadApiError(str(err)))
@@ -520,7 +542,8 @@ def _process_jwt_bundles(
         return self._create_jwt_bundle_set(jwt_bundles_response.bundles)
 
     def _get_spiffe_grpc_channel(self) -> grpc.Channel:
-        grpc_insecure_channel = grpc.insecure_channel(self._config.spiffe_endpoint_socket)
+        target = self._grpc_target(self._config.spiffe_endpoint_socket)
+        grpc_insecure_channel = grpc.insecure_channel(target)
         spiffe_client_interceptor = (
             header_manipulator_client_interceptor.header_adder_interceptor(
                 WORKLOAD_API_HEADER_KEY, WORKLOAD_API_HEADER_VALUE
@@ -551,6 +574,18 @@ def _call_fetch_x509_bundles(self) -> workload_pb2.X509BundlesResponse:
             raise FetchX509BundleError('X.509 Bundles response is empty')
         return item
 
+    def _call_fetch_jwt_bundles(self) -> workload_pb2.JWTBundlesResponse:
+        response = self._spiffe_workload_api_stub.FetchJWTBundles(
+            workload_pb2.JWTBundlesRequest()
+        )
+        try:
+            item = next(response)
+        except StopIteration:
+            raise FetchJwtBundleError('JWT Bundles response is invalid')
+        if len(item.bundles) == 0:
+            raise FetchJwtBundleError('JWT Bundles response is empty')
+        return item
+
     @staticmethod
     def _create_x509_bundle_set(resp_bundles: Mapping[str, bytes]) -> X509BundleSet:
         x509_bundles = [
@@ -588,7 +623,30 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
 
     @staticmethod
     def _check_spiffe_socket_exists(spiffe_socket: str) -> None:
-        if spiffe_socket.startswith('unix:'):
-            spiffe_socket = spiffe_socket[5:]
-        if not os.path.exists(spiffe_socket):
-            raise ArgumentError(f'SPIFFE socket file "{spiffe_socket}" does not exist.')
+        path_to_check = WorkloadApiClient._strip_unix_scheme(spiffe_socket)
+        if not path_to_check:
+            raise ArgumentError('SPIFFE endpoint socket is empty')
+        if not os.path.exists(path_to_check):
+            raise ArgumentError(f'SPIFFE socket file "{path_to_check}" does not exist.')
+
+    @staticmethod
+    def _grpc_target(value: str) -> str:
+        """Returns the gRPC target for UDS, normalizing unix:/// to unix:/."""
+        if value.startswith('unix:'):
+            path = value[5:]
+            if path.startswith('/'):
+                path = '/' + path.lstrip('/')
+            return f'unix:{path}'
+        if value.startswith('/'):
+            return f'unix:{value}'
+        raise ArgumentError(
+            f'Invalid SPIFFE endpoint socket "{value}": only unix domain sockets are supported'
+        )
+
+    @staticmethod
+    def _strip_unix_scheme(value: str) -> str:
+        """Strips unix: scheme and normalizes leading slashes for filesystem checks."""
+        path = value[5:] if value.startswith('unix:') else value
+        if path.startswith('/'):
+            path = '/' + path.lstrip('/')
+        return path