feat: ws token

2026-10-23 19:41:12 +08:00 · 2026-07-06 15:58:53 +08:00
parent 7c3f8b0af5
commit 3df38b5757
13 changed files with 441 additions and 20 deletions
@@ -1,7 +1,9 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-9.4.0-bin.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-9.6.0-bin.zip
 networkTimeout=10000
+retries=0
+retryBackOffMs=500
 validateDistributionUrl=true
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists
@@ -0,0 +1,248 @@
+#!/bin/sh
+
+#
+# Copyright © 2015 the original authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+
+##############################################################################
+#
+#   gradlew start up script for POSIX generated by Gradle.
+#
+#   Important for running:
+#
+#   (1) You need a POSIX-compliant shell to run this script. If your /bin/sh is
+#       noncompliant, but you have some other compliant shell such as ksh or
+#       bash, then to run this script, type that shell name before the whole
+#       command line, like:
+#
+#           ksh gradlew
+#
+#       Busybox and similar reduced shells will NOT work, because this script
+#       requires all of these POSIX shell features:
+#         * functions;
+#         * expansions «$var», «${var}», «${var:-default}», «${var+SET}»,
+#           «${var#prefix}», «${var%suffix}», and «$( cmd )»;
+#         * compound commands having a testable exit status, especially «case»;
+#         * various built-in commands including «command», «set», and «ulimit».
+#
+#   Important for patching:
+#
+#   (2) This script targets any POSIX shell, so it avoids extensions provided
+#       by Bash, Ksh, etc; in particular arrays are avoided.
+#
+#       The "traditional" practice of packing multiple parameters into a
+#       space-separated string is a well documented source of bugs and security
+#       problems, so this is (mostly) avoided, by progressively accumulating
+#       options in "$@", and eventually passing that to Java.
+#
+#       Where the inherited environment variables (DEFAULT_JVM_OPTS, JAVA_OPTS,
+#       and GRADLE_OPTS) rely on word-splitting, this is performed explicitly;
+#       see the in-line comments for details.
+#
+#       There are tweaks for specific operating systems such as AIX, CygWin,
+#       Darwin, MinGW, and NonStop.
+#
+#   (3) This script is generated from the Groovy template
+#       https://github.com/gradle/gradle/blob/3d91ce3b8caaf77ad09f381f43615b715b53f72c/platforms/jvm/plugins-application/src/main/resources/org/gradle/api/internal/plugins/unixStartScript.txt
+#       within the Gradle project.
+#
+#       You can find Gradle at https://github.com/gradle/gradle/.
+#
+##############################################################################
+
+# Attempt to set APP_HOME
+
+# Resolve links: $0 may be a link
+app_path=$0
+
+# Need this for daisy-chained symlinks.
+while
+    APP_HOME=${app_path%"${app_path##*/}"}  # leaves a trailing /; empty if no leading path
+    [ -h "$app_path" ]
+do
+    ls=$( ls -ld "$app_path" )
+    link=${ls#*' -> '}
+    case $link in             #(
+      /*)   app_path=$link ;; #(
+      *)    app_path=$APP_HOME$link ;;
+    esac
+done
+
+# This is normally unused
+# shellcheck disable=SC2034
+APP_BASE_NAME=${0##*/}
+# Discard cd standard output in case $CDPATH is set (https://github.com/gradle/gradle/issues/25036)
+APP_HOME=$( cd -P "${APP_HOME:-./}" > /dev/null && printf '%s\n' "$PWD" ) || exit
+
+# Use the maximum available, or set MAX_FD != -1 to use that value.
+MAX_FD=maximum
+
+warn () {
+    echo "$*"
+} >&2
+
+die () {
+    echo
+    echo "$*"
+    echo
+    exit 1
+} >&2
+
+# OS specific support (must be 'true' or 'false').
+cygwin=false
+msys=false
+darwin=false
+nonstop=false
+case "$( uname )" in                #(
+  CYGWIN* )         cygwin=true  ;; #(
+  Darwin* )         darwin=true  ;; #(
+  MSYS* | MINGW* )  msys=true    ;; #(
+  NONSTOP* )        nonstop=true ;;
+esac
+
+
+
+# Determine the Java command to use to start the JVM.
+if [ -n "$JAVA_HOME" ] ; then
+    if [ -x "$JAVA_HOME/jre/sh/java" ] ; then
+        # IBM's JDK on AIX uses strange locations for the executables
+        JAVACMD=$JAVA_HOME/jre/sh/java
+    else
+        JAVACMD=$JAVA_HOME/bin/java
+    fi
+    if [ ! -x "$JAVACMD" ] ; then
+        die "ERROR: JAVA_HOME is set to an invalid directory: $JAVA_HOME
+
+Please set the JAVA_HOME variable in your environment to match the
+location of your Java installation."
+    fi
+else
+    JAVACMD=java
+    if ! command -v java >/dev/null 2>&1
+    then
+        die "ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.
+
+Please set the JAVA_HOME variable in your environment to match the
+location of your Java installation."
+    fi
+fi
+
+# Increase the maximum file descriptors if we can.
+if ! "$cygwin" && ! "$darwin" && ! "$nonstop" ; then
+    case $MAX_FD in #(
+      max*)
+        # In POSIX sh, ulimit -H is undefined. That's why the result is checked to see if it worked.
+        # shellcheck disable=SC2039,SC3045
+        MAX_FD=$( ulimit -H -n ) ||
+            warn "Could not query maximum file descriptor limit"
+    esac
+    case $MAX_FD in  #(
+      '' | soft) :;; #(
+      *)
+        # In POSIX sh, ulimit -n is undefined. That's why the result is checked to see if it worked.
+        # shellcheck disable=SC2039,SC3045
+        ulimit -n "$MAX_FD" ||
+            warn "Could not set maximum file descriptor limit to $MAX_FD"
+    esac
+fi
+
+# Collect all arguments for the java command, stacking in reverse order:
+#   * args from the command line
+#   * the main class name
+#   * -classpath
+#   * -D...appname settings
+#   * --module-path (only if needed)
+#   * DEFAULT_JVM_OPTS, JAVA_OPTS, and GRADLE_OPTS environment variables.
+
+# For Cygwin or MSYS, switch paths to Windows format before running java
+if "$cygwin" || "$msys" ; then
+    APP_HOME=$( cygpath --path --mixed "$APP_HOME" )
+
+    JAVACMD=$( cygpath --unix "$JAVACMD" )
+
+    # Now convert the arguments - kludge to limit ourselves to /bin/sh
+    for arg do
+        if
+            case $arg in                                #(
+              -*)   false ;;                            # don't mess with options #(
+              /?*)  t=${arg#/} t=/${t%%/*}              # looks like a POSIX filepath
+                    [ -e "$t" ] ;;                      #(
+              *)    false ;;
+            esac
+        then
+            arg=$( cygpath --path --ignore --mixed "$arg" )
+        fi
+        # Roll the args list around exactly as many times as the number of
+        # args, so each arg winds up back in the position where it started, but
+        # possibly modified.
+        #
+        # NB: a `for` loop captures its iteration list before it begins, so
+        # changing the positional parameters here affects neither the number of
+        # iterations, nor the values presented in `arg`.
+        shift                   # remove old arg
+        set -- "$@" "$arg"      # push replacement arg
+    done
+fi
+
+
+# Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
+DEFAULT_JVM_OPTS='"-Xmx64m" "-Xms64m"'
+
+# Collect all arguments for the java command:
+#   * DEFAULT_JVM_OPTS, JAVA_OPTS, and optsEnvironmentVar are not allowed to contain shell fragments,
+#     and any embedded shellness will be escaped.
+#   * For example: A user cannot expect ${Hostname} to be expanded, as it is an environment variable and will be
+#     treated as '${Hostname}' itself on the command line.
+
+set -- \
+        "-Dorg.gradle.appname=$APP_BASE_NAME" \
+        -jar "$APP_HOME/gradle/wrapper/gradle-wrapper.jar" \
+        "$@"
+
+# Stop when "xargs" is not available.
+if ! command -v xargs >/dev/null 2>&1
+then
+    die "xargs is not available"
+fi
+
+# Use "xargs" to parse quoted args.
+#
+# With -n1 it outputs one arg per line, with the quotes and backslashes removed.
+#
+# In Bash we could simply go:
+#
+#   readarray ARGS < <( xargs -n1 <<<"$var" ) &&
+#   set -- "${ARGS[@]}" "$@"
+#
+# but POSIX shell has neither arrays nor command substitution, so instead we
+# post-process each arg (as a line of input to sed) to backslash-escape any
+# character that might be a shell metacharacter, then use eval to reverse
+# that process (while maintaining the separation between arguments), and wrap
+# the whole thing up as a single "set" statement.
+#
+# This will of course break if any of these variables contains a newline or
+# an unmatched quote.
+#
+
+eval "set -- $(
+        printf '%s\n' "$DEFAULT_JVM_OPTS $JAVA_OPTS $GRADLE_OPTS" |
+        xargs -n1 |
+        sed ' s~[^-[:alnum:]+,./:=@_]~\\&~g; ' |
+        tr '\n' ' '
+    )" '"$@"'
+
+exec "$JAVACMD" "$@"
@@ -0,0 +1,82 @@
+@rem
+@rem Copyright 2015 the original author or authors.
+@rem
+@rem Licensed under the Apache License, Version 2.0 (the "License");
+@rem you may not use this file except in compliance with the License.
+@rem You may obtain a copy of the License at
+@rem
+@rem      https://www.apache.org/licenses/LICENSE-2.0
+@rem
+@rem Unless required by applicable law or agreed to in writing, software
+@rem distributed under the License is distributed on an "AS IS" BASIS,
+@rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+@rem See the License for the specific language governing permissions and
+@rem limitations under the License.
+@rem
+@rem SPDX-License-Identifier: Apache-2.0
+@rem
+
+@if "%DEBUG%"=="" @echo off
+@rem ##########################################################################
+@rem
+@rem  gradlew startup script for Windows
+@rem
+@rem ##########################################################################
+
+@rem Set local scope for the variables, and ensure extensions are enabled
+setlocal EnableExtensions
+
+set DIRNAME=%~dp0
+if "%DIRNAME%"=="" set DIRNAME=.
+@rem This is normally unused
+set APP_BASE_NAME=%~n0
+set APP_HOME=%DIRNAME%
+
+@rem Resolve any "." and ".." in APP_HOME to make it shorter.
+for %%i in ("%APP_HOME%") do set APP_HOME=%%~fi
+
+@rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
+set DEFAULT_JVM_OPTS="-Xmx64m" "-Xms64m"
+
+@rem Find java.exe
+if defined JAVA_HOME goto findJavaFromJavaHome
+
+set JAVA_EXE=java.exe
+%JAVA_EXE% -version >NUL 2>&1
+if %ERRORLEVEL% equ 0 goto execute
+
+echo. 1>&2
+echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. 1>&2
+echo. 1>&2
+echo Please set the JAVA_HOME variable in your environment to match the 1>&2
+echo location of your Java installation. 1>&2
+
+"%COMSPEC%" /c exit 1
+
+:findJavaFromJavaHome
+set JAVA_HOME=%JAVA_HOME:"=%
+set JAVA_EXE=%JAVA_HOME%/bin/java.exe
+
+if exist "%JAVA_EXE%" goto execute
+
+echo. 1>&2
+echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME% 1>&2
+echo. 1>&2
+echo Please set the JAVA_HOME variable in your environment to match the 1>&2
+echo location of your Java installation. 1>&2
+
+"%COMSPEC%" /c exit 1
+
+:execute
+@rem Setup the command line
+
+
+
+@rem Execute gradlew
+@rem endlocal doesn't take effect until after the line is parsed and variables are expanded
+@rem which allows us to clear the local environment before executing the java command
+endlocal & "%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -jar "%APP_HOME%\gradle\wrapper\gradle-wrapper.jar" %* & call :exitWithErrorLevel
+
+:exitWithErrorLevel
+@rem Use "%COMSPEC%" /c exit to allow operators to work properly in scripts
+"%COMSPEC%" /c exit %ERRORLEVEL%
@@ -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
@@ -13,6 +13,7 @@ class WebSocketConfig(
    private val exampleHandler: ExampleHandler,
    private val testHandler: TestHandler
 ) {
+
    @Bean
    fun handlerMapping(): HandlerMapping {
        val map = mapOf(
@@ -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))
+            }
 }
@@ -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")
        }
@@ -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))
    }

@@ -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)
+                }
+            }
    }
 }
@@ -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 }
+    }
+}
@@ -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.")
 }
@@ -41,4 +41,7 @@ app:
  flow-limit:
    count: 1000
    limit-time: 3
-    ban-time: 60
+    ban-time: 60
+
+  ws:
+    token-expire-time: 300