330 lines
11 KiB
Plaintext
330 lines
11 KiB
Plaintext
|
/*
|
||
|
|
* simcu::websocket.server —— 事件回调式 WebSocket 服务端（RFC 6455）。
|
||
|
|
*
|
||
|
|
* 用法：
|
||
|
|
* let server = WebSocketServer(bindAt: 8080)
|
||
|
|
* server.on("connection", { (conn) =>
|
||
|
|
* println("新连接: ${conn.remoteAddress}")
|
||
|
|
* conn.on("message", { (msg) => conn.sendText("echo: ${msg.text}") })
|
||
|
|
* })
|
||
|
|
* server.listen()
|
||
|
|
* ...
|
||
|
|
* server.close() // 停止监听并断开所有连接
|
||
|
|
*
|
||
|
|
* 服务端事件：on('connection') / on('message') / on('close') / on('error')。
|
||
|
|
* 连接方法：send / sendText / sendBinary / ping / close / terminate。
|
||
|
|
* 集群方法：broadcast / broadcastText / broadcastBinary / connectionCount。
|
||
|
|
*/
|
||
|
|
package simcu::websocket.server
|
||
|
|
|
||
|
|
import std.collection.ArrayList
|
||
|
|
import std.net.IPSocketAddress
|
||
|
|
import std.net.SocketException
|
||
|
|
import std.net.TcpServerSocket
|
||
|
|
import std.net.TcpSocket
|
||
|
|
import std.sync.Mutex
|
||
|
|
|
||
|
|
import simcu::websocket.common.FrameStream
|
||
|
|
import simcu::websocket.common.MessageType
|
||
|
|
import simcu::websocket.common.WebSocketMessage
|
||
|
|
import simcu::websocket.common.WebSocketException
|
||
|
|
import simcu::websocket.common.WsEvents
|
||
|
|
|
||
|
|
/// 事件回调式 WebSocket 服务端。
|
||
|
|
public class WebSocketServer {
|
||
|
|
private let bindPort: UInt16
|
||
|
|
private let pathFilter: ?String
|
||
|
|
private let maxPayload: Int64
|
||
|
|
private var serverSock: ?TcpServerSocket = None
|
||
|
|
private var running = false
|
||
|
|
private var connections = ArrayList<WebSocketConnection>()
|
||
|
|
private var nextConnId: Int64 = 1
|
||
|
|
private let connLock = Mutex()
|
||
|
|
|
||
|
|
/// 服务端事件：有连接完成握手时触发（参数为连接对象）。
|
||
|
|
public var onConnection: (WebSocketConnection) -> Unit = { conn => () }
|
||
|
|
/// 服务端事件：任意连接收到消息时触发（参数为连接对象与消息）。
|
||
|
|
public var onMessage: (WebSocketConnection, WebSocketMessage) -> Unit = { conn, msg => () }
|
||
|
|
/// 服务端事件：任意连接关闭时触发（参数为连接对象、关闭码、原因）。
|
||
|
|
public var onClose: (WebSocketConnection, Int64, String) -> Unit = { conn, code, reason => () }
|
||
|
|
/// 服务端事件：监听或连接运行时错误时触发（参数为异常）。
|
||
|
|
public var onError: (Exception) -> Unit = { e => () }
|
||
|
|
|
||
|
|
/// @param bindAt 监听端口，0 表示随机空闲端口（listen 后通过 localPort 读取）。
|
||
|
|
/// @param path 仅接受该路径的握手请求（None 表示不限制）。
|
||
|
|
/// @param maxPayload 单条消息最大字节数，默认 64KB。
|
||
|
|
public init(bindAt!: UInt16 = 0, path!: ?String = None, maxPayload!: Int64 = 65536) {
|
||
|
|
this.bindPort = bindAt
|
||
|
|
this.pathFilter = path
|
||
|
|
this.maxPayload = maxPayload
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 注册服务端事件。支持 "connection"。
|
||
|
|
public func on(event: String, listener: (WebSocketConnection) -> Unit): Unit {
|
||
|
|
if (event == WsEvents.connection) {
|
||
|
|
onConnection = listener
|
||
|
|
} else {
|
||
|
|
throw WebSocketException("不支持的事件: ${event}")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 注册服务端事件。支持 "message"。
|
||
|
|
public func on(event: String, listener: (WebSocketConnection, WebSocketMessage) -> Unit): Unit {
|
||
|
|
if (event == WsEvents.message) {
|
||
|
|
onMessage = listener
|
||
|
|
} else {
|
||
|
|
throw WebSocketException("不支持的事件: ${event}")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 注册服务端事件。支持 "close"。
|
||
|
|
public func on(event: String, listener: (WebSocketConnection, 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}")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 启动监听：bind + accept 循环，每连接一个协程处理握手与读循环。
|
||
|
|
public func listen(backlog!: Int64 = 128): Unit {
|
||
|
|
if (running) {
|
||
|
|
throw WebSocketException("服务端已在监听")
|
||
|
|
}
|
||
|
|
let s = TcpServerSocket(bindAt: bindPort)
|
||
|
|
s.backlogSize = backlog
|
||
|
|
s.bind()
|
||
|
|
serverSock = Some(s)
|
||
|
|
running = true
|
||
|
|
spawn { => acceptLoop() }
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 实际监听端口（bindAt=0 时用于发现随机端口）。
|
||
|
|
public prop localPort: UInt16 {
|
||
|
|
get() {
|
||
|
|
let sa = serverSock.getOrThrow().localAddress
|
||
|
|
if (let Some(ip) <- (sa as IPSocketAddress)) {
|
||
|
|
ip.port
|
||
|
|
} else {
|
||
|
|
throw WebSocketException("无法解析本地监听端口")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 当前在线连接数。
|
||
|
|
public prop connectionCount: Int64 {
|
||
|
|
get() {
|
||
|
|
synchronized(connLock) {
|
||
|
|
connections.size
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 停止监听并强制断开所有连接（每连接会触发自身 close 事件）。
|
||
|
|
public func close(): Unit {
|
||
|
|
running = false
|
||
|
|
if (let Some(s) <- serverSock) {
|
||
|
|
try {
|
||
|
|
s.close()
|
||
|
|
} catch (_) {
|
||
|
|
()
|
||
|
|
}
|
||
|
|
serverSock = None
|
||
|
|
}
|
||
|
|
let snapshot = snapshotConnections()
|
||
|
|
for (c in snapshot) {
|
||
|
|
c.terminate()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 向所有在线连接广播消息（单个连接发送失败不影响其余连接）。
|
||
|
|
public func broadcast(message: WebSocketMessage): Unit {
|
||
|
|
let snapshot = snapshotConnections()
|
||
|
|
for (c in snapshot) {
|
||
|
|
try {
|
||
|
|
c.send(message)
|
||
|
|
} catch (_) {
|
||
|
|
()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 广播文本消息。
|
||
|
|
public func broadcastText(text: String): Unit {
|
||
|
|
broadcast(WebSocketMessage.fromText(text))
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 广播二进制消息。
|
||
|
|
public func broadcastBinary(data: Array<Byte>): Unit {
|
||
|
|
broadcast(WebSocketMessage.fromBinary(data))
|
||
|
|
}
|
||
|
|
|
||
|
|
// ===== 内部回调转发（供 WebSocketConnection 调用）=====
|
||
|
|
|
||
|
|
internal func fireMessage(conn: WebSocketConnection, msg: WebSocketMessage): Unit {
|
||
|
|
onMessage(conn, msg)
|
||
|
|
}
|
||
|
|
|
||
|
|
internal func fireClose(conn: WebSocketConnection, code: Int64, reason: String): Unit {
|
||
|
|
onClose(conn, code, reason)
|
||
|
|
}
|
||
|
|
|
||
|
|
internal func fireError(e: Exception): Unit {
|
||
|
|
onError(e)
|
||
|
|
}
|
||
|
|
|
||
|
|
internal func addConnection(conn: WebSocketConnection): Unit {
|
||
|
|
synchronized(connLock) {
|
||
|
|
connections.add(conn)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
internal func allocateConnId(): Int64 {
|
||
|
|
synchronized(connLock) {
|
||
|
|
let id = nextConnId
|
||
|
|
nextConnId += 1
|
||
|
|
id
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
internal func removeConnection(conn: WebSocketConnection): Unit {
|
||
|
|
synchronized(connLock) {
|
||
|
|
var idx: Int64 = -1
|
||
|
|
for (i in 0..connections.size) {
|
||
|
|
if (connections[i].connId == conn.connId) {
|
||
|
|
idx = i
|
||
|
|
break
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if (idx >= 0) {
|
||
|
|
connections.remove(at: idx)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
private func snapshotConnections(): ArrayList<WebSocketConnection> {
|
||
|
|
synchronized(connLock) {
|
||
|
|
let copy = ArrayList<WebSocketConnection>()
|
||
|
|
for (c in connections) {
|
||
|
|
copy.add(c)
|
||
|
|
}
|
||
|
|
copy
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// ===== 监听与握手 =====
|
||
|
|
|
||
|
|
private func acceptLoop(): Unit {
|
||
|
|
while (running) {
|
||
|
|
try {
|
||
|
|
let s = serverSock
|
||
|
|
if (let Some(sock) <- s) {
|
||
|
|
let client = sock.accept()
|
||
|
|
spawn { => handleConnection(client) }
|
||
|
|
} else {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
} catch (e: SocketException) {
|
||
|
|
if (!running) {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
onError(e)
|
||
|
|
} catch (e: Exception) {
|
||
|
|
if (!running) {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
onError(e)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
private func handleConnection(sock: TcpSocket): Unit {
|
||
|
|
let fs = FrameStream(sock, maxPayload, false)
|
||
|
|
try {
|
||
|
|
let request = fs.readHttpHeader()
|
||
|
|
let lines = request.split("\r\n")
|
||
|
|
if (lines.size == 0 || !lines[0].startsWith("GET ")) {
|
||
|
|
reject(sock, 400, "Bad Request")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
let parts = lines[0].split(" ")
|
||
|
|
if (parts.size < 3) {
|
||
|
|
reject(sock, 400, "Bad Request")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
if (let Some(p) <- pathFilter) {
|
||
|
|
if (parts[1] != p) {
|
||
|
|
reject(sock, 404, "Not Found")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
var upgradeOk = false
|
||
|
|
var connectionOk = false
|
||
|
|
var key = ""
|
||
|
|
for (i in 1..lines.size) {
|
||
|
|
let line = lines[i]
|
||
|
|
let ci = line.indexOf(":")
|
||
|
|
if (let Some(c) <- ci) {
|
||
|
|
let name = line[0..c].toAsciiLower()
|
||
|
|
let value = line[c + 1..].trimAscii()
|
||
|
|
if (name == "upgrade") {
|
||
|
|
upgradeOk = value.toAsciiLower().contains("websocket")
|
||
|
|
}
|
||
|
|
if (name == "connection") {
|
||
|
|
connectionOk = value.toAsciiLower().contains("upgrade")
|
||
|
|
}
|
||
|
|
if (name == "sec-websocket-key") {
|
||
|
|
key = value
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if (!upgradeOk || !connectionOk || key == "") {
|
||
|
|
reject(sock, 400, "Bad Request")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
// 101 响应
|
||
|
|
let accept = FrameStream.computeAccept(key)
|
||
|
|
let resp = StringBuilder()
|
||
|
|
resp.append("HTTP/1.1 101 Switching Protocols\r\n")
|
||
|
|
resp.append("Upgrade: websocket\r\n")
|
||
|
|
resp.append("Connection: Upgrade\r\n")
|
||
|
|
resp.append("Sec-WebSocket-Accept: ${accept}\r\n\r\n")
|
||
|
|
sock.write(resp.toString().toArray())
|
||
|
|
sock.flush()
|
||
|
|
|
||
|
|
let conn = WebSocketConnection(Some(this), fs, allocateConnId())
|
||
|
|
addConnection(conn)
|
||
|
|
onConnection(conn)
|
||
|
|
conn.runReadLoop()
|
||
|
|
} catch (e: Exception) {
|
||
|
|
// 握手阶段失败：上报服务端 error 事件
|
||
|
|
onError(e)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
private func reject(sock: TcpSocket, status: Int64, message: String): Unit {
|
||
|
|
try {
|
||
|
|
let resp = "HTTP/1.1 ${status} ${message}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
|
||
|
|
sock.write(resp.toArray())
|
||
|
|
sock.flush()
|
||
|
|
} catch (_) {
|
||
|
|
()
|
||
|
|
}
|
||
|
|
try {
|
||
|
|
sock.close()
|
||
|
|
} catch (_) {
|
||
|
|
()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|