Kubernetes上のJupyter NotebookにOAuth2を追加する

TrueFoundry ユーザーは Jupyter Notebookを展開する AWS、Azure、GCPなどの個人クラウドアカウントで利用できます。この機能により、ユーザーは自身のマシンで機械学習の実験やトレーニングジョブを簡単に実行できます。当初、TrueFoundryを通じてデプロイされたノートブックは、ユーザー名とパスワードの組み合わせで保護されていました。しかし、広範なクライアントからの要望に応え、シングルサインオンを統合しました。これにより、ユーザーはTrueFoundryで使用しているのと同じログイン情報で、ノートブックに簡単にアクセスできるようになりました。このブログ記事では、この機能の実装方法について詳しく説明します。

TrueFoundry上のノートブック

TrueFoundryは、ノートブックのデプロイをオーケストレーションするために、Kubeflow Notebook Controllerのフォークを内部的に使用しています。このコントローラーは、私たちが活用している様々な機能を提供します。例えば：

1. ノートブック仕様の簡素化：Kubeflow Notebook APIはシンプルで、コントローラーがJupyter Notebookのデプロイ作成をオーケストレーションします。
2. 自動停止（Culling）：コントローラーは、一定期間の非アクティブ状態の後、ノートブックを自動的にシャットダウンします。これは、GPUマシンをバックエンドとするノートブックで実験を行うクライアントにとって非常に役立ちます。
3. 永続的なホームディレクトリ：コントローラーは、セッションをまたいでノートブック上のユーザーの進捗状況を保存する永続ボリュームの作成を処理します。
4. 拡張可能なベースイメージ：コントローラーは、TrueFoundryがメンテナンスするJupyter NotebookおよびVS Codeのベースノートブックイメージスイートをサポートしています。ユーザーは、スタートアップスクリプトを追加したり、特定のライブラリをインストールしたりすることで、これらのDockerイメージの機能を拡張できます。

参考までに、シンプルなKubeflow Notebookオブジェクトは次のようになります。

apiVersion: kubeflow.org/v1
kind: Notebook
metadata:
 name: my-notebook
spec:
 template:
   spec:
     containers:
       - name: my-notebook
         image: kubeflownotebookswg/jupyter:master
         args:
           [
             "start.sh",
             "lab",
             "--LabApp.token=''",
             "--LabApp.allow_remote_access='True'",
             "--LabApp.allow_root='True'",
             "--LabApp.ip='*'",
             "--LabApp.base_url=/test/my-notebook/",
             "--port=8888",
             "--no-browser",
           ]

ノートブックの基本認証

OAuth2を実装する前、TrueFoundryは、基本認証を統合することで、公開ノートブックのセキュリティを強化するオプションをユーザーに提供していました。この追加されたセキュリティ層は、許可された個人だけがこれらのノートブックの機密コンテンツにアクセスできるようにするために不可欠でした。この機能を実装するために、TrueFoundryは WebAssembly (Wasm) プラグイン をIstioプロキシ内で、特に Envoy プロキシ.

オープンソースのサービスメッシュであるIstioは、さまざまなサービスワークロード間のネットワーク通信を管理するためのフレームワークを提供します。Istioを使用することで、TrueFoundryはEnvoyプロキシによって管理されるネットワーク層にカスタムロジックを直接注入できるようになりました。このアプローチにより、Jupyter Notebooksとの間で流れるトラフィックを効果的に制御し、セキュリティを確保することができました。基本認証の実装の鍵となったのは、Envoyプロキシ内でWebAssemblyモジュールのデプロイを容易にするIstioの機能であるWasmPluginでした。

この基本認証WasmPluginは、Envoyプロキシ内のネットワークフィルターのシーケンスに統合されます。これらのフィルターは、アクセス制御、変換、データエンリッチメント、監査などに関連する高レベルの機能の実行を可能にし、サービスメッシュ全体のセキュリティと機能性を向上させます。Envoyフィルターチェーンに基本認証フィルターを追加するための仕様の簡略版を以下に示します。

apiVersion: extensions.istio.io/v1alpha1
kind: WasmPlugin
metadata:
 name: basic-auth
 namespace: istio-ingress
spec:
 phase: AUTHN
 pluginConfig:
   basic_auth_rules:
     - credentials:
         - user:pass
       hosts: www.example.com
       prefix: /secret/
 selector:
   matchLabels:
     istio: ingressgateway
 url: oci://ghcr.io/istio-ecosystem/wasm-extensions/basic_auth:1.12.0

ノートブック向けOAuth2

ノートブックにOAuth2を実装するにあたり、Envoyフィルターを利用しましたが、そのアプローチは基本認証の場合とは異なりました。基本認証では、事前に構築されたWasmPluginをフィルターチェーンに簡単に挿入できましたが、OAuth2には、よりカスタマイズされたソリューションが必要でした。これを実現するため、OAuth専用に設計されたHTTPフィルターを採用しました。TrueFoundryでは、当社のシングルサインオンシステムが FusionAuthと統合されており、当社のOAuthプロバイダーとして機能しています。

