KtorでAtProtocolのRefreshSessionをしたい

はじめに

Blueskyのデスクトップアプリ for Macを Compose for Desktop でちまちまと作っている。

通信は Ktor + Ktorfit というKMPのデファクトなライブラリを使っている。

Blueskyのライムライン取得や投稿などはHTTPリクエストを送ることになるが、それら通信の仕様やBlueskyが実装している分散型SNSの仕様は AT Protocolとして定義されている。

atproto.com

通信の流れとしてはよくある感じで、

  1. createSessionでセッションを生成する
  2. 以降はここで得られたトークンをHeaderに入れて通信する（Bearer Authentication）
  3. セッションの有効期限が切れたら refreshSession でリフレッシュ

となる。

Ktorでは各リクエストにトークンを渡す設定は簡単に行える。

install(Auth) {
    bearer {
        loadTokens {
            val accessToken = dataSource.getAccessToken()
            val refreshToken = dataSource.getRefreshToken()
            BearerTokens(accessToken, refreshToken)
        }
    }
}

こんな実装をしておくとヘッダーにAuthorization: Bearer ${your_access_token}が付与される。

ktor.io

また、リフレッシュトークンの仕組みも備わっており、こちらもかなり簡単に実装ができるようになっている。

install(Auth) {
    bearer {
        refreshTokens {
            // トークンをrefreshするようなエンドポイントを叩くなりして新しいトークンを得る
            BearerTokens(newAccessToken, newRefresshToken)
        }
    }
}

ktor.io

Ktorのリフレッシュトークン機構の発動条件

このリフレッシュトークンの仕組みが発動するのは、

  • refreshTokensでトークン更新の設定をしており、
  • リクエストの結果がstatusCode=401のとき

に限られる。

https://github.com/ktorio/ktor/blob/6c391d04548626d77a9a6e9523f9cdeffa6aa5ef/ktor-client/ktor-client-plugins/ktor-client-auth/common/src/io/ktor/client/plugins/auth/Auth.kt#L138-L151

多くのサービスでは問題ないと思うが、トークン有効期限切れのときに401を返さないサービスも存在する。

Blueskyもそうだった。

Blueskyの場合、有効期限が切れているときには次のようなエラーレスポンスが返ってくる。

400 Bad Request
{"error":"ExpiredToken","message":"Token has expired"}

これだとKtorのリフレッシュトークンの仕組みは発動しないため、代替案を考える必要があった。

案1 発動条件を変更する（没）

Ktorは柔軟に設計されている印象を受けていたので、リフレッシュトークンに関してもその発動条件を変更することができるだろうと思った。

ドキュメントと実装をみたが、どうやらできそうになかった。

できても良いような気がしているので、もしできるなら方法を知りたい......。

案2 エラー情報を改変して返す（没）

カスタムプラグイン（後述）を使って「status=400かつerror=ExpiredTokenの場合には401エラーが返ってきた」ことにするよう、偽装できないかと考えた。

結果、こちらも無理そうだった。カスタムプラグインでは返ってきたレスポンスの内容を編集するような操作はできなかった。

案3 カスタムプラグインで実装する（採用）

Ktorにはカスタムプラグインという仕様が存在する。OkHttpのInterceptorのように、リクエストやレスポンス（またはその両方）の前後に任意の処理を挟むことができる。

ktor.io

  1. 通信がstatusCode=400で、かつエラーレスポンスのerrorがExpiredTokenだった場合、
  2. refreshSessionを実行し、
  3. 新たに取得したトークンで400エラーとなった最初の通信をリトライする

というような実装をする。

例としてはこんな感じ。

class RefreshTokenPluginFactory {

    fun create() = createClientPlugin("RefreshTokenPlugin") {
        on(Send) { request ->
            // 目的のリクエストを実行
            val originalCall = proceed(request)

            // 400エラーで失敗し 
            if (originalCall.response.status == HttpStatusCode.BadRequest) {
                val errorResponse = originalCall.response.body<ErrorResponse>()
                
                // エラーが`ExpiredToken`だったとき
                if (errorResponse.error == "ExpiredToken") {
                    
                    // refreshSessionをリクエスト
                    val newSession = originalCall.client.request {
                        method = HttpMethod.Post
                        
                        // see https://github.com/bluesky-social/atproto/blob/main/lexicons/com/atproto/server/refreshSession.json
                        url("https://bsky.social/xrpc/com.atproto.server.refreshSession")
                        
                        // ここはrefreshTokenを使う
                        // see https://atproto.com/specs/xrpc#authentication
                        header("Authorization", "Bearer ${refreshToken}")
                    }.body<Session>
                    
                    // 新しく取得したトークンをもとのリクエストヘッダーに追加
                    request.headers["Authorization"] = "Bearer ${newSession.accessToken}"

                    // 再度リクエスト
                    proceed(request)
                } else {
                    originalCall
                }
            } else {
                originalCall
            }
        }
    }
}

これをinstallすれば良い。

install(refreshTokenPluginFactory.create())

ただ、このままだと色々と問題がある。

RefreshSessionがループしないようにする

上記のサンプル実装ではPOST /refreshSessionをリクエストする場合にもプラグイン処理が呼ばれる。

そのため、POST /refreshSessionがExpiredTokenエラーになるとPOST /refreshSessionをリクエストして...と無限ループに陥る可能性がある。

従ってガードが必要。

if (request.url.toString().contains("refreshSession")) {
    return@on process(request)
}

RefreshTokenが利用されない

Blueskyの場合（というか他のサービスでもそうかも知れないが）、POST /refreshSessionを行う場合にはrefreshTokenを利用する必要がある。

しかし最初に例示したloadTokens{}を使っている場合、上記のようにプラグイン内でrefreshTokenを設定していたとしても、accessTokenで上書きされてしまう。

トークン設定はプラグインの中で行う実装をし、loadTokensは使わないことでこの問題を回避できる。

if (request.url.toString().contains("refreshSession")) {
    return@on process(request)
}

request.headers["Authorization"] = accessToken
val originalCall = process(request)

これは削除する。

install(Auth) {
    bearer {
        loadTokens {
            val accessToken = dataSource.getAccessToken()
            val refreshToken = dataSource.getRefreshToken()
            BearerTokens(accessToken, refreshToken)
        }
    }
}

全体的に実装はもうちょっと整理したほうがいい気がする。

例えば「ヘッダーに情報を詰める処理」と「リフレッシュトークンの処理」はそれぞれ別のプラグインとして実装するのもアリかもしれない。

そもそも、もっと良い方法がありそうな気もするのだけれど。

以上。