grpc custom authentication scheme example (python)

grpc_client.py

import grpc
import helloworld_pb2
import helloworld_pb2_grpc

# you need to use secure port,
# otherwise call credentials won't be transmitted
def run():
    with open('server.crt', 'rb') as f:
        trusted_certs = f.read()

    credentials = grpc.ssl_channel_credentials(root_certificates=trusted_certs)
    # make sure that all headers are in lowecase, otherwise grpc throws an exception
    call_credentials = grpc.metadata_call_credentials(
        lambda context, callback: callback((("x-custom-token", "SampleToken"),), None))
    # use this if you want standard "Authorization" header
    #call_credentials = grpc.access_token_call_credentials("test_access_token")
    composite_credentials = grpc.composite_channel_credentials(credentials, call_credentials)
    channel = grpc.secure_channel('{}:{}'.format('localhost', 50051), composite_credentials)
    stub = helloworld_pb2_grpc.GreeterStub(channel)

    response = stub.SayHello(helloworld_pb2.HelloRequest(name='Evgeny'))

    print("[~] Greeter client received: {}".format(response.message))
    response = stub.SayHelloAgain(helloworld_pb2.HelloRequest(name='Test'))
    print("[~] Greeter client received: {}".format(response.message))

if __name__ == '__main__':
    run()

grpc_server.py

import grpc
import helloworld_pb2
import helloworld_pb2_grpc

class CheckToken(object):
    def __init__(self, empty_response_cls, grpc_status_code=grpc.StatusCode.UNAUTHENTICATED, grpc_details=""):
        self._empty_rsp_cls = empty_response_cls
        self._status = grpc_status_code
        self._details = grpc_details

    def __call__(self, f):
        def wrapped_function(slf, request, context):
            meta = context.invocation_metadata()
            for item in meta:
                if item[0] == "x-custom-token" and item[1] == "SampleToken":
                    # use this if you want to update the context and pass some data to the method handler
                    # context.user_by_token = "sample_user" 
                    return f(slf, request, context)
            context.set_code(self._status)
            context.set_details(self._details)
            return self._empty_rsp_cls()
        return wrapped_function

      
class Greeter(helloworld_pb2_grpc.GreeterServicer):
    @CheckToken(helloworld_pb2.HelloReply, grpc.StatusCode.UNAUTHENTICATED, "Invalid token")
    def SayHello(self, request, context):
        # context.user_by_token is available here
        print("[~] SayHello call")
        return helloworld_pb2.HelloReply(message='Hello, {}!'.format(request.name))

    def SayHelloAgain(self, request, context):
        print("[~] SayHelloAgain call")
        return helloworld_pb2.HelloReply(message="Hello again, {}!".format(request.name))


def serve():
    with open('server.key', 'rb') as f:
        private_key = f.read()
    with open('server.crt', 'rb') as f:
        certificate_chain = f.read()

    server_credentials = grpc.ssl_server_credentials(((private_key, certificate_chain,),))

    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    helloworld_pb2_grpc.add_GreeterServicer_to_server(Greeter(), server)
    server.add_secure_port('[::]:50051', server_credentials)
    server.start()
    try:
        while True:
            time.sleep(_ONE_DAY_IN_SECONDS)
    except KeyboardInterrupt:
        server.stop(0)

if __name__ == '__main__':
    serve()