feat: ws token

backend/src/main/kotlin/com/github/sor2171/backend/config/SecurityConfig.kt

@@ -67,14 +67,14 @@ class SecurityConfig(
                     writeJsonResponse(
                         exchange,
                         HttpStatus.UNAUTHORIZED,
-                        ApiResponse.unauthenticated(ex.message)
+                        ApiResponse.unauthenticated<Any>(ex.message)
                     )
                 }
                     .accessDeniedHandler { exchange, ex ->
                         writeJsonResponse(
                             exchange,
                             HttpStatus.FORBIDDEN,
-                            ApiResponse.forbidden(ex.message)
+                            ApiResponse.forbidden<Any>(ex.message)
                         )
                     }
             }
@@ -112,7 +112,7 @@ class SecurityConfig(
             writeJsonResponse(
                 webFilterExchange.exchange,
                 HttpStatus.UNAUTHORIZED,
-                ApiResponse.unauthenticated(exception.message)
+                ApiResponse.unauthenticated<Any>(exception.message)
             )
         }
         return filter

backend/src/main/kotlin/com/github/sor2171/backend/config/WebSocketConfig.kt

@@ -13,6 +13,7 @@ class WebSocketConfig(
     private val exampleHandler: ExampleHandler,
     private val testHandler: TestHandler
 ) {
+
     @Bean
     fun handlerMapping(): HandlerMapping {
         val map = mapOf(

backend/src/main/kotlin/com/github/sor2171/backend/controller/TestController.kt

@@ -1,5 +1,8 @@
 package com.github.sor2171.backend.controller
 
+import com.github.sor2171.backend.entity.ApiResponse
+import com.github.sor2171.backend.handler.v1.ExampleHandler
+import com.github.sor2171.backend.service.WsTokenService
 import org.springframework.web.bind.annotation.GetMapping
 import org.springframework.web.bind.annotation.RequestMapping
 import org.springframework.web.bind.annotation.RestController
@@ -7,8 +10,19 @@ import reactor.core.publisher.Mono
 
 @RestController
 @RequestMapping("/api/test")
-class TestController {
-
+class TestController(
+    private val wsTokenService: WsTokenService
+) {
     @GetMapping("/hello")
     fun test(): Mono<String> = Mono.just("Hello Kotlin WebFlux!")
+
+    @GetMapping("/wstoken")
+    fun testToken(): Mono<ApiResponse<String>> =
+        wsTokenService.registerToken(ExampleHandler::class)
+            .map { token ->
+                ApiResponse.success(data = token)
+            }
+            .onErrorResume { e ->
+                Mono.just(ApiResponse.innerError(e.message))
+            }
 }

backend/src/main/kotlin/com/github/sor2171/backend/entity/ApiResponse.kt

@@ -6,7 +6,7 @@ import kotlinx.serialization.json.Json
 @Serializable
 data class ApiResponse<T>(
     val code: Int,
-    val data: T,
+    val data: T?,
     val message: String
 ) {
     companion object {
@@ -15,21 +15,24 @@ data class ApiResponse<T>(
         }
 
         // 针对无 data 返回的情况，统一使用 Any?
-        fun success(message: String = "Success"): ApiResponse<Any?> {
+        fun <T> success(message: String = "Success"): ApiResponse<T> {
             return ApiResponse(200, null, message)
         }
 
-        fun <T> failure(code: Int = 401, data: T, message: String?): ApiResponse<T> {
+        fun <T> failure(code: Int, data: T?, message: String?): ApiResponse<T> {
             return ApiResponse(code, data, message ?: "Failure")
         }
 
-        fun unauthenticated(message: String?) =
-            failure(401, null, message)
+        fun <T> unauthenticated(message: String?) =
+            failure<T>(401, null, message)
 
-        fun forbidden(message: String?) =
-            failure(403, null, message)
+        fun <T> forbidden(message: String?) =
+            failure<T>(403, null, message)
 
-        fun logoutFailed(message: String = ""): ApiResponse<Any?> {
+        fun <T> innerError(message: String?) =
+            failure<T>(500, null, message)
+
+        fun <T> logoutFailed(message: String = ""): ApiResponse<T> {
             val formatMessage = if (message.isBlank()) "" else ": $message"
             return failure(400, null, "Logout Failed$formatMessage")
         }

backend/src/main/kotlin/com/github/sor2171/backend/filter/FlowLimitFilter.kt

@@ -45,7 +45,7 @@ class FlowLimitFilter(
         response.statusCode = HttpStatus.FORBIDDEN
         response.headers.contentType = MediaType.APPLICATION_JSON
         val buffer = response.bufferFactory()
-            .wrap(ApiResponse.forbidden("Too many requests").toJsonString().toByteArray())
+            .wrap(ApiResponse.forbidden<Any>("Too many requests").toJsonString().toByteArray())
         return response.writeWith(Mono.just(buffer))
     }
 

backend/src/main/kotlin/com/github/sor2171/backend/handler/v1/ExampleHandler.kt

@@ -1,20 +1,31 @@
 package com.github.sor2171.backend.handler.v1
 
+import com.github.sor2171.backend.service.WsTokenService
 import org.springframework.stereotype.Component
 import org.springframework.web.reactive.socket.WebSocketHandler
 import org.springframework.web.reactive.socket.WebSocketSession
 import reactor.core.publisher.Mono
 
 @Component
-class ExampleHandler : WebSocketHandler {
+class ExampleHandler(
+    private val wsTokenService: WsTokenService,
+) : WebSocketHandler {
     override fun handle(session: WebSocketSession): Mono<Void> {
-        // ws://localhost:8080/ws/v1/example
+        // ws://localhost:8080/ws/v1/example?token=
         val output = session.receive()
             .map { msg ->
                 val userText = msg.payloadAsText
                 session.textMessage("Kotlin springboot4 received: $userText")
             }
-        
-        return session.send(output)
+
+        return wsTokenService
+            .validateToken(session, this::class)
+            .flatMap { valid ->
+                if (!valid) {
+                    session.close()
+                } else {
+                    session.send(output)
+                }
+            }
     }
 }

backend/src/main/kotlin/com/github/sor2171/backend/service/WsTokenService.kt

@@ -0,0 +1,55 @@
+package com.github.sor2171.backend.service
+
+import com.github.sor2171.backend.utils.Const.WS_REQUEST_TOKEN
+import org.springframework.beans.factory.annotation.Value
+import org.springframework.data.redis.core.ReactiveStringRedisTemplate
+import org.springframework.stereotype.Service
+import org.springframework.web.reactive.socket.WebSocketHandler
+import org.springframework.web.reactive.socket.WebSocketSession
+import org.springframework.web.util.UriComponentsBuilder
+import reactor.core.publisher.Mono
+import java.time.Duration
+import java.util.*
+import kotlin.reflect.KClass
+
+@Service
+class WsTokenService(
+    private val template: ReactiveStringRedisTemplate,
+
+    @param:Value($$"${app.ws.token-expire-time}")
+    private val defaultTokenExpireTime: Long,
+) {
+    fun registerToken(
+        wsType: KClass<out WebSocketHandler>,
+        expireTime: Long? = null,
+    ): Mono<String> {
+        val token = UUID.randomUUID().toString()
+        val key = "$WS_REQUEST_TOKEN:${wsType.simpleName}:" + UUID.randomUUID().toString()
+        return template.opsForValue().set(
+            key,
+            "1",
+            Duration.ofSeconds(expireTime ?: defaultTokenExpireTime)
+        ).flatMap { success ->
+            if (success) {
+                Mono.just(token)
+            } else {
+                Mono.error(IllegalStateException("Failed to register token"))
+            }
+        }
+    }
+
+    fun validateToken(
+        session: WebSocketSession,
+        wsType: KClass<out WebSocketHandler>,
+    ): Mono<Boolean> {
+        val token = UriComponentsBuilder
+            .fromUri(session.handshakeInfo.uri)
+            .build()
+            .queryParams
+            .getFirst("token")
+            ?: return Mono.just(false)
+
+        val key = "$WS_REQUEST_TOKEN:${wsType.simpleName}:" + token
+        return template.delete(key).map { it > 0 }
+    }
+}

backend/src/main/kotlin/com/github/sor2171/backend/utils/Const.kt

@@ -11,6 +11,8 @@ object Const {
 
     const val FLOW_LIMIT_COUNTER = "flow:counter:"
     const val FLOW_LIMIT_BLOCK = "flow:block:"
+    
+    const val WS_REQUEST_TOKEN = "ws:token:"
 
     val INTERNAL_ERROR_MONO = Mono.just("something went wrong. Please contact the administrator.")
 }

backend/src/main/resources/application-example.yaml

@@ -41,4 +41,7 @@ app:
   flow-limit:
     count: 1000
     limit-time: 3
-    ban-time: 60
+    ban-time: 60
+
+  ws:
+    token-expire-time: 300