websocket-cj/src/server/WebSocketConnection.cj

/*
 * simcu::websocket.server —— 服务端单条 WebSocket 连接。
 *
 * 由 WebSocketServer 在握手成功后创建并暴露给 on('connection') 回调。
 * 提供 send/sendText/sendBinary/ping/close/terminate，以及连接级事件：
 *   conn.on("message", { (msg) => ... })
 *   conn.on("close", { (code, reason) => ... })
 *   conn.on("error", { (e) => ... })
 */
package simcu::websocket.server

import std.net.IPSocketAddress
import std.net.TcpSocket
import std.sync.Mutex
import std.time.*

import simcu::websocket.common.FrameStream
import simcu::websocket.common.MessageType
import simcu::websocket.common.ReadyState
import simcu::websocket.common.WebSocketMessage
import simcu::websocket.common.WebSocketException
import simcu::websocket.common.WsEvents

/// 服务端单条 WebSocket 连接（事件回调式）。
public class WebSocketConnection {
    private let parent: ?WebSocketServer
    private let frame: FrameStream
    private let stateLock = Mutex()
    private var state: ReadyState = ReadyState.Open

    /// 连接唯一编号（由服务端分配，用于连接集合的引用比较）。
    internal let connId: Int64

    /// 连接级事件：收到消息（不含父级 connection 参数）。
    public var onMessage: (WebSocketMessage) -> Unit = { m => () }
    /// 连接级事件：连接关闭，参数 (关闭码, 关闭原因)。
    public var onClose: (Int64, String) -> Unit = { c, r => () }
    /// 连接级事件：运行时错误。
    public var onError: (Exception) -> Unit = { e => () }

    internal init(parent: ?WebSocketServer, frame: FrameStream, connId: Int64) {
        this.parent = parent
        this.frame = frame
        this.connId = connId
    }

    /// 当前连接状态。
    public prop readyState: ReadyState {
        get() { state }
    }

    /// 连接是否处于 Open。
    public func isOpen(): Bool {
        state == ReadyState.Open
    }

    /// 远端地址（ip:port 或套接字地址字符串）。
    public prop remoteAddress: String {
        get() {
            let sa = frame.remoteAddress()
            if (let Some(ip) <- (sa as IPSocketAddress)) {
                "${ip.address}:${ip.port}"
            } else {
                sa.toString()
            }
        }
    }

    /// 注册连接级事件。支持 "message" / "close" / "error"。
    public func on(event: String, listener: (WebSocketMessage) -> Unit): Unit {
        if (event == WsEvents.message) {
            onMessage = listener
        } else {
            throw WebSocketException("不支持的事件: ${event}")
        }
    }

    /// 注册连接级事件。支持 "close" / "error"。
    public func on(event: String, listener: (Int64, String) -> Unit): Unit {
        if (event == WsEvents.close) {
            onClose = listener
        } else {
            throw WebSocketException("不支持的事件: ${event}")
        }
    }

    /// 注册连接级事件。支持 "error"。
    public func on(event: String, listener: (Exception) -> Unit): Unit {
        if (event == WsEvents.error) {
            onError = listener
        } else {
            throw WebSocketException("不支持的事件: ${event}")
        }
    }

    /// 发送消息。
    public func send(message: WebSocketMessage): Unit {
        ensureOpen()
        if (message.`type` == MessageType.Text) {
            frame.writeFrame(1, message.bytes)
        } else {
            frame.writeFrame(2, message.bytes)
        }
    }

    /// 发送文本消息。
    public func sendText(text: String): Unit {
        ensureOpen()
        frame.writeText(text)
    }

    /// 发送二进制消息。
    public func sendBinary(data: Array<Byte>): Unit {
        ensureOpen()
        frame.writeBinary(data)
    }

    /// 发送 Ping（心跳保活）。
    public func ping(): Unit {
        ensureOpen()
        frame.writePing(Array<Byte>(0, repeat: 0))
    }

    /// 发起关闭握手：发送 close 帧，等待对端回执；5 秒兜底强制断开。
    public func close(code!: Int64 = 1000, reason!: String = ""): Unit {
        synchronized(stateLock) {
            if (state != ReadyState.Open) {
                return
            }
            state = ReadyState.Closing
        }
        try {
            frame.writeClose(code, reason)
        } catch (_) {
            terminate()
            return
        }
        spawn { =>
            sleep(Duration.second * 5)
            if (state == ReadyState.Closing) {
                terminate()
            }
        }
    }

    /// 立即断开底层连接（不发送 close 帧）。
    public func terminate(): Unit {
        synchronized(stateLock) {
            if (state == ReadyState.Closed) {
                return
            }
            state = ReadyState.Closed
        }
        try {
            frame.close()
        } catch (_) {
            ()
        }
    }

    private func ensureOpen(): Unit {
        if (state != ReadyState.Open) {
            throw WebSocketException("连接未打开（readyState=${state}）")
        }
    }

    /// 读循环（由 WebSocketServer 的连接处理协程调用）。
    internal func runReadLoop(): Unit {
        var code: Int64 = 1006
        var reason: String = ""
        try {
            while (true) {
                let m = frame.readMessage()
                match (m) {
                    case None =>
                        break
                    case Some(f) =>
                        if (f.opcode == 8) {
                            let (c, r) = FrameStream.parseClosePayload(f.payload)
                            code = c
                            reason = r
                            if (state == ReadyState.Open) {
                                try {
                                    frame.writeClose(c, r)
                                } catch (_) {
                                    ()
                                }
                            }
                            break
                        } else {
                            let msg = if (f.opcode == 1) {
                                WebSocketMessage(MessageType.Text, f.payload)
                            } else {
                                WebSocketMessage(MessageType.Binary, f.payload)
                            }
                            onMessage(msg)
                            if (let Some(p) <- parent) {
                                p.fireMessage(this, msg)
                            }
                        }
                }
            }
        } catch (e: Exception) {
            code = 1006
            if (state != ReadyState.Closing) {
                onError(e)
                if (let Some(p) <- parent) {
                    p.fireError(e)
                }
            }
        }
        synchronized(stateLock) {
            state = ReadyState.Closed
        }
        try {
            frame.close()
        } catch (_) {
            ()
        }
        if (let Some(p) <- parent) {
            p.removeConnection(this)
        }
        onClose(code, reason)
        if (let Some(p) <- parent) {
            p.fireClose(this, code, reason)
        }
    }
}