3
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

【AWS】Managed Knowledge Baseでの検索フィルターを確認しました

3
Posted at

はじめに

Managed Knowledge Baseと検索フィルターを組み合わせていろいろ試してみました。

Gateway の Interceptor で検索リクエストを書き換えれば、検索の時点で絞り込めるはず。そう思って試したら、思っていたより制約が多かったので書き残しておきます。

前提

こういうのを試しました。
直でGatewayに聞きます。

image.png

Cognito に2人のユーザーを作成します。

ユーザー グループ 見せたい文書
hr-manager@example.com hr-manager public と staff
contractor@example.com contractor public だけ

Knowledge Base には4つの文書を入れて、それぞれに clearance というメタデータを付けました。営業時間と経費規程が public、給与レンジと人事評価が staff です。

メタデータは S3 にサイドカーファイルとして置きます。

{
  "metadataAttributes": {
    "clearance": {
      "value": { "type": "STRING", "stringValue": "staff" },
      "includeForEmbedding": false
    }
  }
}

includeForEmbedding を false にしたのは、clearance の文字列を検索の意味計算に混ぜたくないからです。true だと「staff」という語がクエリにマッチして順位が動いてしまい、アクセス制御としては筋が悪くなります。

やったこと

素の状態だと普通に漏れる

まず Interceptor なしで、contractor として給与を検索してみます。

curl -s -X POST "$GW" -H "Authorization: Bearer $TOKEN" \
  -H "Content-Type: application/json" \
  -H "Accept: application/json, text/event-stream" \
  -d '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{
       "name":"managedkb___Retrieve",
       "arguments":{"retrievalQuery":{"text":"salary band for engineers"}}}}'

結果は2件、どちらも staff の文書でした。給与レンジと人事評価がそのまま返ってきます。

ちなみにこれは検索が失敗しているのではなく、成功しています。給与について聞けば給与の文書が上位に来るのは当たり前で、RAG が正しく動いた結果として機密が出てくる。ここが厄介なところです。

Interceptor で filter を足す

Knowledge Base の Retrieve API は filter を受け取れます。これを Interceptor が途中で差し込みます。

エージェントが送ってくるのはクエリだけ。

"arguments": {
  "retrievalQuery": {"text": "salary band for engineers"}
}

Interceptor がこう書き換えて Knowledge Base に渡します。

図の赤枠の部分で書き換えてからKBに渡す感じです。

image.png

"arguments": {
  "retrievalQuery": {"text": "salary band for engineers"},
  "この下を追加する"
  "retrievalConfiguration": {
    "managedSearchConfiguration": {
      "filter": {"in": {"key": "clearance", "value": ["public"]}}
    }
  }
}

["public"] の部分は、リクエストに付いてきた JWT の cognito:groups から決めています。Interceptor はこれで全部です。

import base64
import json

ALLOWED_CLEARANCE = {
    "hr-manager": ["public", "staff"],
    "contractor": ["public"],
}


def _decode_jwt_payload(token):
    payload = token.split(".")[1]
    payload += "=" * (-len(payload) % 4)
    return json.loads(base64.urlsafe_b64decode(payload))


def _get_header(headers, name):
    # Gateway はヘッダ名を小文字にして渡してくる
    for k, v in (headers or {}).items():
        if k.lower() == name.lower():
            return v
    return ""


def _passthrough(body):
    return {
        "interceptorOutputVersion": "1.0",
        "mcp": {"transformedGatewayRequest": {"body": body}},
    }


def lambda_handler(event, context):
    request = event.get("mcp", {}).get("gatewayRequest", {})
    body = request.get("body", {})
    if body.get("method") != "tools/call":
        return _passthrough(body)

    # 検証済みトークンからロールを取り出す
    auth = _get_header(request.get("headers", {}), "Authorization")
    role = "unknown"
    if auth.startswith("Bearer "):
        claims = _decode_jwt_payload(auth[len("Bearer "):])
        groups = claims.get("cognito:groups") or []
        if groups:
            role = groups[0]
    clearance = ALLOWED_CLEARANCE.get(role, [])   # 未知のロールは何も見せない

    # 検索リクエストに filter を注入する
    params = body.setdefault("params", {})
    args = params.setdefault("arguments", {})
    cfg = args.setdefault("retrievalConfiguration", {})
    search = cfg.setdefault("managedSearchConfiguration", {})
    claimed = search.get("filter")
    search["filter"] = {"in": {"key": "clearance", "value": clearance}}

    print(f"ROLE    {role}")
    print(f"CLAIMED {json.dumps(claimed)}")
    print(f"INJECTED {json.dumps(search['filter'])}")

    return _passthrough(body)

