SpringBoot WebSocket 代理模式、客户端模式

import com.fasterxml.jackson.databind.ObjectMapper
import org.slf4j.LoggerFactory
import org.springframework.web.socket.*
import org.springframework.web.socket.client.WebSocketClient
import org.springframework.web.socket.handler.AbstractWebSocketHandler
import java.util.concurrent.*
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicInteger

/**
 * 包装上游会话及其状态，用于管理授权和心跳
 *
 * @param session WebSocket 上游会话
 * @param authorized 是否已通过授权验证
 * @param downConnected 下游连接是否已建立
 * @param lastHeartbeat 最近心跳时间戳（毫秒）
 */
data class WebSocketProxySession(
    val session: WebSocketSession,
    var authorized: Boolean = false,
    var downConnected: Boolean = false,
    var lastHeartbeat: Long = System.currentTimeMillis()
)

/**
 * 通用 WebSocket 代理抽象类
 *
 * 负责管理上游和下游的连接生命周期、消息转发以及超时清理
 *
 * 使用方式：
 * 1. 实现核心抽象方法：
 *    - registerPath: 定义代理路由
 *    - onUpstreamFirstMessage: 处理上游首条消息并进行授权
 *    - downstreamUri: 获取下游 URI
 *    - transformUpstream: 上游→下游 转换逻辑
 *    - transformDownstream: 下游→上游 转换逻辑
 * 2. 可选覆盖钩子：
 *    - onUpstreamOpen: 上游连接初始化
 *    - onAuthSuccess: 授权成功回调
 *    - onUpstreamFirstMessageIsNull: 授权失败处理
 *    - onSessionClosed: 会话关闭后处理
 *
 * @param objectMapper 用于 JSON 序列化/反序列化
 * @param client WebSocket 客户端，用于建立下游连接
 * @author ThatCoder
 */
