How can I write Security tests for Spring Boot WebSocket endpoints

I'm looking for a sensible way to unit test our WebSocket security implementation. We use Spring TestRestTemplate to test our REST endpoints, it seems this approach almost works with WebSockets. If you accept BAD_REQUEST as a substitue for OK.

Is there a good way to test Spring WebSockets?

This is our handler mapping

@Configuration
class HandlerMappingConfiguration {
    @Bean
    fun webSocketMapping(webSocketHandler: DefaultWebSocketHandler?): HandlerMapping {
        val map: MutableMap<String, WebSocketHandler?> = HashMap()
        map["/"] = webSocketHandler
        map["/protected-ws"] = webSocketHandler

        val handlerMapping = SimpleUrlHandlerMapping()
        handlerMapping.order = 1
        handlerMapping.urlMap = map
        return handlerMapping
    }
}

This is our security configuration

@EnableWebFluxSecurity
@EnableReactiveMethodSecurity
class SecurityConfig(
    private val authenticationConverter: ServerAuthenticationConverter,
    @param:Value("${spring.security.oauth2.resourceserver.jwt.issuer-uri}") private val issuerUri: String,
    @param:Value("${spring.security.oauth2.resourceserver.jwt.audience}") private val audience: String
) {
    @Bean
    fun springWebFilterChain(http: ServerHttpSecurity): SecurityWebFilterChain {
        http
            .csrf().disable()
            .authorizeExchange()
            .pathMatchers("/protected-ws").hasAuthority("SCOPE_read:client-ws")
            .pathMatchers("/protected").hasAuthority("SCOPE_read:client")
            .matchers(EndpointRequest.toAnyEndpoint()).permitAll()
            .anyExchange().authenticated()
            .and()
            .oauth2ResourceServer()
            .bearerTokenConverter(authenticationConverter)
            .jwt()
        return http.build()
    }

    @Bean
    open fun jwtDecoder(): ReactiveJwtDecoder {
        val jwtDecoder = ReactiveJwtDecoders.fromOidcIssuerLocation(issuerUri) as NimbusReactiveJwtDecoder
        jwtDecoder.setJwtValidator(DelegatingOAuth2TokenValidator(withIssuer(), audienceValidator()))
        return jwtDecoder
    }

    private fun withIssuer(): OAuth2TokenValidator<Jwt>? {
        return JwtValidators.createDefaultWithIssuer(
            issuerUri
        )
    }

    private fun audienceValidator(): OAuth2TokenValidator<Jwt> {
        return JwtAudienceValidator(
            audience
        )
    }
}

This is our IntegrationTest class for REST endpoints

@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
class RestSecurityTest(@Autowired val restTemplate: TestRestTemplate) {

    @Test
    fun `Unauthorized user is unauthorized`() {
        val entity = restTemplate.getForEntity<String>("/world", String::class.java)
        assertThat(entity.statusCode).isEqualTo(HttpStatus.UNAUTHORIZED)
    }

    @Test
    @WithMockUser(authorities = ["SCOPE_read:world"])
    fun `Authorized user, without required scope is Forbidden`() {
        val entity = restTemplate.getForEntity<String>("/protected", String::class.java)
        assertThat(entity.statusCode).isEqualTo(HttpStatus.FORBIDDEN)
    }

    @Test
    @WithMockUser(authorities = ["SCOPE_read:client"])
    fun `Authorized user, with required scope is OK`() {
        val entity = restTemplate.getForEntity<String>("/protected", String::class.java)
        assertThat(entity.statusCode).isEqualTo(HttpStatus.OK)
    }
}

We would like to achieve something similar for our WebSocket endpoints... for example

@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
class WebSocketSecurityTest(@Autowired val restTemplate: TestRestTemplate) {

    @Test
    fun `WebSocket unauthorized user is unauthorized`() {
        val entity = restTemplate.getForEntity<String>("/protected-ws", String::class.java)
        assertThat(entity.statusCode).isEqualTo(HttpStatus.UNAUTHORIZED)
    }

    @Test
    @WithMockUser(authorities = ["SCOPE_read:world"])
    fun `WebSocket authorized user, without required scope is forbidden`() {
        val entity = restTemplate.getForEntity<String>("/protected-ws", String::class.java)
        assertThat(entity.statusCode).isEqualTo(HttpStatus.FORBIDDEN)
    }

    @Test
    @WithMockUser(authorities = ["SCOPE_read:client-ws"])
    fun `WebSocket Authorized user, with required scope is OK`() {
        val entity = restTemplate.getForEntity<String>("/protected-ws", String::class.java)
        // You can't open a WebSocket connection with a REST client. But this shows the request is authorized.
        assertThat(entity.statusCode).isEqualTo(HttpStatus.BAD_REQUEST)
    }
}

But this has several limitations:

1. Our OK test has to check for a BAD_REQUEST response!
2. We are unable to verify any actual payloads from the endpoint.

So it's only suitable for failure scenarios. Full source code for this example available on github.

I found an approach based on what I read on STOMP testing, I realised the STOMP client was using a plain WebSocket client and I found I was able to build tests the same way without the STOMP wrapper.

@Test
@WithMockUser(authorities = ["SCOPE_read:client-ws"])
fun `WebSocket authorized user is able to connect to endpoint`() {
    val latch = CountDownLatch(1)
    StandardWebSocketClient().execute(
        URI.create(String.format("ws://localhost:%d/protected-ws", port)),
        TestClientHandler(latch)
    ).subscribe()
    assertThat(latch.await(60, TimeUnit.SECONDS)).isTrue()
}