この記事はSpring Advent Calendar 2022の20日目の記事になりました。

前書き

こんにちは、justInCaseTechnorogiesでバックエンドエンジニアをしている宮田です。
弊チームでは、Spring WebFlux(Kotlin Coroutines) x R2DBCのマルチモジュールプロジェクトで、Row Level Securityを有効にし、かつ効率的に開発を行うための基盤を作成しました。
この記事ではその中で得た知見についてまとめます。

技術スタック

PostgreSQL 13
SpringBoot 2.7.3
Kotlin 1.7.10

Row Level Securityについて

本題に入る前に、Row Level Security（以降RLSと記載）について軽く解説します。

RLSは、簡単に言うと、識別子が一致しない行へのアクセスをDB側で禁止する機能です（DB側の機能であるため、この機能を提供していないDBでは利用できません）。
これを利用することで、「A社の操作によって発行されるクエリはA社以外のデータにアクセスできない」状態が実現できます。

これは、例えばSQLに渡す条件を取り違えたり、where句を付け忘れてしまった場合などに役立ちます。

fun findFooBy(fooId: Int): Foo { /* 何らかのクエリ処理 */ }

fun processFoo(fooId: Int, barId: Int) {
    val foo = findFooBy(barId) // 渡すIDを間違えた！！

    /* 取得できてしまうと処理は継続される... */
}

このような場合に、間違えた条件の先にデータが存在していると、そのデータが破壊されたり、不正なデータが画面に表示されたりといった事象が起きます。
場合によっては、「A社のデータがB社から閲覧できてしまう」ようなことも起きるでしょう。
データ破壊や情報流出は言うまでもなく特大の問題です。

一方、RLSを適用していれば、A社のデータはB社から取得できないため、このようなバグを作り込んでしまった場合にも問題の影響範囲を限定することができます。

アプリケーション側の実装について

ここからは、クライアントからのリクエストを処理する一連の流れを紹介します。

説明内では、RLS用の識別子を単にIdentifierと表現します。
この部分は、実際には識別子に合わせて適切な名前を付ける必要があります。

基本方針

シンプルな対応方針としては、以下のように、Identifierを引数として引き回し、クエリ発行時にそれぞれ設定する形も考えられます。

import io.r2dbc.spi.ConnectionFactory
import kotlinx.coroutines.reactive.awaitSingle
import org.springframework.r2dbc.connection.ConnectionFactoryUtils
import org.springframework.stereotype.Repository
import reactor.core.publisher.Mono
import reactor.kotlin.core.publisher.toMono

@Repository
class Repository(private val cfi: ConnectionFactory) {
    // コネクション発行 -> クエリ処理
    suspend fun query(identifier: String) {
        ConnectionFactoryUtils.getConnection(cfi)
            .flatMap { connection ->
                connection.createStatement("SET app.current_identifier = $1")
                    .bind(0, identifier)
                    .execute()
                    .toMono()
                    .thenReturn(connection)
            }
            .map { /* identifier設定後のコネクションを利用したクエリ処理 */ }
            .awaitSingle()
    }
}

一方、このやり方には以下のような欠点があり、何千と機能を作っていく上では良いやり方とは言えません。

処理が非効率になる
- 特にORMを利用する場合、発行したコネクションから一々初期化処理を行わなければならない
- Transaction中など、1度発行したコネクションを使い回す場合には、Identifierに関するクエリを何度も発行することになり非効率
全体的な記述が冗長になる
- 全クエリ処理にIdentifierの設定処理を書く必要が出る
- 全関数に引数としてIdentifierを設定するのは大変
- Identifierが処理に必須とは限らないため、RLSのために全体で引き回すことには違和感が有る

そこで、基本方針としてはRLSが設定されていない場合と同じ見た目になる（= 日常的なコーディングを行う上での労力が最低限になる）形を目指しました。

基本的な処理の流れ

基本的な処理の流れは以下のようになります。これらの基盤を実装することで、日常的に触れるコードではRLSが設定されていない場合と同じ見た目を実現できます。

スコープ内で利用するIdentifierをReactor Contextにセットする
1でセットしたIdentifierをReactor Contextから読み出す
2で取得したIdentifierをコネクションにセットする

モジュール配置としては、2と3は共通モジュールに、1はモジュール毎に個別の内容を配置する形になります。

以下、基本的な処理の流れに沿ってそれぞれの部分を紹介していきます。

1. スコープ内で利用する`Identifier`を`Reactor Context`にセットする

スコープ内で利用するIdentifierをReactor Contextにセットする際には、WebFilterを利用します。
簡単な例として、リクエストヘッダから文字列のIdentifierを読み取り、Reactor Contextにセットするコードは以下のようになります。