abstract class IWebSocketProxier(
    val objectMapper: ObjectMapper,
    private val client: WebSocketClient
) : AbstractWebSocketHandler() {
    /** 代理接入路径 */
    abstract val registerPath: String

    /** 会话超时时间，默认 10 分钟 */
    open val sessionTimeoutMillis: Long = 10 * 60 * 1000

    private val logger = LoggerFactory.getLogger(this::class.java)
    private val sessions = ConcurrentHashMap<String, WebSocketProxySession>()
    private val downstreamContexts = ConcurrentHashMap<String, DownstreamContext>()
    private val scheduler = Executors.newSingleThreadScheduledExecutor(
        NamedThreadFactory("proxy-session-timeout-")
    )

    init {
        // 定期清理超时会话
        scheduler.scheduleAtFixedRate(
            { cleanupExpired() },
            sessionTimeoutMillis,
            sessionTimeoutMillis,
            TimeUnit.MILLISECONDS
        )
    }

    override fun afterConnectionEstablished(session: WebSocketSession) {
        logger.info("Upstream connected: ${session.id}")
        sessions[session.id] = WebSocketProxySession(session)
        onUpstreamOpen(sessions[session.id]!!)
    }

    override fun handleMessage(session: WebSocketSession, message: WebSocketMessage<*>) {
        val proxy = sessions[session.id] ?: return
        if (!proxy.authorized) {
            val ok = onUpstreamFirstMessage(proxy, message)
            if (!ok) {
                onUpstreamFirstMessageIsNull(proxy)
                closeSession(session.id)
                return
            }
            proxy.authorized = true
            onAuthSuccess(proxy)
            connectDownstream(session.id)
            downstreamContexts[session.id]?.pending?.offer(clone(message))
            return
        }
        val ctx = downstreamContexts[session.id] ?: return
        if (!ctx.downConnected.get()) {
            ctx.pending.offer(clone(message))
        } else {
            ctx.sendToDownstream(transformUpstream(message))
        }
    }

    override fun afterConnectionClosed(session: WebSocketSession, status: CloseStatus) {
        logger.info("Upstream closed: ${session.id}")
        closeSession(session.id)
    }

    /**
     * 向所有上游会话发送心跳，维持长连接
     */
    fun sendHeartbeat() {
        val ping = PingMessage()
        sessions.values.forEach {
            try {
                it.session.sendMessage(ping)
            } catch (_: Exception) {
                // 忽略发送失败
            }
        }
    }

    // ---------- 可覆盖钩子 ----------

    /** 上游连接建立后回调 */
    protected open fun onUpstreamOpen(proxy: WebSocketProxySession) = Unit

    /**
     * 上游首条消息处理并授权
     * @return true 表示通过，false 则触发授权失败
     */
    protected abstract fun onUpstreamFirstMessage(
        proxy: WebSocketProxySession,
        message: WebSocketMessage<*>
    ): Boolean

    /** 授权失败发送给上游的消息 */
    protected open fun onUpstreamFirstMessageIsNull(proxy: WebSocketProxySession) {
        val err = mapOf("finish" to true, "error" to "身份认证失败")
        proxy.session.sendMessage(TextMessage(objectMapper.writeValueAsString(err)))
    }

    /** 授权成功后回调 */
    protected open fun onAuthSuccess(proxy: WebSocketProxySession) = Unit

    /** 根据上游会话获取下游 URI */
    protected abstract fun downstreamUri(proxy: WebSocketProxySession): String

    /** 上游→下游 消息转换 */
    protected abstract fun transformUpstream(message: WebSocketMessage<*>): WebSocketMessage<*>

    /** 下游→上游 消息转换 */
    protected abstract fun transformDownstream(message: WebSocketMessage<*>): WebSocketMessage<*>

    /** 会话关闭后回调 */
    protected open fun onSessionClosed(proxy: WebSocketProxySession) = Unit

    // ---------- 内部逻辑 ----------

    /**
     * 建立下游连接，并将后续消息路由到 DownstreamContext
     */
    private fun connectDownstream(sessionId: String) {
        val proxy = sessions[sessionId]!!
        val ctx = DownstreamContext(proxy)
        downstreamContexts[sessionId] = ctx
        client.execute(object : AbstractWebSocketHandler() {
            override fun afterConnectionEstablished(down: WebSocketSession) {
                logger.info("Downstream connected for: $sessionId")
                ctx.downConnected.set(true)
                ctx.downstream = down
                while (true) {
                    val msg = ctx.pending.poll() ?: break
                    ctx.sendToDownstream(transformUpstream(msg))
                }
            }

            override fun handleMessage(down: WebSocketSession, msg: WebSocketMessage<*>) {
                proxy.session.sendMessage(transformDownstream(msg))
            }

            override fun afterConnectionClosed(down: WebSocketSession, status: CloseStatus) {
                logger.warn("Downstream closed early: ${status.code}")
                closeSession(sessionId)
            }
        }, downstreamUri(proxy))
    }

    /** 关闭并清理指定会话 */
    private fun closeSession(sessionId: String) {
        sessions.remove(sessionId)?.also { onSessionClosed(it) }
        downstreamContexts.remove(sessionId)?.closeAll()
    }

    /** 清理超时会话 */
    private fun cleanupExpired() {
        val now = System.currentTimeMillis()
        sessions.entries
            .filter { now - it.value.lastHeartbeat > sessionTimeoutMillis }
            .forEach { closeSession(it.key) }
    }

    /** 克隆消息以避免并发问题 */
    private fun clone(msg: WebSocketMessage<*>): WebSocketMessage<*> = when (msg) {
        is TextMessage -> TextMessage(msg.payload)
        is BinaryMessage -> BinaryMessage(msg.payload.asReadOnlyBuffer())
        else -> msg
    }

    /**
     * 管理下游消息发送及队列
     */
    private class DownstreamContext(proxy: WebSocketProxySession) {
        @Volatile var downstream: WebSocketSession? = null
        val downConnected = AtomicBoolean(false)
        val pending = ConcurrentLinkedQueue<WebSocketMessage<*>>()
        private val executor: ExecutorService = ThreadPoolExecutor(
            4, 16, 60, TimeUnit.SECONDS,
            LinkedBlockingQueue(1000),
            NamedThreadFactory("proxy-send-${proxy.session.id}")
        )

        /** 将消息异步发送到下游 */
        fun sendToDownstream(msg: WebSocketMessage<*>) {
            executor.execute {
                try {
                    downstream?.sendMessage(msg)
                } catch (e: Exception) {
                    LoggerFactory.getLogger("DownstreamLogger").error("Send downstream failed", e)
                }
            }
        }

        /** 关闭下游并清理资源 */
        fun closeAll() {
            try {
                downstream?.close()
            } catch (_: Exception) {
            }
            executor.shutdownNow()
            pending.clear()
        }
    }

    /** 为线程池生成可读性线程名 */
    private class NamedThreadFactory(prefix: String) : ThreadFactory {
        private val cnt = AtomicInteger(1)
        private val name = "${prefix}-${cnt.getAndIncrement()}"
        override fun newThread(r: Runnable) = Thread(r, name)
    }
}

