@kenmaroです。
普段は主に秘密計算、準同型暗号などの記事について投稿しています。
秘密計算に関連するまとめの記事に関しては以下をご覧ください。

秘密計算エンジニアを始めて２年半が経った。

秘密計算エンジニアを始めて3年が経った。

概要

今回の記事は、

の続編となります。前回は、

grpc をPython で使う際の

SSL/TLS をgrpcに実装する
SSL/TLS + token によるAPIへのアクセスコントロールを実装する

というシナリオでチュートリアルを行いました。

今回は、前回の２つ目の

SSL/TLS + token によるAPIへのアクセスコントロールを実装する

を少し発展させ、

サーバ側に複数トークンを用意し、クライアントによってアクセスできるAPIを切り分ける

という構成のチュートリアルを行います。

コードはこちらに上げています。

今回の内容はブランチ multi_party
となるので、

git clone https://github.com/kenmaro3/grpc_auth.git
git checkout -b multi_party origin/multi_party

としてブランチを移動してください。
また、前回の内容を概ね引き継ぎますので、前回のチュートリアルを行った後に今回のものを行うと良いでしょう。

シナリオの解説

概要にて、

サーバ側に複数トークンを用意し、クライアントによってアクセスできるAPIを切り分ける

と書きましたが、例をあげて少し詳細を記述したいと思います。

今回のチュートリアルコードでは、

proto/myserver.proto

syntax = "proto3";

import "model/message.proto";

service MyServer{
  rpc test0(PB_Message) returns (PB_Message){};
  rpc test1(PB_Message) returns (PB_Message){};
}

として二つのメソッドを用意しています。

このとき、

	test0	test1
ユーザ０	アクセスできる	アクセスできない
ユーザ１	アクセスできる	アクセスできる

というようにアクセスコントロールをしたいと仮定し、これを今回のシナリオとします。

実装の内容

前回トークンを用いたアクセスコントロール自体は実装しました。
具体的には、

custom_header_token/controller.py

class AuthInterceptor(grpc.ServerInterceptor)

を定義しサーバサイドに実装したあと、
クライアントサイドからは

custom_header_token/client.py

    channel = grpc.secure_channel(
    'localhost:5555',
    grpc.composite_channel_credentials(
        grpc.ssl_channel_credentials(trusted_certs),
        grpc.metadata_call_credentials(
            GrpcAuth('mytoken')
        )
    )

としてトークンをコネクション確立時に設定していました。

トークンを複数受け入れて、メソッドによってアクセスを制限するには、以下のように実装を変更します。

custom_header_token/controller.py

class AuthInterceptor(grpc.ServerInterceptor):
    def __init__(self, key):
        self._valid_metadata = ('rpc-auth-header', key)

        def deny(_, context):
            context.abort(grpc.StatusCode.UNAUTHENTICATED, 'Invalid key')

        self._deny = grpc.unary_unary_rpc_method_handler(deny)


    def validate_party(self, party_index, meta, continuation, handler_call_details):
        if meta and meta[0].value == self._valid_metadata[1][party_index]:
            return continuation(handler_call_details)
            #return self._deny
        else:
            logging.info(f"denied: {meta[0].value} vs {self._valid_metadata[party_index]}, {self._valid_metadata}")
            return self._deny


    def intercept_service(self, continuation, handler_call_details):
        meta = handler_call_details.invocation_metadata

        method = handler_call_details.method.split("/")[2]

        if method == "test0":
            return self.validate_party(0, meta, continuation, handler_call_details)
        elif method == "test1":
            return self.validate_party(1, meta, continuation, handler_call_details)
        else:
            raise Exception(f"got method: {handler_call_details.method}, {method}")

ここで、ポイントとして、

intercept_service にて呼ばれているメソッドがなんなのか判断し、
validate_partyを使ってトークンが正しいものか判断

しています。

また、サーバーを立ち上げるときに

custom_header_token/controller.py

    server = grpc.server(
        futures.ThreadPoolExecutor(max_workers=constants.SERVER_MAX_WORKER),
        options=options,
        interceptors=(AuthInterceptor(['mytoken0', 'mytoken1']),)
        # interceptors=(AuthInterceptor('mytoken'),)
    )

として、２つのトークンをハードコードしています。
ここで、mytoken0はユーザ０用のトークンであり、
mytoken1はユーザ１用のトークンであることに言及しておきます。

ここまで準備できたら、client.pyを利用してテストしてみます。

custom_header_token/client.py

    channel = grpc.secure_channel(
        'localhost:5555',
        grpc.composite_channel_credentials(
            grpc.ssl_channel_credentials(trusted_certs),
            grpc.metadata_call_credentials(
                GrpcAuth('mytoken0')
            )
        )
    )

    # channel = grpc.insecure_channel('localhost:5555', options=[])
    stub = myserver_pb2_grpc.MyServerStub(channel)

    res = stub.test0(message_pb2.PB_Message())

mytoken0 はユーザ０が知っているトークンであり、
そのトークンを付加してコネクションを確立すると、test0にアクセスできるので、

hello, world
res: done!

と結果が返ってきます。

ここで、

custom_header_token/client.py

res = stub.test1(message_pb2.PB_Message())

とすると、mytoken0を持つユーザ０は、test1に対してアクセスが許可されていませんから、

status = StatusCode.UNAUTHENTICATED

のエラーが返却され、きちんとアクセスコントロールが実装できていることがわかります。

また、client.pyに渡すトークンを

custom_header_token/client.py

    channel = grpc.secure_channel(
        'localhost:5555',
        grpc.composite_channel_credentials(
            grpc.ssl_channel_credentials(trusted_certs),
            grpc.metadata_call_credentials(
                GrpcAuth('mytoken1')
            )
        )
    )

とすると、今度はtest1にはアクセスできるものの、
test0にはアクセスできないことが確認できます。

これで今回の実装は終了となります。お疲れ様でした。

まとめ

今回は、

サーバ側に複数トークンを用意し、クライアントによってアクセスできるAPIを切り分ける

という構成のチュートリアルを行いました。

１０分以内で試せるソースコードは

のmulti_partyブランチにコミットされていますのでぜひご覧ください。

ここまで来れば、トークン別にアクセスコントロールが実装されたgrpcサーバが作れますので、
結構いろんなシナリオで活用できそうです。

これ以上どう発展させようか少し悩みますが、現在ハードコードされている２つのトークンを、
AzureのKMSから取得してgrpcを立ち上げる
とかもありかなと思いました。

もし余裕があれば取り組んでみようと思っています。

今回はこの辺で。

@kenmaro