Envoyフィルターの仕様は以下のようになります。詳細については、ファイル内のコメントを参照してください。

apiVersion: networking.istio.io/v1alpha3
kind: EnvoyFilter
metadata:
 name: truefoundry-notebook-tfy-oauth2  # EnvoyFilterの名前
 namespace: auth-test  # EnvoyFilterがデプロイされるネームスペース
spec:
 workloadSelector:
   labels:
     truefoundry.com/application: truefoundry-notebook  # 特定のラベルを持つワークロードを対象とするセレクター
 configPatches:
 - applyTo: CLUSTER
   match:
     context: SIDECAR_OUTBOUND
   patch:
     operation: ADD
     value:
       name: tfy-oauth2  # OAuth2認証サービス用クラスター名
       type: LOGICAL_DNS  # サービスディスカバリのタイプ (DNS)
       connect_timeout: 5s  # 接続確立のタイムアウト
       lb_policy: ROUND_ROBIN  # ロードバランシングポリシー
       # その他のロードバランシング設定
 - applyTo: HTTP_FILTER
   match:
     context: SIDECAR_INBOUND
     listener:
       filterChain:
         filter:
           name: "envoy.filters.network.http_connection_manager"
           subFilter:
             name: envoy.filters.http.jwt_authn
   patch:
     operation: INSERT_BEFORE  # このフィルターをJWT認証フィルターの前に挿入
     value:
      name: envoy.filters.http.tfy-oauth  # OAuthフィルター名
      typed_config:
        "@type": type.googleapis.com/envoy.extensions.filters.http.oauth2.v3.OAuth2
        config:
          use_refresh_token: false  # リフレッシュトークンを使用するかどうか
          pass_through_matcher:
            - name: Authorization
              present_match: true  # Authorizationヘッダーが存在する場合にパススルーする
          forward_bearer_token: true  # ベアラートークンをアップストリームに転送する
          auth_type: BASIC_AUTH  # 使用される認証タイプ
          token_endpoint:
            cluster: tfy-oauth2  # トークンエンドポイントのクラスター
            uri: <token-endpoint-uri-of-oauth-provider>
            timeout: 5s  # トークンエンドポイントのタイムアウト
          authorization_endpoint: <authorization-endpoint-uri-of-oauth-provider>
          redirect_uri: https://%REQ(:authority)%/truefoundry-notebook/_auth/callback  # コールバック用のリダイレクトURI
          redirect_path_matcher:
            path:
              exact: /truefoundry-notebook/_auth/callback  # リダイレクトURIのパス
          signout_path:
            path:
              exact: /truefoundry-notebook/_auth/signout  # サインアウトのパス
          credentials:
            client_id: <oauthのクライアントID>
            token_secret:
               # トークンシークレットを取得するための設定
               # シークレットの取得方法についてはこちらをご覧ください:
               # https://www.envoyproxy.io/docs/envoy/latest/configuration/security/secret
            hmac_secret:
               # HMACを取得するための設定

ユーザーがOAuth2フィルターで保護されたサービスに初めてアクセスしようとすると、 authorization_endpointにリダイレクトされます。このエンドポイントは、外部OAuthプロバイダーのURLであり、当社の実装ではFusionAuthベースのTrueFoundryログインモーダルです。このリダイレクトはOAuthプロセスにおける重要なステップであり、ユーザーを安全な場所に誘導し、そこで認証を行い、サービスへのアクセスに必要な権限を付与できるようにします。

ログインが完了すると、FusionAuthはあなたを redirect_uri （フィルター仕様で設定されています）にリダイレクトし、そこに秘密の一時的な認証コードを追加します。このリクエストはフィルターによって傍受され、 token_endpointにリクエストを送信し、コードをJWTトークンと交換します。最後に、フィルターはJWTトークンを含むCookieを設定します。

Cookieが Authorization ヘッダーにJWTを値として設定するため、サービスへのその後のアクセスはHTTPフィルターを通過します。フィルターは、そのようなリクエストを通過させるように設定されています（参照 pass_through_matcher 仕様で）。JWTが有効なトークンであることを検証するために、OAuthプロバイダーに確認するRequestAuthenticationポリシーを作成します。

apiVersion: security.istio.io/v1beta1
kind: RequestAuthentication
metadata:
 # ...
spec:
 selector:
   # ...
 jwtRules:
 - issuer: "truefoundry.com"
   fromHeaders:
   - name: Authorization
     prefix: "Bearer "
   audiences:
     - <client-id>
   jwksUri: <oauth-provider-jwks-uri>
   forwardOriginalToken: true

最後に、どのリクエストに適用するかを指定するAuthorization Policyを追加します。 RequestAuthentication ポート8888へのすべてのリクエストに認証を適用したいと考えています。

apiVersion: security.istio.io/v1beta1
kind: AuthorizationPolicy
metadata:
 name: best-notebook-tfy-oauth2
 namespace: auth-test
spec:
 selector:
   matchLabels:
     truefoundry.com/application: best-notebook
 action: DENY
 rules:
 - from:
   - source:
       notRequestPrincipals: ["*"]
   to:
     - operation:
         ports:
           - "8888"