Beatrust techBlog

Beatrust 株式会社の公式開発者ブログです。

NestJS + Prisma2 で作る RLS の世界

Beatrust で SRE をやっている Yuta(中川 裕太)です.運用がラクにできように色々と改善したり,セキュリティ向上したり,インフラ作ったり API 開発したりしています. 今回のブログでは,Beatrust のセキュリティ向上施策の一つである Row Level Security(RLS)について,実装時に課題になった点とその解決策,また RLS を適用して約3ヶ月運用してくる中で課題となった点とその解決策について紹介します.

TL;DR

  • 大企業に求められるセキュリティ要件を満たすためデータを論理的に分離しました.その実装方法として,実装・運用コストやセキュリティ向上効果などから RLS を選択しました.
  • index を適切に張れば 数10msec 程度の性能劣化で RLS 自体は適用できました.一方で,本来フレームワークが想定する DB connection の管理方法とは異なる方法で実装しているため,そこでの課題が顕著でした.
  • RLS を適用する際は,フレームワークが DB connection を独自に管理・拡張できる機構を持っていることがポイントになるようです.

モチベーション

Beatrust の大きな特徴は,創業まだ間もないスタートアップでありながら,大企業のお客様が抱える課題を解決しようとしている点です.そのため,お客様の多くは大企業であり,必然的に求められるセキュリティのレベルは高くなっています.そこで,データをデータベースレイヤでお客様ごとに論理的に分離するRow Level Security(RLS)と呼ばれる手法を適用することにしました.お客様のデータを論理的に分離する方法は他にも存在しますが,実装コストや運用コスト,それによって得られるセキュリティ向上効果などを鑑みて RLS を採用しました.その検証過程で以下の記事は非常に参考になりましたし,@yudppp さん本人にも色々と助言をいただき,この場をお借りしてお礼申し上げたいと思います.ありがとうございました.

Row Level Securityはマルチテナントの銀の弾丸になりうるのか / Row Level Security is silver bullet for multitenancy? - Speaker Deck

NestJS + Prisma2 による RLS 適用

Beatrust People では,NestJS + Prisma2 で実装されたモノリスでサービスを運用しています.そこで,これらのフレームワーク上で RLS を適用していきます.

データベースレイヤ

まずはデータベースレイヤに RLS を適用していきます.Beatrust ではデータベースとして PostgreSQL を用いているのですが,PostgreSQL に RLS を適用する方法としては大きく2つあります.前提として各テーブルにはアクセスできるテナントを指定した tenant という列を持つとします.

  • tenant = current_user といった形でポリシーを定義し,接続している DB user と一致する tenant を持つ行のみにアクセスできるようにする方法.
  • tenant = current_setting('app.current_tenant', true) といった形でポリシーを定義し,実行時に SET app.current_tenant = 'xxx' といった SQL を発行し,この xxx と一致する tenant を持つ行のみにアクセスできるようにする方法.

Beatrust People の場合,お客様(organization)単位がテナント(tenant)単位になるため,お客様が増える度に DB user を増やす必要がある前者の方法では運用コストがかかってしまいます.そこで,実行時に接続テナントをアプリケーションサイドから指定できる後者の方法を採用しました.これにより,アプリケーション側から明示的に organization を指定しない限りはデータを取得できない状態になります.

具体的には以下のようなポリシーを全テーブルに実装しています.

organizationId を列に持つシンプルなケース
ALTER TABLE "User" ENABLE ROW LEVEL SECURITY;

CREATE POLICY "UserIsolationPolicy" ON "User" USING (
    "organizationId" = current_setting('app.currentOrganizationId', true)
);
organizationId を列に持たず外部キーで制御する必要のある複雑なケース
ALTER TABLE "UserDetailItem" ENABLE ROW LEVEL SECURITY;

CREATE POLICY "UserDetailItemIsolationPolicy" ON "UserDetailItem" USING (
    "userId" IN (
        SELECT "id" FROM "User" WHERE "organizationId" = current_setting('app.currentOrganizationId', true)
    )
);

アプリケーションレイヤ

上記のように実行時に SET 文を発行するアプローチを採用したため,NestJS + Prisma2 にその機構を実装する必要があります.つまり,リクエストごとに独立した DB connection を張り,リクエストの開始時に SET 文を発行し app.currentOrgaizationId に適切な organizationId を設定し,完了時にリリースするという機構です.

