Spring WebFlux2.4.3 - WebSocketHandlerにPrincipalを受け渡す

要約

  • 現時点で、WebFluxでWebSocketをハンドリングするWebSocketHandlerで受け取っているWebSocketSessionHandshakeInfo保有する principleMonoは、Mono.empty()を返している。
  • つまり、そのままだとWebSocketHandlerでは、ログイン中のユーザー情報(Principal)を受け取ることができない。
  • 解決法として、デフォルトで利用されるWebSocketService実装を、ProtocolのUpgrade処理時にSessionをPrincipalに解決できるようにする実装に置き換える。

問題

そのため、以下のようにWebSocketSessionPrincipalの情報を取得しようとしても、完全に処理が停止してしまい、ユーザー情報とFluxを結合させることができない。

public class WebSocketEndpointHandler implements WebSocketHandler {

// ...

    @Override
    public @NonNull Mono<Void> handle(WebSocketSession session) {
        final var fluxWithPrincipal = session.getHandshakeInfo()
                .getPrincipal()
                .flatMap(principal -> session
                        .receive()
                        .map(webSocketMessage ->
                                // 各メッセージとPricipleの結合
                                ));
    }

// ...

}

解決策

解決策としては、ReactorNettyRequestUpgradeStrategy をベースに、オリジナルの RequestUpgradeStrategy を実装し、これを利用するようにWebFluxConfigurerを設定する。

ます、ReactorNettyRequestUpgradeStrategyをオーバーライドして、以下のようにServerSecurityContextRepositoryからPrincipalを解決するようにし、新たなHandshakeInfoを返すようにしたものを用意する。

@RequiredArgsConstructor
public class ReactorNettyRequestUpgradeStrategyWithPrincipal extends ReactorNettyRequestUpgradeStrategy {

    @NonNull
    private final ServerSecurityContextRepository serverSecurityContextRepository;

    @Override
    public Mono<Void> upgrade(ServerWebExchange exchange, WebSocketHandler handler,
                              @Nullable String subProtocol, Supplier<HandshakeInfo> handshakeInfoFactory) {
        return super.upgrade(exchange, handler, subProtocol, () -> {
            final var handShakeInfo = handshakeInfoFactory.get();
            final var gamePlayerPrincipal = serverSecurityContextRepository
                    .load(exchange)
                    .map(SecurityContext::getAuthentication)
                    .map(Authentication::getPrincipal)
                    .cast(GamePlayer.class)
                    .map(player -> (Principal) player::getPlayerId);

            return new HandshakeInfo(
                    handShakeInfo.getUri(),
                    handShakeInfo.getHeaders(),
                    gamePlayerPrincipal,
                    handShakeInfo.getSubProtocol(),
                    handShakeInfo.getRemoteAddress(),
                    handShakeInfo.getAttributes(),
                    handShakeInfo.getLogPrefix()
            );
        });
    }

}

そして、自作の ReactorNettyRequestUpgradeStrategy を利用するように WebFluxConfigurer を設定する。

@Configuration
@EnableWebFlux
public class WebFluxConfig implements WebFluxConfigurer {

// ...

    @Autowired
    private ServerSecurityContextRepository serverSecurityContextRepository;

    @Override
    public WebSocketService getWebSocketService() {
        return new HandshakeWebSocketService(
                new ReactorNettyRequestUpgradeStrategyWithPrincipal(
                        serverSecurityContextRepository));
    }

// ...

}

これにより、WebFlux WebSocketからでもSpring Securityによって解決されたユーザー情報を利用することができるようになる。

参考文献