xref: /aosp_15_r20/external/grpc-grpc/examples/python/auth/async_customized_auth_server.py (revision cc02d7e222339f7a4f6ba5f422e6413f4bd931f2) 
1# Copyright 2020 The gRPC Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Server of the Python AsyncIO example of customizing authentication mechanism."""
15
16import argparse
17import asyncio
18import logging
19from typing import Awaitable, Callable, Tuple
20
21import _credentials
22import grpc
23
24helloworld_pb2, helloworld_pb2_grpc = grpc.protos_and_services(
25    "helloworld.proto"
26)
27
28_LOGGER = logging.getLogger(__name__)
29_LOGGER.setLevel(logging.INFO)
30
31_LISTEN_ADDRESS_TEMPLATE = "localhost:%d"
32_SIGNATURE_HEADER_KEY = "x-signature"
33
34
35class SignatureValidationInterceptor(grpc.aio.ServerInterceptor):
36    def __init__(self):
37        def abort(ignored_request, context: grpc.aio.ServicerContext) -> None:
38            context.abort(grpc.StatusCode.UNAUTHENTICATED, "Invalid signature")
39
40        self._abort_handler = grpc.unary_unary_rpc_method_handler(abort)
41
42    async def intercept_service(
43        self,
44        continuation: Callable[
45            [grpc.HandlerCallDetails], Awaitable[grpc.RpcMethodHandler]
46        ],
47        handler_call_details: grpc.HandlerCallDetails,
48    ) -> grpc.RpcMethodHandler:
49        # Example HandlerCallDetails object:
50        #     _HandlerCallDetails(
51        #       method=u'/helloworld.Greeter/SayHello',
52        #       invocation_metadata=...)
53        method_name = handler_call_details.method.split("/")[-1]
54        expected_metadata = (_SIGNATURE_HEADER_KEY, method_name[::-1])
55        if expected_metadata in handler_call_details.invocation_metadata:
56            return await continuation(handler_call_details)
57        else:
58            return self._abort_handler
59
60
61class SimpleGreeter(helloworld_pb2_grpc.GreeterServicer):
62    async def SayHello(
63        self, request: helloworld_pb2.HelloRequest, unused_context
64    ) -> helloworld_pb2.HelloReply:
65        return helloworld_pb2.HelloReply(message="Hello, %s!" % request.name)
66
67
68async def run_server(port: int) -> Tuple[grpc.aio.Server, int]:
69    # Bind interceptor to server
70    server = grpc.aio.server(interceptors=(SignatureValidationInterceptor(),))
71    helloworld_pb2_grpc.add_GreeterServicer_to_server(SimpleGreeter(), server)
72
73    # Loading credentials
74    server_credentials = grpc.ssl_server_credentials(
75        (
76            (
77                _credentials.SERVER_CERTIFICATE_KEY,
78                _credentials.SERVER_CERTIFICATE,
79            ),
80        )
81    )
82
83    # Pass down credentials
84    port = server.add_secure_port(
85        _LISTEN_ADDRESS_TEMPLATE % port, server_credentials
86    )
87
88    await server.start()
89    return server, port
90
91
92async def main() -> None:
93    parser = argparse.ArgumentParser()
94    parser.add_argument(
95        "--port", nargs="?", type=int, default=50051, help="the listening port"
96    )
97    args = parser.parse_args()
98
99    server, port = await run_server(args.port)
100    logging.info("Server is listening at port :%d", port)
101    await server.wait_for_termination()
102
103
104if __name__ == "__main__":
105    logging.basicConfig(level=logging.INFO)
106    asyncio.run(main())
107