NestJS には injection scopes という考え方があり,Dependency Injection するインスタンスのライフサイクルを指定することができます.これを用いて Scope.REQUEST を指定することでリクエストごとに独立した DB connection を実現することができます.

具体的には以下のように request scope で ContextService から organizationId を取得し各クエリの実行前に SET 文を発行するよう PrismaClient を拡張したクラスを作っています.ここには,各クエリの実行前に任意の処理を行う Prisma2 の Middleware の仕組み を使っています.また,ContextService は request scope のシンプルな key value store として実装しています.さらに,NestJS の MiddleWare の仕組み でリクエストごとに JWT から organizationIdContextService に詰め込むようにしています.

@Injectable({ scope: Scope.REQUEST })
export class PrismaService extends PrismaClient {
  constructor(private readonly contextService: ContextService) {
    super();
    this.$use(async (params, next) => {
      if ((params.action !== 'executeRaw') || (params.args.query && !params.args.query.startsWith('SET'))) {
        const organizationId = this.contextService.get('organizationId');
        await this.$executeRaw(`SET app.currentOrganizationId='${organizationId}';`);
      }
      return next(params);
    });
  }
}

しかし,これだけだと DB connection が close されず,リクエストが来る度に connection が増えてしまいます.NestJS + Primsa の公式サンプル では Lifecycle Events で DB connection を管理する方法で実装していますが,NestJS では injection scope が request scope の場合に Lifecycle Events が発火されないという仕様であるため,このアプローチを取ることができません.そこで,実際のリクエストを処理するロジックの前後に任意の処理を実行できる Interceptors という機構を利用して DB connection の close を行うようにしました.

@Injectable({ scope: Scope.REQUEST })
export class PrismaInterceptor implements NestInterceptor {
  constructor(private readonly prisma: PrismaService) {}

  intercept(
    context: ExecutionContext,
    next: CallHandler<any>
  ): Observable<any> {
    return next.handle().pipe(
      finalize(() => this.prisma.$disconnect())
    );
  }
}

DB connection が request scope だからこその課題

ここまでは事前に検証してあったこともあり比較的スムーズに実装することができました.しかし,実装を進めていくに従い,最初の検証時には思いもしなかった DB connection を request scope にしたからこその課題に遭遇しました.

単体テストで request scope が使えない

一般的に NestJS では以下ようにテストケースごとに一意な service instance を取得してテストするのが お作法なようです.

const moduleRef = await Test.createTestingModule({
  controllers: [CatsController],
  providers: [CatsService],
}).compile();

catsService = await moduleRef.resolve(CatsService);

ただし,上記の request scope だと Lifecycle Events が発火されないことにも関連するのですが,request scope のクラスは resolve しても一意なインスタンスを取得できない仕様なようで,このアプローチでは正しくテストを実施することができませんでした.そこで,NestJS のコミュニティで質問してみたところ,request scope のクラスを mock するアプローチを推奨されました.

一方で,そのアプローチを選択した場合,既存のテストコードの多くを書き換える必要があり,あまり現実的ではありませんでした.そこで,いろいろと調査してみると injection scope を定義側で変えてしまうアプローチを発見しこちらを採用することにしました.