import com.bidr.waterx.transpond.config.extension.fieldJust
import com.bidr.waterx.transpond.config.extension.fieldRemove
import com.bidr.waterx.transpond.config.extension.fieldRename
import com.bidr.waterx.transpond.config.extension.putMap
import com.bidr.waterx.transpond.config.extension.toObjectNode
import com.bidr.waterx.transpond.config.extension.toTextMessage
import com.fasterxml.jackson.databind.ObjectMapper
import com.fasterxml.jackson.databind.node.ArrayNode
import org.slf4j.LoggerFactory
import org.springframework.stereotype.Component
import org.springframework.web.socket.CloseStatus
import org.springframework.web.socket.TextMessage
import org.springframework.web.socket.WebSocketMessage
import org.springframework.web.socket.client.WebSocketClient

/** ASR 代理实现 */
@Component
class AsrProxier(
    objectMapper: ObjectMapper,
    webSocketClient: WebSocketClient,
    private val akService: IAKService
) : IWebSocketProxier(objectMapper, webSocketClient) {
    private val paramProxy = mapOf(
        "id" to "wav_name",
        "finish" to "is_speaking",
        "answer" to "text"
    )

    private val logger = LoggerFactory.getLogger(this::class.java)

    override val registerPath = "/ws/asr"

    // 鉴权服务
    override fun onUpstreamFirstMessage(proxy: WebSocketProxySession, message: WebSocketMessage<*>): Boolean {
        val node = message.toObjectNode(objectMapper) ?: return false
        val ak = node.get("ak")?.asText() ?: return false
        return akService.check(ak)
    }

    override fun downstreamUri(proxy: WebSocketProxySession) = "ws://localhost:10095"

    // 处理上游消息适配成FUNASR接收类型
    override fun transformUpstream(message: WebSocketMessage<*>) = when (message) {
        is TextMessage -> runCatching {
            val forward = message.toObjectNode(objectMapper)?.fieldRename(paramProxy) ?: return message
            forward.get("is_speaking")?.let {
                val finished = it.asBoolean(false)
                if (!finished) forward.putMap(mapOf(
                    "language" to "zn",
                    "itn" to false,
                    "hotwords" to "{\"阿里巴巴\":20,\"hello world\":40}"
                ))
                forward.put("is_speaking", !finished)
            }
            forward.get("mode")?.let {
                if (listOf("mixed","online").contains(it.asText())) {
                    val arr = objectMapper.createArrayNode().add(5).add(10).add(5)
                    forward.set<ArrayNode>("chunk_size", arr)
                    forward.put("chunk_interval", 10)
                    if (it.asText() == "mixed") forward.put("mode", "2pass")
                }
            }
            forward.fieldRemove(listOf("ak"))
            logger.info("transformUpstream: $forward")
            forward.toTextMessage(objectMapper)
        }.getOrDefault(message)
        else -> message
    }

    // 处理下游消息适配成客户接收类型
    override fun transformDownstream(message: WebSocketMessage<*>) = when (message) {
        is TextMessage -> {
            val forward = message.toObjectNode(objectMapper)
                ?.fieldRename(paramProxy.toMutableMap().plus("finish" to "is_final"), true)
                ?: return message
            forward.putMap(mapOf(
                "mode" to when (forward.get("mode")?.asText() ?: "2pass-offline") {
                    "2pass-online" -> "online"
                    "2pass-offline" -> "offline"
                    else -> forward.get("mode").asText()
                },
                "timestamp" to System.currentTimeMillis()
            ))
            forward.fieldJust(paramProxy.keys.plus("mode").toList())
            logger.info("transformDownstream: $forward")
            forward.toTextMessage(objectMapper)
        }
        else -> message
    }

    override fun onUpstreamFirstMessageIsNull(proxy: WebSocketProxySession) {
        super.onUpstreamFirstMessageIsNull(proxy)
        proxy.session.close(CloseStatus.POLICY_VIOLATION)
    }
}