import org.springframework.core.Ordered
import org.springframework.core.annotation.Order
import org.springframework.stereotype.Component
import org.springframework.web.server.ServerWebExchange
import org.springframework.web.server.WebFilter
import org.springframework.web.server.WebFilterChain
import reactor.core.publisher.Mono

const val IDENTIFIER_KEY: String = /* Identifierを管理するためのキー、 */

@Component
@Order(Ordered.LOWEST_PRECEDENCE) // 認証周りより後に実行するため、順序は最低にする
class SetIdentifierFilter : WebFilter {
    // リクエストのヘッダーからIdentifierを抽出する関数
    private fun ServerWebExchange.getIdentifierFromHeader(): String =
        request.headers.getFirst(IDENTIFIER_KEY)!!

    override fun filter(exchange: ServerWebExchange, chain: WebFilterChain): Mono<Void> {
        val identifierMono: Mono<String> = Mono.just(exchange.getIdentifierFromHeader())
        return chain.filter(exchange).contextWrite { it.put(IDENTIFIER_KEY, identifierMono) }
    }
}

filter関数をJWT認証のトークンから読み出す形に修正した場合は以下のようになります。

import org.springframework.security.core.context.ReactiveSecurityContextHolder
import org.springframework.security.oauth2.server.resource.authentication.JwtAuthenticationToken

    /* class SetIdentifierFilter...は省略 */

    override fun filter(exchange: ServerWebExchange, chain: WebFilterChain): Mono<Void> {
        return chain.filter(exchange).contextWrite { context ->
            // ReactiveSecurityContextHolderはchainに連なる処理内でなければ正常に値を読み出せない
            val identifierMono: Mono<String> = ReactiveSecurityContextHolder.getContext().map {
                (it.authentication as JwtAuthenticationToken)
                    .token
                    .getClaimAsString(IDENTIFIER_KEY)
            }
            context.put(IDENTIFIER_KEY, identifierMono)
        }
    }

以下、これらのサンプルコードに関して解説していきます。

セットする内容について

サンプルコードではMono<String>をセットする形に統一しています。
これはReactiveSecurityContextHolder.getContextに連なる処理で得たMono<String>をStringにアンラップしてセットする方法が思いつかなかったためです。

これについては多くの部分で問題になりませんが、どちらかと言えばStringをセットできた方が嬉しい部分がありますので、分かる方がいらっしゃいましたらコメント頂けると助かります。

Reactor Contextへのアクセスに関する注意点

Reactor Contextは非リアクティブプログラミングにおけるThreadLocalのようなもので、リクエスト単位で固有の値を設定・読み出すことができます。

projectreactor.io

ただし、WebFilter内でReactor Contextにセットされた値を読み出す機能を利用する場合は、必ずchainに連なる処理の中で呼び出す必要があります。
サンプルコードでは、contextWrite内で呼び出しているReactiveSecurityContextHolder.getContextがそれです。

これは、WebFilterのReactor Contextとchain内のReactor Contextが異なるために必要な対応です。
アプリケーション内同様にReactor Contextにセットされた値を読み出したい場合は、chainに連なる処理の中で呼び出す必要があります。

認証等の情報から読み出す場合の注意点

他のWebFilterによってセットされる値（特に認証周りなど）からIdentifierを参照したい場合が有ります。
この場合、SetIdentifierFilterは他のWebFilterが設定された後に実行される必要があります。
サンプルコードでは、OrderアノテーションでSetIdentifierFilterの実行が最後になるよう制御しています。

2. 1でセットしたIdentifierをReactor Contextから読み出す

1でセットしたIdentifierをReactor Contextから読み出すコードは以下のようになります。

import kotlinx.coroutines.reactive.awaitSingle
import org.springframework.stereotype.Component
import reactor.core.publisher.Mono

@Component
class IdentifierGetter {
    // Identifier取得の共通関数
    // 有効なスコープ内で実行すればIdentifierが取得できる（コンポーネント内に配置しているのはモックの都合）
    fun getIdentifierMono(): Mono<String> = Mono.deferContextual {
        it.get<Mono<String>>(IDENTIFIER_KEY)
            .onErrorMap { e -> RuntimeException("コンテキストからIdentifierが取得できませんでした", e) }
    }

    suspend fun getIdentifier(): String = getIdentifierMono().awaitSingle()
}

コメントの通り、Component化している理由はモックを容易にするためで、基本的にはトップレベルにも配置できる内容です。
また、アプリケーション側のコードがCoroutineで書かれている場合にも、これらのコードは機能します。

読み出し処理に関する注意点

この読み出し処理はReactor Contextに対して行われるため、Reactor Contextが異なっていたり、Reactor Contextにアクセスできない場所から呼び出すと正常に機能しません。