@Injectable({
  scope: process.env.NODE_ENV === 'test' ? Scope.DEFAULT : Scope.REQUEST
})
export class PrismaService extends PrismaClient {

これで問題なく既存のテストコードは動作するようにはなりました.

全 organization をまたぐ処理との相性が悪い

冒頭にも記載したのですが,Beatrust People は複数機能が1コのアプリケーションに集約されたモノリスとなっています.その複数機能の中には全 organization をまたいで実行される admin 処理やバッチ処理も含まれています.一方で,RLS 適用の思想は request ごとに触れる organization を限定しましょうというもので,思想の段階でコンフリクトが発生してしまっています.そこで半ば強引なアプローチではありますが,getOrganizationList という DB 以外の場所から organization の一覧を取得する関数を作り,以下のようにコントローラレイヤである種のリクエスト処理を実行するような形で実装しています.

const organizationList = await this.organizationsService.getOrganizationList();
for (const { id } of organizationList) {
  this.contextService.set('organizationId', id);
  // 実際の処理
}

非同期処理だと DB connection が途中で切れてしまう

request scope で RLS を扱うという方針を愚直に実装した結果,非同期処理を行う時,処理の途中で DB connection が切れてしまう事象が発生しました.つまり,もし,非同期にリクエストを処理した場合,リクエストが完了した後に DB へのアクセスがある可能性があり,その場合,Interceptors が DB connection を close してしまっていて DB へのアクセスが失敗するケースが発生してしまいます.そうすると,非同期処理という Node.js の一つのメリットを失ってしまうことになります.そこで,非同期処理の中でもちゃんと DB connection を管理できるよう,ガベージコレクションなどで用いられる参照カウントを用いるすることにしました.確かに非同期処理の最後で必ず DB connection を close する処理を入れるようにすれば参照カウントなどのちょっと複雑なロジックを実装する必要はありません.一方で,非同期処理の最後に DB connection を close する処理を入れ忘れた場合,時限爆弾のように DB connection を逼迫し運用中に謎のエラーを発生させる可能性があるため避けたい実装でした.そのため,複雑ではあるものの,確実に DB connection を管理できる参照カウントのアプローチを採用しました.具体的には上記で SET 文を発行し単純に return next(params); で次の SQL 文を実行していたところを,処理の前後で参照カウントを増減させ参照カウントが0になった時のみ実際に DB connection を close するという実装にしました.また,Interceptors では直接 DB connection を close するのではなく,stopService を呼び出すことでサービスクラスでの DB 処理が終わった段階で connection を close してよいことを伝えるようにしました.

@Injectable({ scope: Scope.REQUEST })
export class PrismaService extends PrismaClient {
  constructor(private readonly contextService: ContextService) {
    super();
    this.$use(async (params, next) => {
      this.actionCount += 1;
      if (!this.isRlsSet(params)) await this.setRls();
      const result = await next(params);
      this.actionCount -= 1;
      if (this.shouldStop) this.stop();
      return result;
    });
  }

  private stop() {
    setImmediate(() => {
      if (!this.shouldStop || this.actionCount !== 0) {
        return;
      }
      this.$disconnect();
    });
  }