import com.fasterxml.jackson.databind.ObjectMapper
import org.slf4j.LoggerFactory
import org.springframework.web.socket.*
import org.springframework.web.socket.handler.AbstractWebSocketHandler
import java.util.concurrent.*
import java.util.concurrent.atomic.AtomicInteger

/**
 * 会话、用户信息的包装类
 */
data class WebSocketSenderSession<T>(
    val session: WebSocketSession,
    val user: T,
    /** 心跳超时标志 */
    val lastHeartbeat: Long = System.currentTimeMillis()
)

/**
 * WebSocket 发布者抽象接口，用于构建支持用户认证、心跳维护、消息广播的通用 WebSocket 服务。
 *
 * ### 使用方式
 *
 * #### 必须实现
 * > 继承本类并实现以下核心抽象方法
 * - [registerPath]：注册的路径，WebSocket 接入入口
 * - [onUpstreamFirstMessage]：处理上游客户端首次连接时的消息，一般用于身份验证，返回的用户信息将用于标识会话；若返回 null，连接将被关闭
 * - [onUpstreamMessage]：处理客户端后续发送的消息
 *
 * #### 可选重写
 * - [onUpstreamOpen]：连接建立但未发送任何消息时的初始化回调
 * - [onSessionClosed]：连接关闭后的回调处理
 * - [onUpstreamFirstMessageIsNull]：首次消息认证失败时的回调，默认发送错误信息
 * - [onAuthSuccess]：首次消息认证通过后的回调
 * - [transformPublish]：发送消息前进行的消息变换
 *
 * ### 会话管理
 * - 会话信息以 [WebSocketSenderSession] 包装，包含 `session`、用户信息及心跳时间
 * - 默认 10 分钟未活跃会话将被关闭，可通过 [sessionTimeoutMillis] 调整
 *
 * ### 发布功能
 * - [publishAll]：向所有连接发布消息
 * - [publishByFilter]：根据过滤条件发布消息
 * - [publishSender]：向单个连接发送消息
 * - [sendHeartbeat]：向所有连接发送 Ping 消息，维持长连接
 *
 * @param objectMapper Jackson 用于 JSON 序列化/反序列化
 * @param T 用户类型，需由 [onUpstreamFirstMessage] 提供
 * @author ThatCoder
 */
