如何為 Spring Boot WebSocket 端點編寫安全測試

Question

我正在尋找一種合理的方法來對我們的 WebSocket 安全實現進行單元測試。 我們使用 Spring TestRestTemplate來測試我們的 REST 端點，似乎這種方法幾乎適用於 WebSockets。 如果您接受BAD_REQUEST作為OK的替代品。

有沒有好的方法可以測試 Spring WebSockets？

這是我們的處理程序映射

@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
    }
}

這是我們的安全配置

@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
        )
    }
}

這是我們針對 REST 端點的IntegrationTest class

@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)
    }
}

我們想為我們的 WebSocket 端點實現類似的東西......例如

@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)
    }
}

但這有幾個限制：

1. 我們的OK測試必須檢查BAD_REQUEST響應！
2. 我們無法驗證端點的任何實際有效負載。

所以它只適用於失敗的場景。 此示例的完整源代碼可在 github 上獲得。

Answer 1

我找到了一種基於我在 STOMP 測試中閱讀的內容的方法，我意識到 STOMP 客戶端使用的是普通的 WebSocket 客戶端，我發現我能夠在沒有 STOMP 包裝器的情況下以相同的方式構建測試。

@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()
}