  stopService() {
    this.shouldStop = true;
    this.stop();
  }
}

運用の中でぶちあたった課題

これで晴れて NestJS + Prisma2 での RLS 適用が実装でき,無事にリリースすることができました.ここからはリリースした後,運用をする中でぶちあたった課題,パフォーマンス劣化とメモリリークについて紹介していきます.

特定のクエリが遅い

RLS を適用することでその分処理が増えるため API レスポンスが遅くなることは想定していましたが,それにしても遅すぎるクエリが何本かあり,複数 API を叩くようなケースで UI がもっさりしている印象を与え UX が劣化したように感じるところがありました.調査してみると,以下のように外部キーに対して RLS を適用しているケースでクエリが極端に遅くなっていることが判明しました.

CREATE POLICY "UserDetailItemIsolationPolicy" ON "UserDetailItem" USING (
    "userId" IN (
        SELECT "id" FROM "User" WHERE "organizationId" = current_setting('app.currentOrganizationId', true)
    )
);

このようなケースにおいて実行計画を見てみると以下のようにかなり実行コストがかかっていることが分かりました.

=> EXPLAIN ANALYZE SELECT * FROM "UserDetailItem";
                                                         QUERY PLAN                                                         
----------------------------------------------------------------------------------------------------------------------------
 Seq Scan on "UserDetailItem"  (cost=3313.43..5359.22 rows=10112 width=679) (actual time=56.450..63.843 rows=17451 loops=1)
   Filter: (hashed SubPlan 1)
   Rows Removed by Filter: 2772
   SubPlan 1
     ->  Seq Scan on "User"  (cost=0.00..3193.39 rows=48014 width=27) (actual time=0.022..34.300 rows=48070 loops=1)
           Filter: ("organizationId" = current_setting('app.currentOrganizationId'::text, true))
           Rows Removed by Filter: 2023
 Planning Time: 0.336 ms
 Execution Time: 64.323 ms
(9 rows)

そこで,User.organizationIdUserDetailItem.userId に index を張った上で再度実行計画を見てみると実行コストが 1/100 ほどに減少しました.

=> EXPLAIN ANALYZE SELECT * FROM "UserDetailItem";
----------------------------------------------------------------------------------------------------------------------------------------------
 Seq Scan on "UserDetailItem"  (cost=27.31..2073.10 rows=10112 width=679) (actual time=0.131..8.180 rows=30 loops=1)
   Filter: (hashed SubPlan 1)
   Rows Removed by Filter: 20193
   SubPlan 1
     ->  Index Scan using "User_organizationId_idx" on "User"  (cost=0.29..27.25 rows=25 width=27) (actual time=0.037..0.065 rows=26 loops=1)
           Index Cond: ("organizationId" = current_setting('app.currentOrganizationId'::text, true))
 Planning Time: 0.313 ms
 Execution Time: 8.238 ms
(8 rows)

上記のような形で外部キーに対して RLS を適用しているケースについて適切に index を張った結果,よほど複雑なクエリでない限り数10msec 以内には結果を返すようになり,UI のもっさり感も改善されました.

アクセスが集中すると極端にレイテンシーが劣化する

Beatrust People へのアクセスはある種のスパイク性があり,通常時は 1sec 以内にレスポンスを返していたエンドポイントが 3req/sec を超えるあたりから,急激に 20〜30sec までレイテンシーが劣化するというインシデントが RLS 適用後に発生しました.暫定対処としては,インスタンス数を数倍まで増やすことで対応はできたのですが,そのままではコストが非常にかかってしまうため原因調査と恒久対処を実施しました.調査してみると,DB connection を request scope にしたことの副作用であることが判明しました.多くのサービスクラスで DB connection を扱う必要があるため,request scope である PrismaService への依存が必須でした.NestJS では当然といえば当然なのですが,request scope のサービスクラスに依存するクラスはすべからく request scope になってしまいます.そのため,ほぼ全ての PrismaService に依存するサービスクラスはリクエストの度に生成されるようになってしまいました.一部サービスクラスの中で GCP のクライアントを初期化しているのですが,この初期化処理が例えば 3req/sec の頻度で走ることで GCP へのアクセスが集中しすぎ IO wait で極端にレイテンシーが劣化していることが判明しました.そのため,クリーンアーキテクチャで言うところの infrastructure 層のようなものを作り,GCP のクライアントクラスを default scope で作成し,各サービスクラスはそのクライアントクラスに依存するような構成に変更しました.これにより,以下のようにアクセス集中時のレイテンシーの極端な劣化を改善することができました.

改善前
$ echo "GET localhost:8080" | vegeta attack -rate=10/s -duration=1m | vegeta report
Requests      [total, rate, throughput]         600, 10.02, 5.98
Duration      [total, attack, wait]             1m17s, 59.9s, 17.223s
Latencies     [min, mean, 50, 90, 95, 99, max]  524.745ms, 12.579s, 10.332s, 20.469s, 20.607s, 20.832s, 21.123s
Bytes In      [total, mean]                     52273, 87.12
Bytes Out     [total, mean]                     0, 0.00
Success       [ratio]                           76.83%
Status Codes  [code:count]                      200:461  500:139  
Error Set:
500 Internal Server Error
改善後
$ echo "GET localhost:8080" | vegeta attack -rate=10/s -duration=1m | vegeta report
Requests      [total, rate, throughput]         600, 10.02, 10.00
Duration      [total, attack, wait]             59.993s, 59.9s, 93.033ms
Latencies     [min, mean, 50, 90, 95, 99, max]  84.263ms, 103.318ms, 91.247ms, 108.46ms, 195.897ms, 304.007ms, 392.469ms
Bytes In      [total, mean]                     9600, 16.00
Bytes Out     [total, mean]                     0, 0.00
Success       [ratio]                           100.00%
Status Codes  [code:count]                      200:600  
Error Set:

メモリリーク

RLS を適用した後,しばらく運用していると,メモリリークのような挙動が見て取れました.そこで,stg 環境にて過負荷をかけ続けたところ以下のように明らかなメモリリークを観測することができました.

そこで原因を調査してみると,上記のレイテンシー劣化の話と関連して,少なからずメモリリークしているライブラリたちがリクエストごとに生成と破棄を繰り返すことで,メモリリークの影響がより顕著になったということのようでした.そこで,以下2つのアプローチを取りました.

