Skip to content

Commit

Permalink
fix(framework) Avoid processing requests from non-FleetStub (#4900)
Browse files Browse the repository at this point in the history
Co-authored-by: Javier <[email protected]>
  • Loading branch information
panh99 and jafermarq authored Feb 5, 2025
1 parent 2ec0561 commit cc00892
Showing 1 changed file with 12 additions and 3 deletions.
15 changes: 12 additions & 3 deletions src/py/flwr/server/superlink/fleet/grpc_rere/server_interceptor.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,11 @@
MAX_TIMESTAMP_DIFF = TIMESTAMP_TOLERANCE + SYSTEM_TIME_TOLERANCE


def _unary_unary_rpc_terminator(message: str) -> grpc.RpcMethodHandler:
def _unary_unary_rpc_terminator(
message: str, code: Any = grpc.StatusCode.UNAUTHENTICATED
) -> grpc.RpcMethodHandler:
def terminate(_request: GrpcMessage, context: grpc.ServicerContext) -> GrpcMessage:
context.abort(grpc.StatusCode.UNAUTHENTICATED, message)
context.abort(code, message)
raise RuntimeError("Should not reach this point") # Make mypy happy

return grpc.unary_unary_rpc_method_handler(terminate)
Expand All @@ -68,7 +70,7 @@ def __init__(self, state_factory: LinkStateFactory, auto_auth: bool = False):
self.state_factory = state_factory
self.auto_auth = auto_auth

def intercept_service(
def intercept_service( # pylint: disable=too-many-return-statements
self,
continuation: Callable[[Any], Any],
handler_call_details: grpc.HandlerCallDetails,
Expand All @@ -79,6 +81,13 @@ def intercept_service(
metadata sent by the node. Continue RPC call if node is authenticated, else,
terminate RPC call by setting context to abort.
"""
# Filter out non-Fleet service calls
if not handler_call_details.method.startswith("/flwr.proto.Fleet/"):
return _unary_unary_rpc_terminator(
"This request should be sent to a different service.",
grpc.StatusCode.FAILED_PRECONDITION,
)

state = self.state_factory.state()
metadata_dict = dict(handler_call_details.invocation_metadata)

Expand Down

0 comments on commit cc00892

Please sign in to comment.