abstract class IWebSocketPublisher<T>(
    private val objectMapper: ObjectMapper
) : AbstractWebSocketHandler() {

    abstract val registerPath: String

    /** 会话超时毫秒数 默认十分钟 */
    val sessionTimeoutMillis: Long = 10*60*1000

    private val logger = LoggerFactory.getLogger(this::class.java)

    /** 所有会话管理器 */
    val sessions = ConcurrentHashMap<String, WebSocketSenderSession<T>>()
    private val scheduler = Executors.newSingleThreadScheduledExecutor(NamedThreadFactory("session-timeout-"))

    init {
        // 定期清理超时会话
        scheduler.scheduleAtFixedRate({ cleanupExpired() }, sessionTimeoutMillis, sessionTimeoutMillis, TimeUnit.MILLISECONDS)
    }

    override fun afterConnectionEstablished(session: WebSocketSession) {
        logger.info("Client connected: ${session.id}")
        onUpstreamOpen(session)
    }

    override fun handleMessage(session: WebSocketSession, message: WebSocketMessage<*>) {
        // 首次消息处理授权与注册
        if (!sessions.containsKey(session.id)) {
            val user = onUpstreamFirstMessage(session, message)
            if (user == null) {
                onUpstreamFirstMessageIsNull(session)
                session.close(CloseStatus.POLICY_VIOLATION)
                return
            }
            sessions[session.id] = WebSocketSenderSession(session, user)
            onAuthSuccess(sessions[session.id]!!)
            logger.info("Session registered: ${session.id} -> $user")
            return
        }
        // 心跳更新或具体消息处理
        onUpstreamMessage(sessions[session.id]!!, message)
    }

    override fun afterConnectionClosed(session: WebSocketSession, status: CloseStatus) {
        logger.info("Client closed: ${session.id} (${status.reason})")
        sessions.remove(session.id)?.let { onSessionClosed(it) }
    }

    /**
     * 全局发布消息
     * @param message 消息
     */
    fun publishAll(message: WebSocketMessage<*>) {
        sessions.values.forEach { sender -> send(sender, message) }
    }

    /**
     * 按过滤器发布
     * @param filter 过滤器
     * @param message 消息
     */
    fun publishByFilter(filter: (WebSocketSenderSession<T>) -> Boolean, message: WebSocketMessage<*>) {
        sessions.values.filter(filter).forEach { send(it, message) }
    }

    /**
     * 发送消息给指定会话
     * @param sender 会话
     * @param message 消息
     */
    fun publishSender(sender: WebSocketSenderSession<T>, message: WebSocketMessage<*>) {
        send(sender, message)
    }

    /** 发送心跳 */
    fun sendHeartbeat() {
        val ping = PingMessage()
        sessions.values.forEach {
            try { it.session.sendMessage(ping) } catch (_: Exception) {}
        }
    }

    // ========== 子类扩展点 ===========

    /**
     * 会话首次创建时调用
     * @param session 会话
     */
    protected open fun onUpstreamOpen(session: WebSocketSession) = Unit

    /**
     * 会话消息
     * @param sender 会话对象
     * @param message 消息
     */
    protected abstract fun onUpstreamMessage(sender: WebSocketSenderSession<T>, message: WebSocketMessage<*>)

    /**
     * 会话关闭时调用
     * @param sender 会话对象
     */
    protected open fun onSessionClosed(sender: WebSocketSenderSession<T>) = Unit

    /**
     * 会话首条消息
     *
     * 通常在验证用户权限时调用
     * @param session 会话
     * @param message 消息
     * @return 用户信息 如果返回null则触发 onUpstreamFirstMessageIsNull
     * @see onUpstreamFirstMessageIsNull
     */
    protected abstract fun onUpstreamFirstMessage(
        session: WebSocketSession,
        message: WebSocketMessage<*>
    ): T?

    /**
     * 会话首条消息处理为空时调用
     * @param session 会话
     */
    protected open fun onUpstreamFirstMessageIsNull(session: WebSocketSession) {
        val err = mapOf("error" to "身份认证失败")
        session.sendMessage(TextMessage(objectMapper.writeValueAsString(err)))
    }

    /**
     * 认证成功后执行
     * @param sender 会话
     */
    protected open fun onAuthSuccess(sender: WebSocketSenderSession<T>) = Unit

    /** 清理超时会话 */
    private fun cleanupExpired() {
        val now = System.currentTimeMillis()
        sessions.values.filter { now - it.lastHeartbeat > sessionTimeoutMillis }
            .forEach {
                try { it.session.close(CloseStatus.SESSION_NOT_RELIABLE) } catch (_: Exception) {}
                sessions.remove(it.session.id)
                logger.info("Session timeout removed: ${it.session.id}")
            }
    }

    private fun send(sender: WebSocketSenderSession<T>, message: WebSocketMessage<*>) {
        if (!sender.session.isOpen) return
        try {
            sender.session.sendMessage(transformPublish(message, sender))
        } catch (e: Exception) {
            logger.error("Publish to ${sender.session.id} failed", e)
        }
    }

    protected open fun transformPublish(
        message: WebSocketMessage<*>,
        sender: WebSocketSenderSession<T>
    ): WebSocketMessage<*> = message

    private class NamedThreadFactory(prefix: String) : ThreadFactory {
        private val cnt = AtomicInteger(1)
        private val name = prefix + cnt.getAndIncrement()
        override fun newThread(r: Runnable): Thread {
            return Thread(r, name)
        }
    }
}