  • ライブラリをバージョンアップし,なるべくメモリリークが改善していることを期待する.
  • それでも解決しないものは上記と同様のアプローチで default scope でなるべく生成するようにする.

これにより,以下のように完全にメモリリークを解消できたわけではないですが,かなり改善することはできました.

さらなる改善

RLS を適用してから約3ヶ月が経ち,いくつかのトラブルはありましたが,現状では問題なく運用できています.ここからは今後進めていきたい改善ポイントについて紹介します.

Prisma2(ORM)の見直し

ORM として使っている Prisma2 は Query engine を別プロセスで立ち上げることで connection pooling などの面倒な処理を engine process で管理し JavaScript のクライアント側で実施すべき処理を軽減するアプローチを取っています.これ自身は Prisma2 をシンプルに使う場合にはよくできた仕組みなのですが,RLS のように自身で DB connection を管理したい場合に JavaScript のクライアント側からいじる方法がないことを意味しています.そのため,上記のようにややむりくり感のある実装をする必要がありました.

実装上の問題だけであれば,まだ何とかなりますが,実はここにパフォーマンス上の問題もはらんでいます.Prisma2 の設計思想の中で DB connection を request scope で管理するため,今まで記載してきたように Prisma2 の client を request scope で生成するアプローチを選んでいます.これはつまり,リクエストごとに Query engine もその都度プロセスとして生成されることを意味しています.プロセスを生成することは一般的にコストが高く,現状で実はリクエストにかかる処理の大半は SQL の実行ではなく,このプロセスと DB connection の生成時間となっています.さらに,これはリクエストごとに DB connection が生成されることを意味しており,DB connection 数がコントロールできないことを示しています.現状のアクセス数では DB インスタンスのコネクションを使い切るほどのアクセスは来ていないため問題ありませんが,将来的にサービスが成長していく過程で確実にボトルネックになることは見えているのでぜひとも改善したいポイントとなっています.

Prisma2 を使う限り,設計思想レベルでのコンフリクトがあるため,例えば Prisma が version 1⇒2 で大きく思想が変わったように 2⇒3 で大きく思想が変わる可能性もあるため,そこは定期的にウォッチしつつ,DB connection の管理方法を自分たちで拡張できるような ORM を継続的に探し検証しつつ,適する ORM が見つかったら適宜移行していこうと考えています.

節度あるマイクロサービス化

上記のように複数 organization をまたぐ admin 処理やバッチ処理を実現するため,organizationId の一覧を DB とそれ以外の場所の2ヶ所で管理するというアプローチを取っています.このアプローチを採用した以上,将来的に負債化する可能性が高く改善することが必須だと考えています.確かに,早すぎるマイクロサービス化はアンチパターンと言われますが,まずは admin 処理やバッチ処理から先行して,サービス境界を見極めマイクロサービス化を進めていこうと考えています.サービス境界をちゃんと整理することでサービスが触るべきデータが明確化され DB user の持つべき権限も見えてくると考えています.それに従い,RLS のポリシーを再実装しサービスも分割していくことで,技術的負債を解消できるかなと考えています.また,そうすることで実装が simple になり,誤った実装によるセキュリティ劣化も副次的に防ぐことになり,よりセキュアなサービスを作ることができると信じています.

具体的には,organization の一覧を触るべきサービスと organization に閉じた処理を実施すべきサービスを分けていこうと考えています.さらに,organization をまたぐ処理については,前者の organization の一覧を触るサービスから後者の organization に閉じたサービスを呼び出すようなアプローチを検討しています.

最後に

Beatrust に入社して最初の大仕事である Beatrust People への RLS 適用ですが,無事にリリースでき今では安定して運用できるレベルになって本当によかったなと感じています.実装を始めた当初の一番の懸念は RLS を適用することで SQL の実行が遅くなりサービス全体で速度的な劣化が見られることでした.しかし,実際に実装してみると,そこのボトルネックは 100msec にも満たない程度で BtoC のように msec 単位でレイテンシーを追っているサービスでない限り,適切な index を張りさえすればそこまで死活問題ではないことが分かりました.一方で,NestJS + Prisma2 というフレームワークと RLS の相性の悪さから DB connection をやや無理のある方法で管理せざるを得なかったことが実装上も性能上も問題となったことは大きな学びでした.やはり,RLS のような低レイヤーをいじる必要がある場合には,easy なフレームワークではなく simple なフレームワークを選択すべきだなと改めて感じました.今後,この記事を読んで RLS を実装しようと考えている方は,ぜひとも自分たちが採用しているフレームワークで DB connection を独自に管理・拡張できる機構が備わっているかに気をつけていただけると記事を書いた意味があって嬉しいなと考えています.

Beatrust では,大企業に求められるセキュリティレベルを担保しつつ,創業間もないスタートアップだからこそ爆速で仮説検証を回さなければならない,というとてもユニークなチャレンジをしている会社だと自負しています.もしこの記事や Beatrust 自体に興味を持っていただけたら,以下からお声がけいただけると幸いです.カジュアル面談のご希望もお気軽にぜひ!

beatrust.com