前者は、例えば共通処理を別Coroutine Scopeで実行したような場合に問題となります。
これに関してはReactor Contextの値を引き継ぐことで対処できます。

qiita.com

後者はReactorを利用していない・Reactive Streamsのみ利用している外部ライブラリに連なる処理で実行されたような場合に問題となります。
自分が知っている限りではjOOQがこれに当たります。
詳しくは以下の記事にまとめています。

qiita.com

3. 2で取得したIdentifierをコネクションにセットする

Identifierをコネクションにセットする際は、ConnectionPoolのpostAllocateを利用するのが楽だと思います。
最低限の設定だけを抜き出したR2dbcConfigurationは以下のようになります。

import io.r2dbc.pool.ConnectionPool
import io.r2dbc.pool.ConnectionPoolConfiguration
import io.r2dbc.spi.ConnectionFactories
import io.r2dbc.spi.ConnectionFactory
import org.springframework.boot.autoconfigure.r2dbc.R2dbcProperties
import org.springframework.boot.context.properties.ConfigurationProperties
import org.springframework.context.annotation.Bean
import org.springframework.context.annotation.Configuration
import org.springframework.data.r2dbc.config.AbstractR2dbcConfiguration
import reactor.kotlin.core.publisher.toMono

@Configuration
class R2dbcConfiguration(private val identifierGetter: IdentifierGetter) : AbstractR2dbcConfiguration() {
    @ConfigurationProperties("spring.r2dbc")
    @Bean
    fun r2dbcProperties(): R2dbcProperties = R2dbcProperties()

    @Bean
    override fun connectionFactory(): ConnectionFactory {
        // URL にユーザー名とパスワードが含まれてないと接続ができない
        val defaultConnectionFactory = ConnectionFactories.get(r2dbcProperties().url)

        val config = ConnectionPoolConfiguration.builder(defaultConnectionFactory)
            .postAllocate { connection ->
                identifierGetter.getIdentifierMono()
                    .flatMap {
                        connection.createStatement("SET app.current_identifier = $1")
                            .bind(0, it)
                            .execute()
                            .toMono()
                    }
                    .then()
            }
            .build()

        return ConnectionPool(config)
    }
}

application.propertiesのspring.r2dbcの下にはurlプロパティのみ設定している想定です。

spring.r2dbc.url=r2dbc:postgresql://${ユーザー名}:${パスワード}@localhost/${DB名}

この設定を行なったConnectionPool（ConnectionFactory）から発行されるConnectionは、アプリケーション内では常に適切なIdentifierが設定された状態になります。

終わりに

この記事では、Spring WebFlux x R2DBCのマルチモジュールプロジェクトでのRLS基盤作成についてまとめました。

弊チームが開発を行なっていた時点では、Spring WebFlux x R2DBCでRow Level Securityを有効にする例は世界的に見ても少なく、実装時点ではサンプルコードも中々見つけられない状況でした。
記事ではかなりあっさりした書き方をしていますが、開発を進める中では、知見不足で調査に時間がかかったり、jOOQ固有の問題でハマったりと、多くの苦労が有りました。

同じような構成での開発を行う方へ、この記事が何かのお役に立てば幸いです。

justInCaseでは、さまざまな職種を採用中です。
保険SaaSの開発にご興味がある方は採用ページよりお気軽にご応募ください！
Spring WebFluxを用いたバックエンド開発をやってみたい方大歓迎です！

justincase.jp

justInCaseTechnologiesのテックブログ

Spring WebFlux・R2DBCのマルチモジュールプロジェクトでRow Level Security基盤を構築した話

前書き

技術スタック

Row Level Securityについて

アプリケーション側の実装について

基本方針

基本的な処理の流れ

1. スコープ内で利用する`Identifier`を`Reactor Context`にセットする

セットする内容について

Reactor Contextへのアクセスに関する注意点

認証等の情報から読み出す場合の注意点

2. 1でセットしたIdentifierをReactor Contextから読み出す

読み出し処理に関する注意点

3. 2で取得したIdentifierをコネクションにセットする

終わりに

前書き

技術スタック

Row Level Securityについて

アプリケーション側の実装について

基本方針

基本的な処理の流れ

1. スコープ内で利用するIdentifierをReactor Contextにセットする

セットする内容について

Reactor Contextへのアクセスに関する注意点

認証等の情報から読み出す場合の注意点

2. 1でセットしたIdentifierをReactor Contextから読み出す

読み出し処理に関する注意点

3. 2で取得したIdentifierをコネクションにセットする

終わりに

1. スコープ内で利用する`Identifier`を`Reactor Context`にセットする