import com.bidr.waterx.transpond.config.extension.toObjectNode
import com.bidr.waterx.transpond.config.extension.toTextMessage
import com.fasterxml.jackson.databind.ObjectMapper
import org.springframework.stereotype.Component
import org.springframework.web.socket.TextMessage
import org.springframework.web.socket.WebSocketMessage
import org.springframework.web.socket.WebSocketSession

data class ChatUser(val userid: String, val name: String)

/**
 * 聊天室发布者
 */
@Component
class ChatPublisher(private val objectMapper: ObjectMapper) : IWebSocketPublisher<ChatUser>(objectMapper) {

    override val registerPath = "/ws/chat"

    override fun onUpstreamFirstMessage(session: WebSocketSession, message: WebSocketMessage<*>): ChatUser? {
        val message = message.toObjectNode() ?: return null
        val ak = message.get("ak")?.asText() ?: return null
        val name = message.get("name")?.asText() ?: return null
        if (ak != "123456") return null
        // 创建用户
        val user = ChatUser( session.id, name)
        // 给该用户发送欢迎信息
        session.sendMessage(TextMessage("Hi, $name. Please chat friendly!"))
        // 群发用户入群提示
        publishAll(TextMessage("$name've joined the chat room."))
        return user
    }

    override fun onSessionClosed(sender: WebSocketSenderSession<ChatUser>) {
        // 群发用户离开提示
        publishAll(TextMessage("${sender.user.name} has left the chat room."))
    }

    override fun onUpstreamMessage(sender: WebSocketSenderSession<ChatUser>, message: WebSocketMessage<*>) {
        // 转发用户消息至群聊
        publishAll(objectMapper.createObjectNode().apply {
            put("type", "chat")
            putPOJO("user", sender.user)
            putPOJO("message", message.toObjectNode())
        }.toTextMessage())
    }
}

package cn.uwant.auto.config

import IWebSocketProxier
import IWebSocketPublisher
import jakarta.websocket.ContainerProvider
import org.springframework.context.annotation.Bean
import org.springframework.context.annotation.Configuration
import org.springframework.web.socket.client.WebSocketClient
import org.springframework.web.socket.client.standard.StandardWebSocketClient
import org.springframework.web.socket.config.annotation.EnableWebSocket
import org.springframework.web.socket.config.annotation.WebSocketConfigurer
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry
import org.springframework.context.annotation.Lazy
import kotlin.collections.map

/**
 * WebSocket配置
 * @author ThatCoder
 */
@Configuration
@EnableWebSocket
class WebSocketConfig(
    @Lazy private val proxies: List<IWebSocketProxier>,
    @Lazy private val publishers: List<IWebSocketPublisher<*>>
) : WebSocketConfigurer {
    override fun registerWebSocketHandlers(registry: WebSocketHandlerRegistry) {
        proxies.map {
            registry.addHandler(it, it.registerPath).setAllowedOrigins("*")
        }
        publishers.map {
            registry.addHandler(it, it.registerPath).setAllowedOrigins("*")
        }
    }
    @Bean
    fun webSocketClient(): WebSocketClient {
        val container = ContainerProvider.getWebSocketContainer()
        return StandardWebSocketClient(container)
    }
}