ヘッダを読むために、ターゲット側で passRequestHeaders を true にしておく必要があります。デフォルトは false で、その場合 headers が渡ってきません。

なお署名の検証はしていません。ここに到達している時点で Gateway の JWT authorizer が公開鍵で検証を済ませているためです。改ざんしたトークンを投げると Interceptor は呼ばれずに 403 で終わります。CloudWatch のログに記録が1行も残らないので、順序も確認できました。

動いた

まったく同じクエリ「salary band for engineers」を、2人が投げます。リクエストの中身は1バイトも違いません。違うのは Authorization ヘッダのトークンだけです。

contractor の場合。

1件
  [public] expense-policy.txt
      Expense reimbursement policy. Expense reports are submitted through the ...

hr-manager の場合。

2件
  [staff] salary-bands.txt
      Confidential compensation data. Engineer salary band L4 is 9,000,000 to ...
  [staff] performance-reviews.txt
      Confidential performance records. In the 2026 review cycle, three employ...

給与レンジを聞いているので、hr-manager には給与文書が返ります。contractor には返らず、代わりに関連度が次点の経費規程が出てきます。

大事なのは、contractor に対して給与文書が「取得されてから捨てられた」のではないことです。Knowledge Base の検索候補に最初から入っていないので、ベクトル類似度の計算すらされていません。

CloudWatch のログを見ると、何が起きたか分かります。

ROLE    contractor
CLAIMED null
INJECTED {"in": {"key": "clearance", "value": ["public"]}}
ROLE    hr-manager
CLAIMED null
INJECTED {"in": {"key": "clearance", "value": ["public", "staff"]}}

CLAIMED null は「エージェントは filter を送ってきていない」、INJECTED が Interceptor の入れた値です。トークンのロールによって値が変わっています。

自分でフィルタを書いてきた場合

contractor が staff を狙って、自分でフィルタを組み立てて送ってみます。

"retrievalConfiguration": {"managedSearchConfiguration": {
  "filter": {"in": {"key": "clearance", "value": ["public", "staff"]}}}}

結果は変わりません。

1件
  [public] expense-policy.txt

ログには、送られてきた値がそのまま記録されています。

ROLE    contractor
CLAIMED {"in": {"value": ["public", "staff"], "key": "clearance"}}
INJECTED {"in": {"key": "clearance", "value": ["public"]}}

CLAIMED に嘘が乗っていて、INJECTED で正しい値に差し替わっている。

ポイントは、送られてきた値を検証していないことです。「この filter は妥当か」を判定しようとすると、orAll で入れ子にされたケースやキー名を変えてきたケースまで全部潰す必要が出てきます。読まずに捨てて上書きすれば、そもそも判定が要りません。

なお実運用でこれをやるなら、CLAIMED に値が入っていた時点で異常として記録しておくと監査に使えます。正しく作られたエージェントなら、そもそも filter を送ってこないはずなので。

ハマったところ

登録していないパスは書けない

最初、Interceptor が filter を足そうとしたらこう怒られました。

IllegalArgumentException - Invalid request parameters:
cannot set parameter(s): [/retrievalConfiguration]

コネクタターゲットは、設定してよい項目のホワイトリストを持っています。それが parameterOverrides です。ここに登録したパスしか、リクエストに含められません。

'parameterOverrides': [
    {'path': '$.retrievalQuery.text', 'visible': True},
    {'path': '$.retrievalConfiguration.managedSearchConfiguration.filter',
     'visible': True},
]

面白いのは、Lambda ターゲットだと真逆だったことです。ツールスキーマに定義していない totally_bogus みたいな項目でも、送れば Lambda まで素通りしました。同じ Gateway でもターゲットの種類で検証の厳しさが違います。

もうひとつ。Gateway は「Interceptor が書いた」か「エージェントが書いた」かを区別していません。同じリクエストとして同じ検証にかけます。Interceptor に特権はない、ということです。

さいごに

ブログは短ければ短いほどいいと最近感じています

3
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
3
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?