feat(ws): release generic WebSocket client and server v0.31.0-preview
This commit is contained in:
149
ws/README.md
Normal file
149
ws/README.md
Normal file
@@ -0,0 +1,149 @@
|
||||
# WebSocket
|
||||
|
||||
`ws` 提供通用 Go Client、Server 和单次物理连接 Session。核心使用标准 RFC 6455 WebSocket，可连接浏览器或第三方 WS 服务。`hub`、`rpc`、`observe` 按需引入；不预设用户、租户、设备、错误码或数据存储。
|
||||
|
||||
- [回显、分组广播与浏览器示例](../examples/websocket/README.md)
|
||||
- [双向 RPC 与订阅恢复示例](../examples/websocketRPC/README.md)
|
||||
- [RPC 协议与订阅规则](rpc/README.md)
|
||||
- [日志与指标适配](observe/README.md)
|
||||
- [测试与基准报告](../reports/websocket-validation-2026-10-29.md)
|
||||
|
||||
## 独立 Server 与 Client
|
||||
|
||||
```go
|
||||
server, err := ws.NewServer(ws.ServerConfig{
|
||||
MaxConnections: 4096,
|
||||
// 省略时允许同源和无 Origin 的非浏览器请求。
|
||||
AllowedOrigins: []string{"https://console.example.com"},
|
||||
})
|
||||
if err != nil { return err }
|
||||
mux.Handle("/ws", server.Handler(ws.Handlers{
|
||||
OnMessage: func(ctx context.Context, session *ws.Session, msg ws.Message) error {
|
||||
return session.Send(ctx, msg)
|
||||
},
|
||||
}))
|
||||
// HTTP 宿主退出后、依赖释放前调用 server.Shutdown(ctx)。
|
||||
|
||||
client, err := ws.NewClient(ws.ClientConfig{ProxyMode: "direct"})
|
||||
if err != nil { return err }
|
||||
session, err := client.Dial(workCtx, "wss://service.example.com/ws", ws.Handlers{
|
||||
OnMessage: func(ctx context.Context, session *ws.Session, msg ws.Message) error {
|
||||
// msg.Type 是 ws.Text 或 ws.Binary；此回调中处理业务消息。
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if err != nil { return err }
|
||||
err = session.SendJSON(ctx, map[string]string{"message": "hello"})
|
||||
// 退出前 client.Shutdown(ctx)，确保回调结束后再关闭其依赖。
|
||||
```
|
||||
|
||||
`NewServer` / `NewClient` 只验证并构造，无监听或拨号。Server 是 `http.Handler`，复用 HTTP 监听器及 TLS；客户端 `TLSConfig` 支持系统信任链、自定义 CA 与客户端证书。注入的证书/CA/回调对象由调用方保持并发安全且不再修改。TLS 默认正常验证，不自动跳过验证。
|
||||
|
||||
`Server.Accept(w,r,handlers)` 会写出握手成功或错误响应，成功后启动托管会话并返回。原 HTTP Handler 可以立即返回；不要保留整个 `*httpx.Context`、Request 或可变鉴权对象。Server 的会话生命周期由 `ServerConfig.Context`（默认 Background）管理，握手 context 仅用于握手。Client.Dial 的 ctx 持有会话寿命，其配置的 HandshakeTimeout 只限制握手。
|
||||
|
||||
## 握手与安全配置
|
||||
|
||||
`ServerConfig` 支持 `Session`、`MaxConnections`、`HandshakeTimeout`、`AllowedOrigins`、`RequireOrigin`、`Subprotocols`、`RequireSubprotocol`、`EnableCompression`、`HandshakeRate`、`HandshakeBurst`。连接上限包含进行中的握手。速率为每秒令牌数、burst 为突发数，零 rate 默认关闭；超过速率返回 429，连接上限/停机返回 503。HTTP 读头超时和头大小限制仍由宿主配置。
|
||||
|
||||
Origin 校验与 CORS 独立：
|
||||
|
||||
- 默认只接受与请求的真实 scheme/host/port 一致的 Origin；缺少 Origin 默认允许，便于 Go 等非浏览器客户端，设置 RequireOrigin 可拒绝。
|
||||
- 显式 AllowedOrigins 是完整来源白名单，支持明确的 `http://` / `https://` 和端口；不支持通配 `*`、路径、userinfo、null 或重复 Origin。默认端口归一化。
|
||||
- 不信任 Forwarded/X-Forwarded-* 来推断来源。TLS 在代理终止时，应填写实际外部 Origin 白名单。
|
||||
- Origin 不是身份认证。`Prepare(ctx,*http.Request)(Handlers,error)` 在升级前验证业务身份，同时可为每连接创建 handlers。失败默认 401；返回 `*ws.RejectError{StatusCode:403}` 可选择 4xx。响应不包含底层错误。
|
||||
- 浏览器原生 WebSocket 不能自定义 Authorization Header，业务可采用已受保护的 Cookie 或短期握手票据。若使用 Cookie，继续严格检查 Origin；避免在 URL 中放长期凭证，并配置代理访问日志脱敏。
|
||||
|
||||
RequireSubprotocol 要求协商成功，否则握手失败；纯 WS 默认无必选子协议。RPC 使用 `go-omni.rpc.v1`，两端显式声明并要求它。
|
||||
|
||||
## Session 的收发与关闭
|
||||
|
||||
Session 提供 `ID / Context / Done / Err / Subprotocol / Stats / Send / TrySend / SendJSON / Close / Abort / Wait`。
|
||||
|
||||
| 操作 | 语义 |
|
||||
| --- | --- |
|
||||
| `Send(ctx,Message)` | 有界队列入场后等待写入结果；队列满立即返回 ErrQueueFull，不无限等待空位。 |
|
||||
| `TrySend(Message)` | 只确认入队，随后仍可能写失败；用于广播、单向通知。 |
|
||||
| `SendJSON(ctx,value)` | JSON 编码后按 Text 发送，适用于应用自己的协议。 |
|
||||
| `Close(code,reason)` | 幂等启动关闭，取消回调 context，尽力发送已接纳且未取消的出站消息，再发关闭帧；不等待自身回调。 |
|
||||
| `Abort()` | 立即关闭 socket，仍需 Wait 才能证明回调退出。 |
|
||||
| `Wait(ctx)` / `Done()` | 等待读、写、分发、监控及 OnClose 全部结束；Wait 返回运行错误或等待超时。 |
|
||||
|
||||
文本严格验证 UTF-8；二进制保持字节内容。入队复制发送 payload，Send/TrySend 返回后调用方可以复用原切片。不得在调用进行中并发修改它。接收 payload 不复用，回调可以保留，但应用自行承担保留后的内存。业务不得直接抢占底层 reader/writer。
|
||||
|
||||
单个会话 OnOpen、OnMessage 顺序执行，OnClose 在 I/O 和消息回调结束后执行；不同会话可并发。回调必须响应 ctx 取消，不要在回调内调用自身 Wait、Client/Server.Shutdown 或 App.Close，可调用 Session.Close 或 Server.Stop。OnOpen 中不要等待必须依赖 OnMessage 才完成的操作；RPC Call/订阅恢复应在 WaitReady 后从独立受管理任务调用。
|
||||
|
||||
排队时 Send ctx 取消会跳过该消息；写入期间取消可能已部分送出，此时终止物理连接。成功写出也不代表对端业务处理成功。框架不缓存离线消息、不自动重放；业务需要 ACK、去重、持久化时显式设计。Close 不是业务事务提交屏障，会取消当前业务并丢弃尚未分发的入站消息。
|
||||
|
||||
## 资源上限与心跳
|
||||
|
||||
`SessionConfig` 零字段使用默认值，负值或超过支持边界的值在构造前报错。明确关闭心跳使用 `DisableHeartbeat`，零时间不表示禁用。
|
||||
|
||||
| 配置 | 默认 |
|
||||
| --- | --- |
|
||||
| MaxMessageBytes | 1 MiB，包括解压后的消息 |
|
||||
| ReadQueueSize / WriteQueueSize | 每方向 64 条 |
|
||||
| ReadQueueBytes / WriteQueueBytes | 每方向 4 MiB |
|
||||
| WriteTimeout | 10 秒 |
|
||||
| PingInterval / PongTimeout | 30 / 60 秒 |
|
||||
| CloseTimeout | 5 秒 |
|
||||
|
||||
队列计数包含正在执行的回调/正在写的消息；此外还需计入一条正在解码的消息、协议缓冲和应用持有的对象。实例总内存还受连接数影响。接收队列溢出关闭该连接；发送满明确报错，广播逐连接返回错误。有限队列允许持续处理控制帧，不让慢业务无限阻塞 reader。压缩默认关闭；启用后仍执行解压大小限制，但会消耗额外 CPU，应结合并发连接限额验证。
|
||||
|
||||
协议 Ping/Pong 由底层持续处理；浏览器自动回复 Pong。心跳监测连接活性，不证明业务健康。
|
||||
|
||||
## Client 重连与网络设置
|
||||
|
||||
`ClientConfig` 支持 Session、HandshakeTimeout（默认 10s）、MaxConnections（默认 64）、Subprotocols/RequireSubprotocol、EnableCompression、ProxyMode/ProxyURL、TLSConfig、Headers、HeaderProvider、Prepare、Reconnect、ShouldReconnect。
|
||||
|
||||
- `Dial` 只拨一次。`Run(ctx,url,handlers)` 管理一条重连链；同 Client 可以并发运行多条连接，统一限制和关闭。
|
||||
- `Reconnect.Enabled` 显式开启；默认 minBackoff=100ms、maxBackoff=30s、stableAfter=30s。指数退避带抖动；短暂反复断线不重置预算。MaxAttempts/MaxElapsed 可限制重试次数和时间；开启重连后两者为零表示不限。
|
||||
- 默认不重试 401/403、协议不兼容、证书验证错误、正常 1000 关闭；远端先发 1001（going away）、1012 等重启关闭码和临时网络故障可重试。本地主动 Close/Client.Stop 不触发重连。ShouldReconnect 可改变分类，不能绕过显式关闭和预算。
|
||||
- 每次握手调用 HeaderProvider(ctx)，用于刷新凭证；返回头会复制。Prepare(ctx) 可为每代创建新的 RPC Peer/handlers，禁止复用已经绑定的 Peer。
|
||||
- ProxyMode 默认为 environment，可显式 direct 或 url；固定代理支持 HTTP CONNECT、socks5、socks5h。HTTPS 代理当前明确拒绝，不会降级直连。
|
||||
- 自有 HTTP 包装链 `httpc.Do`、`telemetry.Transport`、`resilience.Transport` 不直接用于 WS 升级；WS 独立处理握手、消息与连接寿命。
|
||||
|
||||
`Client.Close` / `Stop` 停止入场、取消握手/退避并发起会话关闭；`Drain` / `Shutdown` 等待所有 Run、握手、会话和回调结束。Server 同样提供 Stop/Drain/Shutdown。终止后对象不可重新启动。
|
||||
|
||||
## Hub 与扩展
|
||||
|
||||
`hub.New(Config)` 提供 `Add / Remove / Join / Leave / Send / SendGroup / Broadcast / Kick / Stats / Close`。Session 实现其 `Connection` 接口，也可实现其他有界非阻塞连接适配器。组键是业务自定义字符串，Join 前由业务鉴权。
|
||||
|
||||
Send/Broadcast/SendGroup 返回 `Result{Attempted,Queued,Errors}`，逐连接尝试入队，慢连接不会等待网络写出。Queued 不等于远端收到。Hub 借用连接，Close 只清注册和 watcher；Kick 才主动关闭，Session.Done 后自动清理。上限可配置连接数、组数、总成员关系、每连接组数量。
|
||||
|
||||
Hub 是单进程注册表。跨节点广播由业务消息总线投递到各节点的本地 Hub；此版不含分布式路由、离线存储或 exactly-once 保证。通过接口扩展时仍需保持 TrySend 有界、非阻塞、并发安全。
|
||||
|
||||
## GoOmni 接入与停机
|
||||
|
||||
独立 App：
|
||||
|
||||
```go
|
||||
application := omni.New()
|
||||
server, err := ws.NewServer(ws.ServerConfig{})
|
||||
if err != nil { return err }
|
||||
if err = wsomni.Mount(application, "/ws", server, handlers); err != nil { return err }
|
||||
// Mount 同时登记 Server.Stop/Drain，App.OnShutdown 的资源晚于会话释放。
|
||||
return application.RunContext(ctx, ":8080")
|
||||
```
|
||||
|
||||
`wsomni` 的 import 是 `git.yuntuo.hk/k2safe/go-omni/ws/omni`。App 新增通用 `OnStop(name,hook)` / `OnDrain(name,hook)`：HTTP 结束 → 全部 Stop → 全部 Drain → OnShutdown。复用 lifecycle 配置及 ErrResourcesRetained 规则；只在用到后台组件时创建额外 Runtime。
|
||||
|
||||
生成项目在 `svc.Init` 创建 server，Context 设为 svc.WorkContext，并登记 `svc.OnStop("ws",server.Stop)` / `svc.OnDrain("ws",server.Drain)`。项目可手写以下方法，新版生成 NewApp 自动调用：
|
||||
|
||||
```go
|
||||
func (c *ServiceContext) RegisterCustomRoutes(app *omni.App) error {
|
||||
app.GET("/ws", wsomni.Handler(c.WebSocket, c.WebSocketHandlers))
|
||||
return nil
|
||||
}
|
||||
```
|
||||
|
||||
不要同时把同一组件登记到 App 和 svc；依赖归属与排空屏障必须一致。自定义路由在生成路由与内建端点之后注册，重复路径由路由器拒绝，不会覆盖现有接口。重新运行 `api generate` 更新入口，无需修改受管文件；普通项目没有此方法时行为不变。握手成功会记录 HTTP 101，HTTP 日志只计算握手时间。
|
||||
|
||||
ServerConfig/ClientConfig 直接组合进项目的 Config/Default/Validate；文件参数仍采用显式 `${NAME}` 引用。框架不新增 WS DSL 或自动业务 Token。
|
||||
|
||||
兼容范围：默认 nethttp 的 HTTP/1.1 Upgrade、标准 HTTP/1.1 反向代理、WS/WSS。自研 native 和独立实验 http2 Runner 目前不支持 WS；Mount 在配置时拒绝，普通 Handler 遇到不支持接管的 ResponseWriter 返回 501。不宣称支持 HTTP/2 extended CONNECT。切换 Runner 时仍需遵守此边界。
|
||||
|
||||
## 错误与观测
|
||||
|
||||
`errors.Is` 判断 ErrClosed、ErrNotConnected、ErrQueueFull、ErrMessageTooLarge、ErrInvalidMessage、ErrUnsupportedTransport、ErrConnectionLimit；`errors.As` 获取 HandshakeError.StatusCode、CloseError.Code、OpError.Op。Cause 保留诊断链，但可能包含敏感数据；默认错误文本不包含 URL、认证头、payload 或对端 reason，不要无条件递归记录 Cause。
|
||||
|
||||
Server.Stats 提供握手数、活跃连接、累计接受/拒绝；Client.Stats 提供活跃连接、握手、Run数、关闭状态及累计 DialAttempts、HandshakeFailures、Reconnects；Session.Stats 提供消息/字节和队列占用。`observe` 提供安全 logx 日志和跨会话统计快照，可由项目导出到 metrics/OTel。连接 ID 不适合作为指标标签；不要把整条长连接当成持续数小时的 HTTP Span。业务消息 tracing 由业务协议显式传播。
|
||||
633
ws/client.go
Normal file
633
ws/client.go
Normal file
@@ -0,0 +1,633 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"errors"
|
||||
"io"
|
||||
"math/rand"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
var ErrConnectionLimit = errors.New("ws: connection limit reached")
|
||||
|
||||
// ReconnectConfig controls Run only. Dial never retries. Zero MaxAttempts and
|
||||
// MaxElapsed mean unlimited retries/time when Enabled is explicitly true.
|
||||
// A stable connection resets both retry budgets; short flaps do not.
|
||||
type ReconnectConfig struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
MinBackoff time.Duration `json:"minBackoff"`
|
||||
MaxBackoff time.Duration `json:"maxBackoff"`
|
||||
StableAfter time.Duration `json:"stableAfter"`
|
||||
MaxAttempts int `json:"maxAttempts"`
|
||||
MaxElapsed time.Duration `json:"maxElapsed"`
|
||||
}
|
||||
|
||||
type ClientConfig struct {
|
||||
Session SessionConfig `json:"session"`
|
||||
HandshakeTimeout time.Duration `json:"handshakeTimeout"`
|
||||
MaxConnections int `json:"maxConnections"`
|
||||
Subprotocols []string `json:"subprotocols"`
|
||||
RequireSubprotocol bool `json:"requireSubprotocol"`
|
||||
EnableCompression bool `json:"enableCompression"`
|
||||
// ProxyMode is environment (default), direct, or url. ProxyURL may contain
|
||||
// credentials and must not be logged. Explicit proxy URLs ignore NO_PROXY.
|
||||
// Supported proxy schemes are http, socks5, and socks5h (remote DNS).
|
||||
ProxyMode string `json:"proxyMode"`
|
||||
ProxyURL string `json:"proxyUrl"`
|
||||
Reconnect ReconnectConfig `json:"reconnect"`
|
||||
Headers http.Header `json:"-"`
|
||||
// HeaderProvider is called before each handshake. Return a fresh header map;
|
||||
// the client snapshots it. Provider values replace static Headers values.
|
||||
// Implementations must honor context cancellation and support concurrent calls.
|
||||
HeaderProvider func(context.Context) (http.Header, error) `json:"-"`
|
||||
// Prepare creates per-session handlers after a successful handshake. When
|
||||
// set, its result replaces Dial/Run handlers. It must honor cancellation.
|
||||
Prepare func(context.Context) (Handlers, error) `json:"-"`
|
||||
TLSConfig *tls.Config `json:"-"`
|
||||
// ShouldReconnect overrides default error classification, not budgets or
|
||||
// explicit client/context shutdown. It must be safe for concurrent Run calls.
|
||||
ShouldReconnect func(error) bool `json:"-"`
|
||||
}
|
||||
|
||||
func (cfg ClientConfig) Validate() error { _, _, err := normalizeClientConfig(cfg); return err }
|
||||
|
||||
func normalizeClientConfig(cfg ClientConfig) (ClientConfig, func(*http.Request) (*url.URL, error), error) {
|
||||
var err error
|
||||
cfg.Session, err = normalizeSessionConfig(cfg.Session)
|
||||
if err != nil {
|
||||
return cfg, nil, err
|
||||
}
|
||||
if cfg.HandshakeTimeout < 0 || cfg.MaxConnections < 0 || cfg.MaxConnections > 65536 || cfg.Reconnect.MinBackoff < 0 || cfg.Reconnect.MaxBackoff < 0 || cfg.Reconnect.StableAfter < 0 || cfg.Reconnect.MaxElapsed < 0 || cfg.Reconnect.MaxAttempts < 0 {
|
||||
return cfg, nil, errors.New("ws: invalid client limits or durations")
|
||||
}
|
||||
if cfg.HandshakeTimeout == 0 {
|
||||
cfg.HandshakeTimeout = 10 * time.Second
|
||||
}
|
||||
if cfg.MaxConnections == 0 {
|
||||
cfg.MaxConnections = 64
|
||||
}
|
||||
if cfg.Reconnect.MinBackoff == 0 {
|
||||
cfg.Reconnect.MinBackoff = 100 * time.Millisecond
|
||||
}
|
||||
if cfg.Reconnect.MaxBackoff == 0 {
|
||||
cfg.Reconnect.MaxBackoff = 30 * time.Second
|
||||
}
|
||||
if cfg.Reconnect.StableAfter == 0 {
|
||||
cfg.Reconnect.StableAfter = 30 * time.Second
|
||||
}
|
||||
if cfg.Reconnect.MinBackoff > cfg.Reconnect.MaxBackoff {
|
||||
return cfg, nil, errors.New("ws: minBackoff exceeds maxBackoff")
|
||||
}
|
||||
cfg.Headers, err = clientHeaders(cfg.Headers)
|
||||
if err != nil {
|
||||
return cfg, nil, err
|
||||
}
|
||||
if len(cfg.Subprotocols) > 32 {
|
||||
return cfg, nil, errors.New("ws: too many subprotocols")
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
for _, protocol := range cfg.Subprotocols {
|
||||
if !headerToken(protocol) || len(protocol) > 256 || seen[protocol] {
|
||||
return cfg, nil, errors.New("ws: invalid or duplicate subprotocol")
|
||||
}
|
||||
seen[protocol] = true
|
||||
}
|
||||
if cfg.RequireSubprotocol && len(cfg.Subprotocols) == 0 {
|
||||
return cfg, nil, errors.New("ws: required subprotocol list is empty")
|
||||
}
|
||||
cfg.Subprotocols = append([]string(nil), cfg.Subprotocols...)
|
||||
if cfg.TLSConfig != nil {
|
||||
cfg.TLSConfig = cfg.TLSConfig.Clone()
|
||||
}
|
||||
if cfg.ProxyMode != "url" && cfg.ProxyURL != "" {
|
||||
return cfg, nil, errors.New("ws: proxyUrl requires url proxyMode")
|
||||
}
|
||||
var proxy func(*http.Request) (*url.URL, error)
|
||||
switch cfg.ProxyMode {
|
||||
case "", "environment":
|
||||
proxy = func(r *http.Request) (*url.URL, error) {
|
||||
u, err := http.ProxyFromEnvironment(r)
|
||||
if err != nil || u == nil {
|
||||
return u, err
|
||||
}
|
||||
return supportedProxy(u)
|
||||
}
|
||||
case "direct":
|
||||
case "url":
|
||||
u, e := url.Parse(cfg.ProxyURL)
|
||||
if e != nil || u == nil || u.Hostname() == "" || u.Opaque != "" || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || (u.Path != "" && u.Path != "/") || !validPort(u) {
|
||||
return cfg, nil, errors.New("ws: invalid proxy URL")
|
||||
}
|
||||
u, e = supportedProxy(u)
|
||||
if e != nil {
|
||||
return cfg, nil, e
|
||||
}
|
||||
proxy = http.ProxyURL(u)
|
||||
default:
|
||||
return cfg, nil, errors.New("ws: invalid proxyMode")
|
||||
}
|
||||
return cfg, proxy, nil
|
||||
}
|
||||
|
||||
func supportedProxy(u *url.URL) (*url.URL, error) {
|
||||
copyURL := *u
|
||||
switch copyURL.Scheme {
|
||||
case "http", "socks5":
|
||||
case "socks5h":
|
||||
copyURL.Scheme = "socks5"
|
||||
default:
|
||||
return nil, errors.New("ws: proxy scheme must be http, socks5 or socks5h")
|
||||
}
|
||||
return ©URL, nil
|
||||
}
|
||||
|
||||
func validPort(u *url.URL) bool {
|
||||
if strings.HasSuffix(u.Host, ":") {
|
||||
return false
|
||||
}
|
||||
if p := u.Port(); p != "" {
|
||||
n, err := strconv.Atoi(p)
|
||||
return err == nil && n > 0 && n <= 65535
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func headerToken(s string) bool {
|
||||
if s == "" {
|
||||
return false
|
||||
}
|
||||
for i := 0; i < len(s); i++ {
|
||||
c := s[i]
|
||||
if c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' || strings.ContainsRune("!#$%&'*+-.^_`|~", rune(c)) {
|
||||
continue
|
||||
}
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
func clientHeaders(input http.Header) (http.Header, error) {
|
||||
if len(input) > 128 {
|
||||
return nil, errors.New("ws: too many handshake headers")
|
||||
}
|
||||
out := make(http.Header, len(input))
|
||||
size := 0
|
||||
for key, values := range input {
|
||||
name := http.CanonicalHeaderKey(key)
|
||||
if !headerToken(key) || len(key) > 256 {
|
||||
return nil, errors.New("ws: invalid handshake header name")
|
||||
}
|
||||
switch name {
|
||||
case "Host", "Connection", "Upgrade", "Sec-Websocket-Key", "Sec-Websocket-Version", "Sec-Websocket-Extensions", "Sec-Websocket-Protocol", "Content-Length", "Transfer-Encoding", "Trailer":
|
||||
return nil, errors.New("ws: reserved handshake header")
|
||||
}
|
||||
if _, ok := out[name]; ok {
|
||||
return nil, errors.New("ws: duplicate handshake header")
|
||||
}
|
||||
for _, value := range values {
|
||||
size += len(key) + len(value)
|
||||
for i := 0; i < len(value); i++ {
|
||||
if value[i] < 32 && value[i] != '\t' || value[i] == 127 {
|
||||
return nil, errors.New("ws: invalid handshake header value")
|
||||
}
|
||||
}
|
||||
}
|
||||
if size > 64<<10 {
|
||||
return nil, errors.New("ws: handshake headers exceed limit")
|
||||
}
|
||||
out[name] = append([]string(nil), values...)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Client owns its sessions. Concurrent Dial and Run calls are independent and
|
||||
// share only limits and lifecycle. One Run owns one reconnecting connection;
|
||||
// every successful handshake produces a new Session with a new ID.
|
||||
type Client struct {
|
||||
config ClientConfig
|
||||
dialer websocket.Dialer
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
mu sync.Mutex
|
||||
closed bool
|
||||
active int
|
||||
connections int
|
||||
dialAttempts, handshakeFailures, reconnects uint64
|
||||
sessions map[string]*Session
|
||||
releases map[string]func()
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
func NewClient(cfg ClientConfig) (*Client, error) {
|
||||
cfg, proxy, err := normalizeClientConfig(cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
return &Client{config: cfg, dialer: websocket.Dialer{Proxy: proxy, HandshakeTimeout: cfg.HandshakeTimeout, TLSClientConfig: cfg.TLSConfig, Subprotocols: cfg.Subprotocols, EnableCompression: cfg.EnableCompression}, ctx: ctx, cancel: cancel, sessions: make(map[string]*Session), releases: make(map[string]func()), done: make(chan struct{})}, nil
|
||||
}
|
||||
|
||||
func (c *Client) begin(connection bool) error {
|
||||
for {
|
||||
c.mu.Lock()
|
||||
if c.closed {
|
||||
c.mu.Unlock()
|
||||
return ErrClosed
|
||||
}
|
||||
if connection && c.connections >= c.config.MaxConnections {
|
||||
// Done may wake a caller before its registry watcher gets CPU time.
|
||||
// Retire finished sessions synchronously before reporting overload.
|
||||
var release func()
|
||||
for id, s := range c.sessions {
|
||||
select {
|
||||
case <-s.Done():
|
||||
release = c.releases[id]
|
||||
default:
|
||||
}
|
||||
if release != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
c.mu.Unlock()
|
||||
if release != nil {
|
||||
release()
|
||||
continue
|
||||
}
|
||||
return ErrConnectionLimit
|
||||
}
|
||||
c.active++
|
||||
if connection {
|
||||
c.connections++
|
||||
c.dialAttempts++
|
||||
}
|
||||
c.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
}
|
||||
func (c *Client) end(connection bool, id string) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if id != "" {
|
||||
delete(c.sessions, id)
|
||||
delete(c.releases, id)
|
||||
}
|
||||
c.active--
|
||||
if connection {
|
||||
c.connections--
|
||||
}
|
||||
if c.closed && c.active == 0 {
|
||||
close(c.done)
|
||||
}
|
||||
}
|
||||
|
||||
// Dial performs exactly one handshake. Its timeout ends at the handshake;
|
||||
// ctx continues to own the resulting session's lifetime. No message is replayed.
|
||||
func (c *Client) Dial(ctx context.Context, endpoint string, handlers Handlers) (*Session, error) {
|
||||
return c.dial(ctx, endpoint, handlers, 0)
|
||||
}
|
||||
func (c *Client) dial(ctx context.Context, endpoint string, handlers Handlers, budget time.Duration) (*Session, error) {
|
||||
if ctx == nil {
|
||||
return nil, errors.New("ws: nil context")
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
u, err := url.Parse(endpoint)
|
||||
if err != nil || u == nil || u.Hostname() == "" || u.User != nil || u.Fragment != "" || u.Opaque != "" || (u.Scheme != "ws" && u.Scheme != "wss") || !validPort(u) {
|
||||
return nil, errors.New("ws: invalid endpoint")
|
||||
}
|
||||
if err = c.begin(true); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
transferred := false
|
||||
sessionCtx, cancelSession := context.WithCancel(ctx)
|
||||
stopSessionShutdown := context.AfterFunc(c.ctx, cancelSession)
|
||||
defer func() {
|
||||
if !transferred {
|
||||
c.mu.Lock()
|
||||
c.handshakeFailures++
|
||||
c.mu.Unlock()
|
||||
stopSessionShutdown()
|
||||
cancelSession()
|
||||
c.end(true, "")
|
||||
}
|
||||
}()
|
||||
handshakeTimeout := c.config.HandshakeTimeout
|
||||
if budget > 0 && budget < handshakeTimeout {
|
||||
handshakeTimeout = budget
|
||||
}
|
||||
hctx, cancel := context.WithTimeout(ctx, handshakeTimeout)
|
||||
stop := context.AfterFunc(c.ctx, cancel)
|
||||
defer func() { stop(); cancel() }()
|
||||
headers := c.config.Headers.Clone()
|
||||
if c.config.HeaderProvider != nil {
|
||||
provided, e := provideHeaders(c.config.HeaderProvider, hctx)
|
||||
if e != nil {
|
||||
return nil, &OpError{Op: "headers", Cause: e}
|
||||
}
|
||||
provided, e = clientHeaders(provided)
|
||||
if e != nil {
|
||||
return nil, &OpError{Op: "headers", Cause: e}
|
||||
}
|
||||
for k, v := range provided {
|
||||
headers[k] = v
|
||||
}
|
||||
}
|
||||
dialer := c.dialer
|
||||
// Gorilla applies the context deadline to handshake I/O, but canceling a
|
||||
// deadline-less context after TCP connect does not interrupt its HTTP read.
|
||||
// Retain the raw socket only through the handshake so Close/Stop cancels
|
||||
// stalled HTTP and proxy handshakes immediately, without timing out sessions.
|
||||
var stopHandshake func() bool
|
||||
dialer.NetDialContext = func(dialctx context.Context, network, address string) (net.Conn, error) {
|
||||
conn, err := (&net.Dialer{Timeout: c.config.HandshakeTimeout, KeepAlive: 30 * time.Second}).DialContext(dialctx, network, address)
|
||||
if err == nil {
|
||||
stopHandshake = context.AfterFunc(hctx, func() { _ = conn.Close() })
|
||||
}
|
||||
return conn, err
|
||||
}
|
||||
if dialer.TLSClientConfig != nil {
|
||||
dialer.TLSClientConfig = dialer.TLSClientConfig.Clone()
|
||||
}
|
||||
conn, response, err := dialer.DialContext(hctx, endpoint, headers)
|
||||
if stopHandshake != nil {
|
||||
stopHandshake()
|
||||
}
|
||||
if err == nil && hctx.Err() != nil {
|
||||
_ = conn.Close()
|
||||
err = hctx.Err()
|
||||
}
|
||||
if err != nil {
|
||||
status := 0
|
||||
if response != nil {
|
||||
status = response.StatusCode
|
||||
if response.Body != nil {
|
||||
_ = response.Body.Close()
|
||||
}
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
if c.ctx.Err() != nil {
|
||||
return nil, ErrClosed
|
||||
}
|
||||
if hctx.Err() != nil {
|
||||
err = hctx.Err()
|
||||
} else if deadline, ok := hctx.Deadline(); ok && !time.Now().Before(deadline) {
|
||||
var timed net.Error
|
||||
if errors.As(err, &timed) && timed.Timeout() {
|
||||
err = errors.Join(context.DeadlineExceeded, err)
|
||||
}
|
||||
}
|
||||
return nil, &HandshakeError{StatusCode: status, Cause: err}
|
||||
}
|
||||
if c.config.RequireSubprotocol && conn.Subprotocol() == "" {
|
||||
_ = conn.Close()
|
||||
return nil, &HandshakeError{StatusCode: http.StatusSwitchingProtocols, Cause: websocket.ErrBadHandshake}
|
||||
}
|
||||
if c.config.Prepare != nil {
|
||||
handlers, err = prepareClient(c.config.Prepare, sessionCtx)
|
||||
if err != nil {
|
||||
_ = conn.Close()
|
||||
return nil, &OpError{Op: "prepare", Cause: err}
|
||||
}
|
||||
}
|
||||
c.mu.Lock()
|
||||
if c.closed {
|
||||
c.mu.Unlock()
|
||||
_ = conn.Close()
|
||||
return nil, ErrClosed
|
||||
}
|
||||
s := newSession(conn, c.config.Session, handlers, sessionCtx)
|
||||
c.sessions[s.ID()] = s
|
||||
var releaseOnce sync.Once
|
||||
release := func() { releaseOnce.Do(func() { cancelSession(); c.end(true, s.ID()) }) }
|
||||
c.releases[s.ID()] = release
|
||||
transferred = true
|
||||
stopSessionShutdown()
|
||||
c.mu.Unlock()
|
||||
s.start()
|
||||
go func() { <-s.Done(); release() }()
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func provideHeaders(fn func(context.Context) (http.Header, error), ctx context.Context) (headers http.Header, err error) {
|
||||
defer func() {
|
||||
if recover() != nil {
|
||||
err = errors.New("ws: header provider panicked")
|
||||
}
|
||||
}()
|
||||
return fn(ctx)
|
||||
}
|
||||
func prepareClient(fn func(context.Context) (Handlers, error), ctx context.Context) (handlers Handlers, err error) {
|
||||
defer func() {
|
||||
if recover() != nil {
|
||||
err = errors.New("ws: prepare callback panicked")
|
||||
}
|
||||
}()
|
||||
return fn(ctx)
|
||||
}
|
||||
|
||||
// Run blocks until canceled, closed, or retry policy/budget stops it. Reconnect
|
||||
// must be enabled explicitly. Handlers are reused; OnOpen must restore any
|
||||
// application subscriptions. Credentials are refreshed by HeaderProvider.
|
||||
func (c *Client) Run(ctx context.Context, endpoint string, handlers Handlers) error {
|
||||
if ctx == nil {
|
||||
return errors.New("ws: nil context")
|
||||
}
|
||||
if err := c.begin(false); err != nil {
|
||||
return err
|
||||
}
|
||||
defer c.end(false, "")
|
||||
runctx, cancel := context.WithCancel(ctx)
|
||||
stop := context.AfterFunc(c.ctx, cancel)
|
||||
defer func() { stop(); cancel() }()
|
||||
attempts := 0
|
||||
episode := time.Now()
|
||||
var previousError error = context.DeadlineExceeded
|
||||
repeat := false
|
||||
for {
|
||||
var remaining time.Duration
|
||||
if c.config.Reconnect.Enabled && c.config.Reconnect.MaxElapsed > 0 {
|
||||
remaining = c.config.Reconnect.MaxElapsed - time.Since(episode)
|
||||
if remaining <= 0 {
|
||||
return previousError
|
||||
}
|
||||
}
|
||||
started := time.Now()
|
||||
if repeat {
|
||||
c.mu.Lock()
|
||||
c.reconnects++
|
||||
c.mu.Unlock()
|
||||
}
|
||||
repeat = true
|
||||
s, err := c.dial(runctx, endpoint, handlers, remaining)
|
||||
if err == nil {
|
||||
started = time.Now()
|
||||
err = s.Wait(runctx)
|
||||
if time.Since(started) >= c.config.Reconnect.StableAfter {
|
||||
attempts = 0
|
||||
episode = time.Now()
|
||||
}
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if c.ctx.Err() != nil {
|
||||
return nil
|
||||
}
|
||||
if !c.config.Reconnect.Enabled || !c.retryable(err) {
|
||||
return err
|
||||
}
|
||||
previousError = err
|
||||
rc := c.config.Reconnect
|
||||
if rc.MaxAttempts > 0 && attempts >= rc.MaxAttempts {
|
||||
return err
|
||||
}
|
||||
delay := retryBackoff(rc, attempts)
|
||||
if rc.MaxElapsed > 0 && time.Since(episode)+delay >= rc.MaxElapsed {
|
||||
return err
|
||||
}
|
||||
attempts++
|
||||
timer := time.NewTimer(delay)
|
||||
select {
|
||||
case <-runctx.Done():
|
||||
timer.Stop()
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
return nil
|
||||
case <-timer.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func retryBackoff(cfg ReconnectConfig, attempt int) time.Duration {
|
||||
d := cfg.MinBackoff
|
||||
for i := 0; i < attempt && d < cfg.MaxBackoff; i++ {
|
||||
if d > cfg.MaxBackoff/2 {
|
||||
d = cfg.MaxBackoff
|
||||
} else {
|
||||
d *= 2
|
||||
}
|
||||
}
|
||||
if d > 1 {
|
||||
d = d/2 + time.Duration(rand.Int63n(int64(d-d/2))+1)
|
||||
}
|
||||
return d
|
||||
}
|
||||
func (c *Client) retryable(err error) bool {
|
||||
if err == nil || errors.Is(err, ErrClosed) || errors.Is(err, context.Canceled) {
|
||||
return false
|
||||
}
|
||||
if c.config.ShouldReconnect != nil {
|
||||
return c.config.ShouldReconnect(err)
|
||||
}
|
||||
var handshake *HandshakeError
|
||||
if errors.As(err, &handshake) {
|
||||
switch handshake.StatusCode {
|
||||
case 429, 500, 502, 503, 504:
|
||||
return true
|
||||
case 0:
|
||||
default:
|
||||
return false
|
||||
}
|
||||
if errors.Is(err, websocket.ErrBadHandshake) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
var unknown x509.UnknownAuthorityError
|
||||
var invalid x509.CertificateInvalidError
|
||||
var hostname x509.HostnameError
|
||||
if errors.As(err, &unknown) || errors.As(err, &invalid) || errors.As(err, &hostname) {
|
||||
return false
|
||||
}
|
||||
var closeError *CloseError
|
||||
if errors.As(err, &closeError) {
|
||||
switch closeError.Code {
|
||||
case 1001, 1006, 1011, 1012, 1013:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
var network net.Error
|
||||
if errors.As(err, &network) || errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
|
||||
return true
|
||||
}
|
||||
return handshake != nil || errors.Is(err, ErrNotConnected)
|
||||
}
|
||||
|
||||
// Close is idempotent and nonblocking: it rejects new work, cancels handshakes
|
||||
// and reconnect waits, and requests session close. Drain waits for all work,
|
||||
// including user callbacks. A callback that ignores cancellation cannot be killed.
|
||||
func (c *Client) Close() error {
|
||||
c.mu.Lock()
|
||||
if c.closed {
|
||||
c.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
c.closed = true
|
||||
sessions := make([]*Session, 0, len(c.sessions))
|
||||
for _, s := range c.sessions {
|
||||
sessions = append(sessions, s)
|
||||
}
|
||||
if c.active == 0 {
|
||||
close(c.done)
|
||||
}
|
||||
c.mu.Unlock()
|
||||
c.cancel()
|
||||
for _, s := range sessions {
|
||||
_ = s.Close(1001, "client shutdown")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (c *Client) Stop(context.Context) error { return c.Close() }
|
||||
func (c *Client) Drain(ctx context.Context) error {
|
||||
if ctx == nil {
|
||||
return errors.New("ws: nil context")
|
||||
}
|
||||
select {
|
||||
case <-c.done:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
c.mu.Lock()
|
||||
sessions := make([]*Session, 0, len(c.sessions))
|
||||
for _, s := range c.sessions {
|
||||
sessions = append(sessions, s)
|
||||
}
|
||||
c.mu.Unlock()
|
||||
for _, s := range sessions {
|
||||
s.Abort()
|
||||
}
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
func (c *Client) Shutdown(ctx context.Context) error { _ = c.Close(); return c.Drain(ctx) }
|
||||
|
||||
type ClientStats struct {
|
||||
ActiveConnections int
|
||||
Handshakes int
|
||||
Running int
|
||||
Closed bool
|
||||
// DialAttempts includes admitted attempts canceled before the HTTP handshake.
|
||||
DialAttempts uint64
|
||||
HandshakeFailures uint64
|
||||
// Reconnects counts subsequent Dial attempts made by Run, successful or not.
|
||||
Reconnects uint64
|
||||
}
|
||||
|
||||
func (c *Client) Stats() ClientStats {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return ClientStats{ActiveConnections: len(c.sessions), Handshakes: c.connections - len(c.sessions), Running: c.active - c.connections, Closed: c.closed, DialAttempts: c.dialAttempts, HandshakeFailures: c.handshakeFailures, Reconnects: c.reconnects}
|
||||
}
|
||||
494
ws/client_test.go
Normal file
494
ws/client_test.go
Normal file
@@ -0,0 +1,494 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
func clientTestConfig() ClientConfig {
|
||||
return ClientConfig{ProxyMode: "direct", Session: SessionConfig{CloseTimeout: 100 * time.Millisecond}}
|
||||
}
|
||||
func clientForTest(t *testing.T, cfg ClientConfig) *Client {
|
||||
t.Helper()
|
||||
c, err := NewClient(cfg)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
if err := c.Shutdown(ctx); err != nil {
|
||||
t.Errorf("client shutdown: %v", err)
|
||||
}
|
||||
})
|
||||
return c
|
||||
}
|
||||
func echoWebSocket(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
for {
|
||||
kind, data, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if err = conn.WriteMessage(kind, data); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
func wsURL(s *httptest.Server) string { return "ws" + strings.TrimPrefix(s.URL, "http") }
|
||||
func receiveClientMessage(t *testing.T, ch <-chan Message) Message {
|
||||
t.Helper()
|
||||
select {
|
||||
case m := <-ch:
|
||||
return m
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("message timeout")
|
||||
return Message{}
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientHandshakeDeadlineDoesNotExpireSession(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(echoWebSocket))
|
||||
defer server.Close()
|
||||
cfg := clientTestConfig()
|
||||
cfg.HandshakeTimeout = 40 * time.Millisecond
|
||||
c := clientForTest(t, cfg)
|
||||
messages := make(chan Message, 2)
|
||||
s, err := c.Dial(context.Background(), wsURL(server), Handlers{OnMessage: func(_ context.Context, _ *Session, m Message) error { messages <- m; return nil }})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
timer := time.NewTimer(100 * time.Millisecond)
|
||||
<-timer.C
|
||||
for _, kind := range []MessageType{Text, Binary} {
|
||||
if err = s.Send(context.Background(), Message{Type: kind, Data: []byte("alive")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := receiveClientMessage(t, messages)
|
||||
if got.Type != kind || string(got.Data) != "alive" {
|
||||
t.Fatalf("unexpected echo: %#v", got)
|
||||
}
|
||||
}
|
||||
_ = s.Close(1000, "")
|
||||
}
|
||||
|
||||
func TestClientTLSAndMutualTLS(t *testing.T) {
|
||||
for _, mutual := range []bool{false, true} {
|
||||
t.Run(fmt.Sprint(mutual), func(t *testing.T) {
|
||||
server := httptest.NewUnstartedServer(http.HandlerFunc(echoWebSocket))
|
||||
if mutual {
|
||||
server.TLS = &tls.Config{ClientAuth: tls.RequireAnyClientCert}
|
||||
}
|
||||
server.StartTLS()
|
||||
defer server.Close()
|
||||
roots := x509.NewCertPool()
|
||||
roots.AddCert(server.Certificate())
|
||||
cfg := clientTestConfig()
|
||||
cfg.TLSConfig = &tls.Config{RootCAs: roots, MinVersion: tls.VersionTLS12}
|
||||
if mutual {
|
||||
cfg.TLSConfig.Certificates = append([]tls.Certificate(nil), server.TLS.Certificates...)
|
||||
}
|
||||
c := clientForTest(t, cfg)
|
||||
cfg.TLSConfig.ServerName = "wrong-after-construction.invalid"
|
||||
messages := make(chan Message, 1)
|
||||
s, err := c.Dial(context.Background(), wsURL(server), Handlers{OnMessage: func(_ context.Context, _ *Session, m Message) error { messages <- m; return nil }})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.Send(context.Background(), Message{Type: Text, Data: []byte("tls")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(receiveClientMessage(t, messages).Data) != "tls" {
|
||||
t.Fatal("bad echo")
|
||||
}
|
||||
_ = s.Close(1000, "")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientHeadersRefreshReconnectAndFreshSessions(t *testing.T) {
|
||||
var accepted atomic.Int32
|
||||
firstOpen := make(chan struct{})
|
||||
var firstOpenOnce sync.Once
|
||||
headerValues := make(chan string, 4)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
headerValues <- r.Header.Get("Authorization")
|
||||
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
if accepted.Add(1) == 1 {
|
||||
select {
|
||||
case <-firstOpen:
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(1012, "restart"), time.Now().Add(time.Second))
|
||||
return
|
||||
}
|
||||
for {
|
||||
if _, _, err = conn.ReadMessage(); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
var providerCalls atomic.Int32
|
||||
var prepareCalls atomic.Int32
|
||||
ids := make(chan string, 4)
|
||||
cfg := clientTestConfig()
|
||||
cfg.Reconnect = ReconnectConfig{Enabled: true, MinBackoff: time.Millisecond, MaxBackoff: 5 * time.Millisecond, StableAfter: time.Second, MaxAttempts: 3}
|
||||
cfg.HeaderProvider = func(context.Context) (http.Header, error) {
|
||||
n := providerCalls.Add(1)
|
||||
return http.Header{"Authorization": {fmt.Sprint(n)}}, nil
|
||||
}
|
||||
cfg.Prepare = func(ctx context.Context) (Handlers, error) {
|
||||
prepareCalls.Add(1)
|
||||
return Handlers{OnOpen: func(sessionCtx context.Context, s *Session) error {
|
||||
if ctx.Err() != nil || sessionCtx.Err() != nil {
|
||||
return errors.New("prepare context canceled early")
|
||||
}
|
||||
ids <- s.ID()
|
||||
firstOpenOnce.Do(func() { close(firstOpen) })
|
||||
return nil
|
||||
}}, nil
|
||||
}
|
||||
c := clientForTest(t, cfg)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
finished := make(chan error, 1)
|
||||
go func() { finished <- c.Run(ctx, wsURL(server), Handlers{}) }()
|
||||
var first, second string
|
||||
select {
|
||||
case first = <-ids:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("first open timeout")
|
||||
}
|
||||
select {
|
||||
case second = <-ids:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("reopen timeout")
|
||||
}
|
||||
if first == second {
|
||||
t.Fatal("reused physical session ID")
|
||||
}
|
||||
if <-headerValues != "1" || <-headerValues != "2" {
|
||||
t.Fatal("headers not refreshed")
|
||||
}
|
||||
if prepareCalls.Load() != 2 {
|
||||
t.Fatal("prepare not recreated")
|
||||
}
|
||||
cancel()
|
||||
select {
|
||||
case err := <-finished:
|
||||
if !errors.Is(err, context.Canceled) {
|
||||
t.Fatalf("run: %v", err)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("run cancel timeout")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientDoesNotReconnectAuthenticationOrProtocolFailure(t *testing.T) {
|
||||
for _, status := range []int{400, 401, 403, 404} {
|
||||
t.Run(fmt.Sprint(status), func(t *testing.T) {
|
||||
var calls atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls.Add(1)
|
||||
w.WriteHeader(status)
|
||||
_, _ = w.Write([]byte("secret response"))
|
||||
}))
|
||||
defer server.Close()
|
||||
cfg := clientTestConfig()
|
||||
cfg.Reconnect.Enabled = true
|
||||
c := clientForTest(t, cfg)
|
||||
err := c.Run(context.Background(), wsURL(server)+"/?token=TOPSECRET", Handlers{})
|
||||
var handshake *HandshakeError
|
||||
if !errors.As(err, &handshake) || handshake.StatusCode != status {
|
||||
t.Fatalf("error: %v", err)
|
||||
}
|
||||
if calls.Load() != 1 || strings.Contains(err.Error(), "SECRET") || strings.Contains(err.Error(), "secret") {
|
||||
t.Fatalf("retry/leak: calls=%d error=%v", calls.Load(), err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientCloseCancelsHandshakeAndBackoff(t *testing.T) {
|
||||
t.Run("handshake", func(t *testing.T) {
|
||||
entered := make(chan struct{})
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { close(entered); <-r.Context().Done() }))
|
||||
defer server.Close()
|
||||
c := clientForTest(t, clientTestConfig())
|
||||
finished := make(chan error, 1)
|
||||
go func() { _, err := c.Dial(context.Background(), wsURL(server), Handlers{}); finished <- err }()
|
||||
<-entered
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := c.Shutdown(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := <-finished; !errors.Is(err, ErrClosed) {
|
||||
t.Fatalf("dial: %v", err)
|
||||
}
|
||||
})
|
||||
t.Run("backoff", func(t *testing.T) {
|
||||
entered := make(chan struct{})
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { close(entered); w.WriteHeader(503) }))
|
||||
defer server.Close()
|
||||
cfg := clientTestConfig()
|
||||
cfg.Reconnect = ReconnectConfig{Enabled: true, MinBackoff: time.Hour, MaxBackoff: time.Hour}
|
||||
c := clientForTest(t, cfg)
|
||||
finished := make(chan error, 1)
|
||||
go func() { finished <- c.Run(context.Background(), wsURL(server), Handlers{}) }()
|
||||
<-entered
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := c.Shutdown(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := <-finished; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestClientConcurrentDialLimitAndClose(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(echoWebSocket))
|
||||
defer server.Close()
|
||||
cfg := clientTestConfig()
|
||||
cfg.MaxConnections = 2
|
||||
c := clientForTest(t, cfg)
|
||||
var wg sync.WaitGroup
|
||||
var successes atomic.Int32
|
||||
for i := 0; i < 20; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_, err := c.Dial(context.Background(), wsURL(server), Handlers{})
|
||||
if err == nil {
|
||||
successes.Add(1)
|
||||
} else if !errors.Is(err, ErrConnectionLimit) {
|
||||
t.Errorf("dial: %v", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
if successes.Load() != 2 {
|
||||
t.Fatalf("connections=%d", successes.Load())
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := c.Shutdown(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := c.Dial(ctx, wsURL(server), Handlers{}); !errors.Is(err, ErrClosed) {
|
||||
t.Fatalf("closed dial: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientRetryBudgetAndValidation(t *testing.T) {
|
||||
var calls atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1); w.WriteHeader(503) }))
|
||||
defer server.Close()
|
||||
cfg := clientTestConfig()
|
||||
cfg.Reconnect = ReconnectConfig{Enabled: true, MinBackoff: time.Millisecond, MaxBackoff: 2 * time.Millisecond, MaxAttempts: 2}
|
||||
c := clientForTest(t, cfg)
|
||||
if err := c.Run(context.Background(), wsURL(server), Handlers{}); err == nil {
|
||||
t.Fatal("wanted failure")
|
||||
}
|
||||
if calls.Load() != 3 {
|
||||
t.Fatalf("retry budget: %d", calls.Load())
|
||||
}
|
||||
for _, endpoint := range []string{"https://example.test", "ws://user:password@example.test/", "ws://", "ws://example.test/#fragment"} {
|
||||
if _, err := c.Dial(context.Background(), endpoint, Handlers{}); err == nil || strings.Contains(err.Error(), "password") {
|
||||
t.Fatalf("invalid endpoint: %v", err)
|
||||
}
|
||||
}
|
||||
for _, cfg := range []ClientConfig{{HandshakeTimeout: -1}, {ProxyMode: "bad"}, {ProxyMode: "direct", ProxyURL: "http://proxy"}, {Headers: http.Header{"Authorization": {"bad\r\nheader"}}}, {Headers: http.Header{"Sec-WebSocket-Key": {"fake"}}}, {RequireSubprotocol: true}, {Subprotocols: []string{"bad protocol"}}} {
|
||||
if _, err := NewClient(cfg); err == nil {
|
||||
t.Fatal("accepted invalid configuration")
|
||||
}
|
||||
}
|
||||
for i := 0; i < 100; i++ {
|
||||
d := retryBackoff(ReconnectConfig{MinBackoff: time.Second, MaxBackoff: 30 * time.Second}, i)
|
||||
if d <= 0 || d > 30*time.Second {
|
||||
t.Fatalf("unbounded delay %v", d)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientElapsedBudgetInterruptsStalledHandshake(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { <-r.Context().Done() }))
|
||||
defer server.Close()
|
||||
cfg := clientTestConfig()
|
||||
cfg.Reconnect = ReconnectConfig{Enabled: true, MaxElapsed: 50 * time.Millisecond}
|
||||
c := clientForTest(t, cfg)
|
||||
started := time.Now()
|
||||
err := c.Run(context.Background(), wsURL(server), Handlers{})
|
||||
if !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("budget error: %v", err)
|
||||
}
|
||||
if time.Since(started) > time.Second {
|
||||
t.Fatal("handshake exceeded reconnect budget")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientSingleSlotRapidReconnect(t *testing.T) {
|
||||
var calls atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
code := 1012
|
||||
if calls.Add(1) >= 20 {
|
||||
code = 1000
|
||||
}
|
||||
_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(code, ""), time.Now().Add(time.Second))
|
||||
_, _, _ = conn.ReadMessage()
|
||||
}))
|
||||
defer server.Close()
|
||||
cfg := clientTestConfig()
|
||||
cfg.MaxConnections = 1
|
||||
cfg.Reconnect = ReconnectConfig{Enabled: true, MinBackoff: time.Nanosecond, MaxBackoff: time.Nanosecond, MaxAttempts: 25}
|
||||
c := clientForTest(t, cfg)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
defer cancel()
|
||||
if err := c.Run(ctx, wsURL(server), Handlers{}); err != nil {
|
||||
t.Fatalf("rapid reconnect: %v", err)
|
||||
}
|
||||
if calls.Load() != 20 {
|
||||
t.Fatalf("connections: %d", calls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientSubprotocol(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := (&websocket.Upgrader{Subprotocols: []string{"omni.v1"}}).Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
for {
|
||||
if _, _, err = conn.ReadMessage(); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
cfg := clientTestConfig()
|
||||
cfg.Subprotocols = []string{"omni.v1"}
|
||||
cfg.RequireSubprotocol = true
|
||||
c := clientForTest(t, cfg)
|
||||
s, err := c.Dial(context.Background(), wsURL(server), Handlers{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if s.Subprotocol() != "omni.v1" {
|
||||
t.Fatal("wrong protocol")
|
||||
}
|
||||
_ = s.Close(1000, "")
|
||||
cfg.Subprotocols = []string{"missing"}
|
||||
other := clientForTest(t, cfg)
|
||||
if _, err = other.Dial(context.Background(), wsURL(server), Handlers{}); err == nil {
|
||||
t.Fatal("required protocol not enforced")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientHTTPProxyAuthenticationAndTunnel(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Header.Get("Proxy-Authorization") != "" {
|
||||
t.Error("proxy credential leaked to target")
|
||||
}
|
||||
echoWebSocket(w, r)
|
||||
}))
|
||||
defer server.Close()
|
||||
var connected atomic.Int32
|
||||
proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodConnect || r.Header.Get("Proxy-Authorization") != "Basic dXNlcjpwYXNz" {
|
||||
w.WriteHeader(http.StatusProxyAuthRequired)
|
||||
return
|
||||
}
|
||||
upstream, err := net.Dial("tcp", r.Host)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
w.WriteHeader(502)
|
||||
return
|
||||
}
|
||||
conn, rw, err := w.(http.Hijacker).Hijack()
|
||||
if err != nil {
|
||||
upstream.Close()
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
connected.Add(1)
|
||||
_, _ = conn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))
|
||||
go func() { defer conn.Close(); defer upstream.Close(); _, _ = io.Copy(upstream, rw) }()
|
||||
go func() { defer conn.Close(); defer upstream.Close(); _, _ = io.Copy(conn, upstream) }()
|
||||
}))
|
||||
defer proxy.Close()
|
||||
cfg := clientTestConfig()
|
||||
cfg.ProxyMode = "url"
|
||||
cfg.ProxyURL = strings.Replace(proxy.URL, "http://", "http://user:pass@", 1)
|
||||
c := clientForTest(t, cfg)
|
||||
messages := make(chan Message, 1)
|
||||
s, err := c.Dial(context.Background(), wsURL(server), Handlers{OnMessage: func(_ context.Context, _ *Session, m Message) error { messages <- m; return nil }})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.Send(context.Background(), Message{Type: Text, Data: []byte("proxied")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(receiveClientMessage(t, messages).Data) != "proxied" || connected.Load() != 1 {
|
||||
t.Fatal("proxy not used")
|
||||
}
|
||||
_ = s.Close(1000, "")
|
||||
}
|
||||
|
||||
func BenchmarkClientRoundTrip(b *testing.B) {
|
||||
server := httptest.NewServer(http.HandlerFunc(echoWebSocket))
|
||||
defer server.Close()
|
||||
c, err := NewClient(clientTestConfig())
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
defer c.Shutdown(context.Background())
|
||||
received := make(chan struct{}, 1)
|
||||
s, err := c.Dial(context.Background(), wsURL(server), Handlers{OnMessage: func(context.Context, *Session, Message) error { received <- struct{}{}; return nil }})
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
message := Message{Type: Text, Data: []byte("benchmark")}
|
||||
b.ReportAllocs()
|
||||
b.SetBytes(int64(len(message.Data)))
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if err := s.Send(context.Background(), message); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
<-received
|
||||
}
|
||||
b.StopTimer()
|
||||
_ = s.Close(1000, "")
|
||||
}
|
||||
83
ws/config.go
Normal file
83
ws/config.go
Normal file
@@ -0,0 +1,83 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SessionConfig bounds connection memory and I/O. Zero fields use defaults;
|
||||
// DisableHeartbeat explicitly disables automatic ping/pong deadlines.
|
||||
type SessionConfig struct {
|
||||
MaxMessageBytes int64 `json:"maxMessageBytes"`
|
||||
ReadQueueSize int `json:"readQueueSize"`
|
||||
ReadQueueBytes int64 `json:"readQueueBytes"`
|
||||
WriteQueueSize int `json:"writeQueueSize"`
|
||||
WriteQueueBytes int64 `json:"writeQueueBytes"`
|
||||
WriteTimeout time.Duration `json:"writeTimeout"`
|
||||
PingInterval time.Duration `json:"pingInterval"`
|
||||
PongTimeout time.Duration `json:"pongTimeout"`
|
||||
CloseTimeout time.Duration `json:"closeTimeout"`
|
||||
DisableHeartbeat bool `json:"disableHeartbeat"`
|
||||
}
|
||||
|
||||
func DefaultSessionConfig() SessionConfig {
|
||||
return SessionConfig{MaxMessageBytes: 1 << 20, ReadQueueSize: 64, ReadQueueBytes: 4 << 20,
|
||||
WriteQueueSize: 64, WriteQueueBytes: 4 << 20, WriteTimeout: 10 * time.Second,
|
||||
PingInterval: 30 * time.Second, PongTimeout: 60 * time.Second, CloseTimeout: 5 * time.Second}
|
||||
}
|
||||
|
||||
func (c SessionConfig) defaults() SessionConfig {
|
||||
d := DefaultSessionConfig()
|
||||
if c.MaxMessageBytes == 0 {
|
||||
c.MaxMessageBytes = d.MaxMessageBytes
|
||||
}
|
||||
if c.ReadQueueSize == 0 {
|
||||
c.ReadQueueSize = d.ReadQueueSize
|
||||
}
|
||||
if c.ReadQueueBytes == 0 {
|
||||
c.ReadQueueBytes = d.ReadQueueBytes
|
||||
}
|
||||
if c.WriteQueueSize == 0 {
|
||||
c.WriteQueueSize = d.WriteQueueSize
|
||||
}
|
||||
if c.WriteQueueBytes == 0 {
|
||||
c.WriteQueueBytes = d.WriteQueueBytes
|
||||
}
|
||||
if c.WriteTimeout == 0 {
|
||||
c.WriteTimeout = d.WriteTimeout
|
||||
}
|
||||
if c.PingInterval == 0 {
|
||||
c.PingInterval = d.PingInterval
|
||||
}
|
||||
if c.PongTimeout == 0 {
|
||||
c.PongTimeout = d.PongTimeout
|
||||
}
|
||||
if c.CloseTimeout == 0 {
|
||||
c.CloseTimeout = d.CloseTimeout
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
func (c SessionConfig) Validate() error {
|
||||
if c.MaxMessageBytes < 0 || c.MaxMessageBytes > 1<<30 || c.ReadQueueSize < 0 || c.ReadQueueSize > 65536 || c.ReadQueueBytes < 0 || c.ReadQueueBytes > 1<<30 || c.WriteQueueSize < 0 || c.WriteQueueSize > 65536 || c.WriteQueueBytes < 0 || c.WriteQueueBytes > 1<<30 {
|
||||
return fmt.Errorf("ws: message and queue limits must be nonnegative and bounded")
|
||||
}
|
||||
if c.WriteTimeout < 0 || c.PingInterval < 0 || c.PongTimeout < 0 || c.CloseTimeout < 0 {
|
||||
return fmt.Errorf("ws: timeouts must be nonnegative")
|
||||
}
|
||||
if c.WriteTimeout > 24*time.Hour || c.PingInterval > 24*time.Hour || c.PongTimeout > 24*time.Hour || c.CloseTimeout > 24*time.Hour {
|
||||
return fmt.Errorf("ws: timeout exceeds 24 hours")
|
||||
}
|
||||
c = c.defaults()
|
||||
if !c.DisableHeartbeat && c.PongTimeout <= c.PingInterval {
|
||||
return fmt.Errorf("ws: pongTimeout must exceed pingInterval")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func normalizeSessionConfig(c SessionConfig) (SessionConfig, error) {
|
||||
if err := c.Validate(); err != nil {
|
||||
return SessionConfig{}, err
|
||||
}
|
||||
return c.defaults(), nil
|
||||
}
|
||||
48
ws/errors.go
Normal file
48
ws/errors.go
Normal file
@@ -0,0 +1,48 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrClosed = errors.New("ws: session closed")
|
||||
ErrNotConnected = errors.New("ws: not connected")
|
||||
ErrQueueFull = errors.New("ws: queue full")
|
||||
ErrMessageTooLarge = errors.New("ws: message too large")
|
||||
ErrUnsupportedTransport = errors.New("ws: transport does not support websocket upgrade")
|
||||
ErrServerClosed = errors.New("ws: server closed")
|
||||
ErrInvalidMessage = errors.New("ws: invalid message")
|
||||
)
|
||||
|
||||
// CloseError deliberately excludes the peer-provided reason, which is untrusted
|
||||
// and may contain credentials. Cause remains available through errors.Unwrap.
|
||||
type CloseError struct {
|
||||
Code int
|
||||
Cause error
|
||||
}
|
||||
|
||||
func (e *CloseError) Error() string { return fmt.Sprintf("ws: connection closed (code %d)", e.Code) }
|
||||
func (e *CloseError) Unwrap() error { return e.Cause }
|
||||
|
||||
// HandshakeError exposes the HTTP status without putting a URL or response body
|
||||
// in the error text. Inspect Cause explicitly only when it is safe to do so.
|
||||
type HandshakeError struct {
|
||||
StatusCode int
|
||||
Cause error
|
||||
}
|
||||
|
||||
func (e *HandshakeError) Error() string {
|
||||
return fmt.Sprintf("ws: handshake failed (HTTP %d)", e.StatusCode)
|
||||
}
|
||||
func (e *HandshakeError) Unwrap() error { return e.Cause }
|
||||
|
||||
// OpError preserves the underlying cause without including its potentially
|
||||
// sensitive text in logs. Op is a fixed operation name, never a URL or payload.
|
||||
type OpError struct {
|
||||
Op string
|
||||
Cause error
|
||||
}
|
||||
|
||||
func (e *OpError) Error() string { return "ws: " + e.Op + " failed" }
|
||||
func (e *OpError) Unwrap() error { return e.Cause }
|
||||
319
ws/hub/hub.go
Normal file
319
ws/hub/hub.go
Normal file
@@ -0,0 +1,319 @@
|
||||
// Package hub provides a bounded, process-local WebSocket registry. Connection
|
||||
// and group IDs are opaque application keys; callers authorize membership.
|
||||
// It does not persist messages or provide cross-process delivery.
|
||||
package hub
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"git.yuntuo.hk/k2safe/go-omni/ws"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrClosed = errors.New("ws/hub: closed")
|
||||
ErrNotFound = errors.New("ws/hub: connection not found")
|
||||
ErrDuplicate = errors.New("ws/hub: duplicate connection ID")
|
||||
ErrLimit = errors.New("ws/hub: resource limit reached")
|
||||
ErrInvalid = errors.New("ws/hub: invalid connection or group")
|
||||
)
|
||||
|
||||
// Connection is implemented by *ws.Session. Implementations must be safe for
|
||||
// concurrent use, have a stable ID, and make TrySend nonblocking with bounded
|
||||
// memory. Done closes only after connection work finishes.
|
||||
type Connection interface {
|
||||
ID() string
|
||||
Done() <-chan struct{}
|
||||
TrySend(ws.Message) error
|
||||
Close(int, string) error
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
MaxConnections int `json:"maxConnections"`
|
||||
MaxGroups int `json:"maxGroups"`
|
||||
MaxMemberships int `json:"maxMemberships"`
|
||||
MaxGroupsPerConnection int `json:"maxGroupsPerConnection"`
|
||||
}
|
||||
|
||||
func (cfg Config) Validate() error {
|
||||
if cfg.MaxConnections < 0 || cfg.MaxConnections > 1000000 || cfg.MaxGroups < 0 || cfg.MaxGroups > 1000000 || cfg.MaxMemberships < 0 || cfg.MaxMemberships > 10000000 || cfg.MaxGroupsPerConnection < 0 || cfg.MaxGroupsPerConnection > 10000 {
|
||||
return errors.New("ws/hub: invalid limits")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type entry struct {
|
||||
connection Connection
|
||||
groups map[string]struct{}
|
||||
removed chan struct{}
|
||||
}
|
||||
type Hub struct {
|
||||
mu sync.RWMutex
|
||||
config Config
|
||||
connections map[string]*entry
|
||||
groups map[string]map[string]*entry
|
||||
memberships int
|
||||
closed bool
|
||||
done chan struct{}
|
||||
watchers sync.WaitGroup
|
||||
}
|
||||
|
||||
func New(cfg Config) (*Hub, error) {
|
||||
if err := cfg.Validate(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if cfg.MaxConnections == 0 {
|
||||
cfg.MaxConnections = 10000
|
||||
}
|
||||
if cfg.MaxGroups == 0 {
|
||||
cfg.MaxGroups = 1024
|
||||
}
|
||||
if cfg.MaxMemberships == 0 {
|
||||
cfg.MaxMemberships = 100000
|
||||
}
|
||||
if cfg.MaxGroupsPerConnection == 0 {
|
||||
cfg.MaxGroupsPerConnection = 32
|
||||
}
|
||||
return &Hub{config: cfg, connections: make(map[string]*entry), groups: make(map[string]map[string]*entry), done: make(chan struct{})}, nil
|
||||
}
|
||||
|
||||
func validKey(key string) bool {
|
||||
return key != "" && len(key) <= 256 && !strings.ContainsAny(key, "\r\n\x00")
|
||||
}
|
||||
|
||||
// Add borrows a connection and starts one bounded watcher to remove it when
|
||||
// Done closes. Remove and Close stop those watchers without closing connections.
|
||||
func (h *Hub) Add(connection Connection) error {
|
||||
if connection == nil {
|
||||
return ErrInvalid
|
||||
}
|
||||
v := reflect.ValueOf(connection)
|
||||
if (v.Kind() == reflect.Pointer || v.Kind() == reflect.Interface || v.Kind() == reflect.Map || v.Kind() == reflect.Func || v.Kind() == reflect.Slice || v.Kind() == reflect.Chan) && v.IsNil() {
|
||||
return ErrInvalid
|
||||
}
|
||||
id := connection.ID()
|
||||
done := connection.Done()
|
||||
if !validKey(id) || done == nil {
|
||||
return ErrInvalid
|
||||
}
|
||||
select {
|
||||
case <-done:
|
||||
return ws.ErrClosed
|
||||
default:
|
||||
}
|
||||
h.mu.Lock()
|
||||
if h.closed {
|
||||
h.mu.Unlock()
|
||||
return ErrClosed
|
||||
}
|
||||
if _, ok := h.connections[id]; ok {
|
||||
h.mu.Unlock()
|
||||
return ErrDuplicate
|
||||
}
|
||||
if len(h.connections) >= h.config.MaxConnections {
|
||||
h.mu.Unlock()
|
||||
return ErrLimit
|
||||
}
|
||||
e := &entry{connection: connection, groups: make(map[string]struct{}), removed: make(chan struct{})}
|
||||
h.connections[id] = e
|
||||
h.watchers.Add(1)
|
||||
h.mu.Unlock()
|
||||
go func() {
|
||||
defer h.watchers.Done()
|
||||
select {
|
||||
case <-done:
|
||||
h.mu.Lock()
|
||||
if h.connections[id] == e {
|
||||
h.removeLocked(id, e)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
case <-e.removed:
|
||||
case <-h.done:
|
||||
}
|
||||
}()
|
||||
return nil
|
||||
}
|
||||
func (h *Hub) removeLocked(id string, e *entry) {
|
||||
delete(h.connections, id)
|
||||
close(e.removed)
|
||||
for name := range e.groups {
|
||||
delete(h.groups[name], id)
|
||||
h.memberships--
|
||||
if len(h.groups[name]) == 0 {
|
||||
delete(h.groups, name)
|
||||
}
|
||||
}
|
||||
}
|
||||
func (h *Hub) Remove(id string) bool {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
e := h.connections[id]
|
||||
if e == nil {
|
||||
return false
|
||||
}
|
||||
h.removeLocked(id, e)
|
||||
return true
|
||||
}
|
||||
|
||||
func (h *Hub) Join(id, group string) error {
|
||||
if !validKey(group) {
|
||||
return ErrInvalid
|
||||
}
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.closed {
|
||||
return ErrClosed
|
||||
}
|
||||
e := h.connections[id]
|
||||
if e == nil {
|
||||
return ErrNotFound
|
||||
}
|
||||
if _, ok := e.groups[group]; ok {
|
||||
return nil
|
||||
}
|
||||
if len(e.groups) >= h.config.MaxGroupsPerConnection || h.memberships >= h.config.MaxMemberships {
|
||||
return ErrLimit
|
||||
}
|
||||
if h.groups[group] == nil {
|
||||
if len(h.groups) >= h.config.MaxGroups {
|
||||
return ErrLimit
|
||||
}
|
||||
h.groups[group] = make(map[string]*entry)
|
||||
}
|
||||
h.groups[group][id] = e
|
||||
e.groups[group] = struct{}{}
|
||||
h.memberships++
|
||||
return nil
|
||||
}
|
||||
func (h *Hub) Leave(id, group string) bool {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
e := h.connections[id]
|
||||
if e == nil {
|
||||
return false
|
||||
}
|
||||
if _, ok := e.groups[group]; !ok {
|
||||
return false
|
||||
}
|
||||
delete(e.groups, group)
|
||||
delete(h.groups[group], id)
|
||||
h.memberships--
|
||||
if len(h.groups[group]) == 0 {
|
||||
delete(h.groups, group)
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// Result reports queue acceptance, never remote receipt or business success.
|
||||
// Errors contains one entry per rejected destination; no messages are silently
|
||||
// dropped. Concurrent removal may race a snapshot and yield a closed error.
|
||||
type Result struct {
|
||||
Attempted int
|
||||
Queued int
|
||||
Errors map[string]error
|
||||
}
|
||||
|
||||
func send(entries map[string]*entry, message ws.Message) Result {
|
||||
result := Result{Attempted: len(entries)}
|
||||
for id, e := range entries {
|
||||
if err := e.connection.TrySend(message); err != nil {
|
||||
if result.Errors == nil {
|
||||
result.Errors = make(map[string]error)
|
||||
}
|
||||
result.Errors[id] = err
|
||||
} else {
|
||||
result.Queued++
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
func (h *Hub) Send(id string, message ws.Message) (Result, error) {
|
||||
h.mu.RLock()
|
||||
if h.closed {
|
||||
h.mu.RUnlock()
|
||||
return Result{}, ErrClosed
|
||||
}
|
||||
e := h.connections[id]
|
||||
h.mu.RUnlock()
|
||||
if e == nil {
|
||||
return Result{}, ErrNotFound
|
||||
}
|
||||
return send(map[string]*entry{id: e}, message), nil
|
||||
}
|
||||
func (h *Hub) snapshot(group string) (map[string]*entry, error) {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
if h.closed {
|
||||
return nil, ErrClosed
|
||||
}
|
||||
source := h.connections
|
||||
if group != "" {
|
||||
source = h.groups[group]
|
||||
}
|
||||
out := make(map[string]*entry, len(source))
|
||||
for id, e := range source {
|
||||
out[id] = e
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (h *Hub) Broadcast(message ws.Message) (Result, error) {
|
||||
entries, err := h.snapshot("")
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
return send(entries, message), nil
|
||||
}
|
||||
func (h *Hub) SendGroup(group string, message ws.Message) (Result, error) {
|
||||
if !validKey(group) {
|
||||
return Result{}, ErrInvalid
|
||||
}
|
||||
entries, err := h.snapshot(group)
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
return send(entries, message), nil
|
||||
}
|
||||
func (h *Hub) Kick(id string, code int, reason string) error {
|
||||
h.mu.RLock()
|
||||
if h.closed {
|
||||
h.mu.RUnlock()
|
||||
return ErrClosed
|
||||
}
|
||||
e := h.connections[id]
|
||||
h.mu.RUnlock()
|
||||
if e == nil {
|
||||
return ErrNotFound
|
||||
}
|
||||
return e.connection.Close(code, reason)
|
||||
}
|
||||
|
||||
type Stats struct {
|
||||
Connections int
|
||||
Groups int
|
||||
Memberships int
|
||||
}
|
||||
|
||||
func (h *Hub) Stats() Stats {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
return Stats{len(h.connections), len(h.groups), h.memberships}
|
||||
}
|
||||
|
||||
// Close unregisters borrowed connections, releases membership memory and waits
|
||||
// for internal watchers only. Use Session/Server shutdown to close connections.
|
||||
func (h *Hub) Close() error {
|
||||
h.mu.Lock()
|
||||
if !h.closed {
|
||||
h.closed = true
|
||||
close(h.done)
|
||||
for id, e := range h.connections {
|
||||
h.removeLocked(id, e)
|
||||
}
|
||||
h.groups = make(map[string]map[string]*entry)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
h.watchers.Wait()
|
||||
return nil
|
||||
}
|
||||
271
ws/hub/hub_test.go
Normal file
271
ws/hub/hub_test.go
Normal file
@@ -0,0 +1,271 @@
|
||||
package hub
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.yuntuo.hk/k2safe/go-omni/ws"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
type testConnection struct {
|
||||
id string
|
||||
done chan struct{}
|
||||
once sync.Once
|
||||
full bool
|
||||
sent atomic.Int64
|
||||
}
|
||||
|
||||
func (c *testConnection) ID() string { return c.id }
|
||||
func (c *testConnection) Done() <-chan struct{} { return c.done }
|
||||
func (c *testConnection) TrySend(ws.Message) error {
|
||||
select {
|
||||
case <-c.done:
|
||||
return ws.ErrClosed
|
||||
default:
|
||||
}
|
||||
if c.full {
|
||||
return ws.ErrQueueFull
|
||||
}
|
||||
c.sent.Add(1)
|
||||
return nil
|
||||
}
|
||||
func (c *testConnection) Close(int, string) error { c.once.Do(func() { close(c.done) }); return nil }
|
||||
func newConnection(id string) *testConnection {
|
||||
return &testConnection{id: id, done: make(chan struct{})}
|
||||
}
|
||||
func testHub(t *testing.T, cfg Config) *Hub {
|
||||
t.Helper()
|
||||
h, err := New(cfg)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = h.Close() })
|
||||
return h
|
||||
}
|
||||
|
||||
func TestHubRoutingBackpressureAndCleanup(t *testing.T) {
|
||||
h := testHub(t, Config{})
|
||||
a, b, c := newConnection("a"), newConnection("b"), newConnection("c")
|
||||
b.full = true
|
||||
for _, connection := range []*testConnection{a, b, c} {
|
||||
if err := h.Add(connection); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
for _, id := range []string{"a", "b"} {
|
||||
if err := h.Join(id, "group"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
message := ws.Message{Type: ws.Text, Data: []byte("hello")}
|
||||
result, err := h.SendGroup("group", message)
|
||||
if err != nil || result.Attempted != 2 || result.Queued != 1 || !errors.Is(result.Errors["b"], ws.ErrQueueFull) {
|
||||
t.Fatalf("group: %#v %v", result, err)
|
||||
}
|
||||
if c.sent.Load() != 0 {
|
||||
t.Fatal("group leaked")
|
||||
}
|
||||
result, err = h.Broadcast(message)
|
||||
if err != nil || result.Attempted != 3 || result.Queued != 2 {
|
||||
t.Fatalf("broadcast: %#v %v", result, err)
|
||||
}
|
||||
if err = h.Kick("a", 1000, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
deadline := time.After(time.Second)
|
||||
for h.Stats().Connections != 2 {
|
||||
select {
|
||||
case <-deadline:
|
||||
t.Fatal("auto cleanup timeout")
|
||||
default:
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
if got := h.Stats(); got.Memberships != 1 || got.Groups != 1 {
|
||||
t.Fatalf("stats: %#v", got)
|
||||
}
|
||||
if !h.Remove("b") {
|
||||
t.Fatal("remove")
|
||||
}
|
||||
if got := h.Stats(); got.Memberships != 0 || got.Groups != 0 {
|
||||
t.Fatalf("group leak: %#v", got)
|
||||
}
|
||||
if _, err = h.Send("missing", message); !errors.Is(err, ErrNotFound) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = h.Close()
|
||||
if _, err = h.Broadcast(message); !errors.Is(err, ErrClosed) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
select {
|
||||
case <-c.done:
|
||||
t.Fatal("hub closed borrowed connection")
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func TestHubLimitsAndReusedID(t *testing.T) {
|
||||
h := testHub(t, Config{MaxConnections: 2, MaxGroups: 1, MaxMemberships: 1, MaxGroupsPerConnection: 1})
|
||||
first := newConnection("a")
|
||||
if err := h.Add(first); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := h.Add(first); !errors.Is(err, ErrDuplicate) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := h.Add(newConnection("b")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := h.Add(newConnection("c")); !errors.Is(err, ErrLimit) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := h.Join("a", "one"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := h.Join("a", "one"); err != nil {
|
||||
t.Fatal("idempotent join", err)
|
||||
}
|
||||
for _, pair := range [][2]string{{"a", "two"}, {"b", "one"}} {
|
||||
if err := h.Join(pair[0], pair[1]); !errors.Is(err, ErrLimit) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
h.Remove("a")
|
||||
second := newConnection("a")
|
||||
if err := h.Add(second); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = first.Close(1000, "")
|
||||
if err := h.Join("a", "replacement"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := h.Send("a", ws.Message{Type: ws.Text}); err != nil {
|
||||
t.Fatal("old watcher removed replacement", err)
|
||||
}
|
||||
var nilConnection *testConnection
|
||||
if err := h.Add(nilConnection); !errors.Is(err, ErrInvalid) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHubConcurrentLifecycle(t *testing.T) {
|
||||
h := testHub(t, Config{MaxConnections: 100})
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 50; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
c := newConnection(fmt.Sprint(i))
|
||||
if err := h.Add(c); err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
for j := 0; j < 10; j++ {
|
||||
_ = h.Join(c.ID(), "shared")
|
||||
_, _ = h.Broadcast(ws.Message{Type: ws.Text})
|
||||
h.Leave(c.ID(), "shared")
|
||||
}
|
||||
_ = c.Close(1000, "")
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
_ = h.Close()
|
||||
if h.Stats() != (Stats{}) {
|
||||
t.Fatal("registry not released")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHubRealSessionDeliveryAndDisconnect(t *testing.T) {
|
||||
received := make(chan string, 1)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
for {
|
||||
_, data, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
received <- string(data)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
client, err := ws.NewClient(ws.ClientConfig{ProxyMode: "direct", Session: ws.SessionConfig{CloseTimeout: 100 * time.Millisecond}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := client.Shutdown(ctx); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}()
|
||||
s, err := client.Dial(context.Background(), "ws"+strings.TrimPrefix(server.URL, "http"), ws.Handlers{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h := testHub(t, Config{})
|
||||
if err = h.Add(s); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = h.Join(s.ID(), "room"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
result, err := h.SendGroup("room", ws.Message{Type: ws.Text, Data: []byte("queued")})
|
||||
if err != nil || result.Queued != 1 {
|
||||
t.Fatalf("send: %#v %v", result, err)
|
||||
}
|
||||
select {
|
||||
case got := <-received:
|
||||
if got != "queued" {
|
||||
t.Fatal(got)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("delivery timed out")
|
||||
}
|
||||
_ = s.Close(1000, "")
|
||||
deadline := time.After(time.Second)
|
||||
for h.Stats().Connections != 0 {
|
||||
select {
|
||||
case <-deadline:
|
||||
t.Fatal("disconnected Session still registered")
|
||||
default:
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkHubBroadcast(b *testing.B) {
|
||||
for _, count := range []int{1, 100, 1000} {
|
||||
b.Run(fmt.Sprint(count), func(b *testing.B) {
|
||||
h, _ := New(Config{})
|
||||
defer h.Close()
|
||||
for i := 0; i < count; i++ {
|
||||
if err := h.Add(newConnection(fmt.Sprint(i))); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
message := ws.Message{Type: ws.Text, Data: []byte("benchmark")}
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
result, err := h.Broadcast(message)
|
||||
if err != nil || result.Queued != count {
|
||||
b.Fatal("broadcast failed")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
100
ws/limits_test.go
Normal file
100
ws/limits_test.go
Normal file
@@ -0,0 +1,100 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
func TestSessionConfigUpperBounds(t *testing.T) {
|
||||
for _, c := range []SessionConfig{{MaxMessageBytes: 1<<30 + 1}, {ReadQueueSize: 65537}, {WriteQueueSize: 65537}, {ReadQueueBytes: 1<<30 + 1}, {WriteQueueBytes: 1<<30 + 1}, {CloseTimeout: 25 * time.Hour}} {
|
||||
if c.Validate() == nil {
|
||||
t.Fatal("oversized config accepted")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientReconnectsRemoteGoingAway(t *testing.T) {
|
||||
var handshakes atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
if handshakes.Add(1) == 1 {
|
||||
_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(1001, ""), time.Now().Add(time.Second))
|
||||
}
|
||||
for {
|
||||
if _, _, err = conn.ReadMessage(); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
client, err := NewClient(ClientConfig{ProxyMode: "direct", Reconnect: ReconnectConfig{Enabled: true, MinBackoff: time.Millisecond, MaxBackoff: time.Millisecond}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := client.Shutdown(ctx); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
defer cancel()
|
||||
connected := make(chan struct{}, 1)
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
done <- client.Run(ctx, "ws"+strings.TrimPrefix(server.URL, "http"), Handlers{OnOpen: func(context.Context, *Session) error {
|
||||
if handshakes.Load() >= 2 {
|
||||
connected <- struct{}{}
|
||||
}
|
||||
return nil
|
||||
}})
|
||||
}()
|
||||
select {
|
||||
case <-connected:
|
||||
case err := <-done:
|
||||
t.Fatalf("Run ended before reconnect: %v", err)
|
||||
case <-ctx.Done():
|
||||
t.Fatal(ctx.Err())
|
||||
}
|
||||
_ = client.Close()
|
||||
select {
|
||||
case err := <-done:
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
case <-ctx.Done():
|
||||
t.Fatal(ctx.Err())
|
||||
}
|
||||
if stats := client.Stats(); stats.DialAttempts != 2 || stats.Reconnects != 1 || stats.HandshakeFailures != 0 {
|
||||
t.Fatal(stats)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientHandshakeStats(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Error(w, "denied", http.StatusUnauthorized) }))
|
||||
defer server.Close()
|
||||
client, err := NewClient(ClientConfig{ProxyMode: "direct"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer client.Close()
|
||||
_, err = client.Dial(context.Background(), "ws"+strings.TrimPrefix(server.URL, "http"), Handlers{})
|
||||
if err == nil {
|
||||
t.Fatal("expected rejection")
|
||||
}
|
||||
if stats := client.Stats(); stats.DialAttempts != 1 || stats.HandshakeFailures != 1 || stats.Reconnects != 0 {
|
||||
t.Fatal(stats)
|
||||
}
|
||||
}
|
||||
44
ws/observe/README.md
Normal file
44
ws/observe/README.md
Normal file
@@ -0,0 +1,44 @@
|
||||
# WebSocket 可选观测
|
||||
|
||||
`ws/observe` 为 `ws.Client`、`ws.Server` 的会话提供连接日志和聚合快照。核心 `ws` 不依赖全局 Logger；观测器不启动后台 goroutine，不给每条消息增加日志、锁或 span。
|
||||
|
||||
```go
|
||||
logger, err := logx.New(logx.Default())
|
||||
if err != nil { return err }
|
||||
defer logger.Close() // 必须晚于 Client/Server 的 Shutdown 完成
|
||||
|
||||
observer := observe.New(logger) // nil 表示只统计，不记日志
|
||||
handlers := observer.Wrap(ws.Handlers{
|
||||
OnMessage: func(ctx context.Context, session *ws.Session, message ws.Message) error {
|
||||
return session.Send(ctx, message)
|
||||
},
|
||||
})
|
||||
server, err := ws.NewServer(ws.ServerConfig{})
|
||||
if err != nil { return err }
|
||||
mux.Handle("/ws", server.Handler(handlers))
|
||||
|
||||
// 由项目按自己的采样周期读取，接入 metrics、Prometheus 或 OTel。
|
||||
snapshot := observer.Snapshot()
|
||||
```
|
||||
|
||||
导入路径为 `git.yuntuo.hk/k2safe/go-omni/ws/observe`。同一个 Observer 和同一组包装后的 Handlers 可供多个连接并发复用；每个会话有独立的开始时间。每个会话仅包装一次。使用客户端重连和 `ClientConfig.Prepare` 时，在每次 Prepare 返回前调用 `observer.Wrap(handlers)`。
|
||||
|
||||
快照包含：
|
||||
|
||||
- `Opened / Closed / Active`：已观察到的连接总数、关闭回调结束数、当前处于受观察生命周期中的连接数。连接可能在 OnOpen 执行前断开，此时仍在 OnClose 计入 Opened 和 Closed，日志时长为 0。
|
||||
- `MessagesSent / MessagesReceived / BytesSent / BytesReceived`：已关闭会话的最终统计，加上当前会话的即时统计。发送成功表示写入完成；不是对端收到消息或业务处理成功的确认。统计口径与 `Session.Stats()` 一致。
|
||||
- `WriteQueuedMessages / WriteQueuedBytes / ReadQueuedMessages / ReadQueuedBytes`：调用快照时，各活跃连接队列水位之和。接收侧包括仍在处理、尚未释放的消息。
|
||||
- `SampledPeak*`：历次 Snapshot 观察到的聚合峰值；不是精确的历史峰值，采样间隔内的突增可能未被记录。
|
||||
- `CloseCallbackPanics`：下游 OnClose 回调发生 panic 的次数。
|
||||
|
||||
Snapshot 为 O(活跃会话数)，只在采样时读取 `Session.Stats()`，不会替换消息回调。避免在每条消息上采样；采样周期由项目根据连接规模决定。快照不会保存已关闭连接的 ID，也不引入高基数指标标签。自行导出 OTel 时，建议记录握手、连接时长、消息数与字节数；业务消息追踪应显式传递上下文，不自动延长握手 span 到整条连接的寿命。
|
||||
|
||||
日志只由适配器写入固定事件名、Session ID、时长、关闭码和固定阶段。不会展开错误文本、Cause、URL、请求头、消息体或对端关闭原因；业务自定义 `OpError.Op` 也通过固定白名单筛选。用户传入 Logger 自身已经附加的属性仍由用户负责。默认日志级别为 Info。
|
||||
|
||||
OnClose 的收尾使用 defer，即使业务回调 panic，也会归档消息统计、移除活跃记录、记录固定的 `on_close_panic` 阶段，再把 panic 交还核心处理。此处“关闭”表示 I/O 已停止且下游 OnClose 已返回或 panic；核心随后才关闭 `Session.Done()`。因此 `Snapshot.Active == 0` 不是生命周期完成屏障，停机必须等待 `Client/Server.Shutdown` 或 `Session.Done`，再关闭 Logger。观测器不会在 OnClose 中等待自身 Done，也不会在 Done 后继续写日志。
|
||||
|
||||
运行真实 TCP 与并发测试：
|
||||
|
||||
```sh
|
||||
go test -race ./ws/observe
|
||||
```
|
||||
194
ws/observe/observe.go
Normal file
194
ws/observe/observe.go
Normal file
@@ -0,0 +1,194 @@
|
||||
// Package observe adds opt-in, payload-free WebSocket lifecycle logs and
|
||||
// aggregate snapshots. The ws core has no logger dependency. Observers do not
|
||||
// own loggers, start goroutines, or create per-message tracing spans.
|
||||
package observe
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"git.yuntuo.hk/k2safe/go-omni/logx"
|
||||
"git.yuntuo.hk/k2safe/go-omni/ws"
|
||||
)
|
||||
|
||||
// Snapshot contains process-local totals for sessions observed by this Observer.
|
||||
// Active covers observed callback lifetimes: removal happens after the user's
|
||||
// OnClose returns, immediately before Session.Done. It is not a drain barrier.
|
||||
// Queue peaks are sampled only when Snapshot is called, not lifetime maxima.
|
||||
type Snapshot struct {
|
||||
Opened uint64 `json:"opened"`
|
||||
Closed uint64 `json:"closed"`
|
||||
Active int `json:"active"`
|
||||
CloseCallbackPanics uint64 `json:"closeCallbackPanics"`
|
||||
MessagesSent uint64 `json:"messagesSent"`
|
||||
MessagesReceived uint64 `json:"messagesReceived"`
|
||||
BytesSent uint64 `json:"bytesSent"`
|
||||
BytesReceived uint64 `json:"bytesReceived"`
|
||||
WriteQueuedMessages int `json:"writeQueuedMessages"`
|
||||
WriteQueuedBytes int64 `json:"writeQueuedBytes"`
|
||||
ReadQueuedMessages int `json:"readQueuedMessages"`
|
||||
ReadQueuedBytes int64 `json:"readQueuedBytes"`
|
||||
SampledPeakWriteQueuedMessages int `json:"sampledPeakWriteQueuedMessages"`
|
||||
SampledPeakWriteQueuedBytes int64 `json:"sampledPeakWriteQueuedBytes"`
|
||||
SampledPeakReadQueuedMessages int `json:"sampledPeakReadQueuedMessages"`
|
||||
SampledPeakReadQueuedBytes int64 `json:"sampledPeakReadQueuedBytes"`
|
||||
}
|
||||
|
||||
type observation struct{ started time.Time }
|
||||
|
||||
type Observer struct {
|
||||
log *logx.Logger
|
||||
mu sync.Mutex
|
||||
active map[*ws.Session]observation
|
||||
totals Snapshot
|
||||
}
|
||||
|
||||
// New borrows log without changing or closing it. nil disables logs and retains
|
||||
// metrics. Keep the logger open until Client/Server shutdown completes.
|
||||
func New(log *logx.Logger) *Observer {
|
||||
return &Observer{log: log, active: make(map[*ws.Session]observation)}
|
||||
}
|
||||
|
||||
// Wrap can be reused across concurrent sessions. It leaves OnMessage unchanged;
|
||||
// message accounting comes from Session.Stats during snapshots and closure.
|
||||
// Apply it once per session. With ClientConfig.Prepare, wrap each returned set
|
||||
// of handlers so each reconnect uses the same observer and its own start time.
|
||||
func (o *Observer) Wrap(next ws.Handlers) ws.Handlers {
|
||||
if o == nil {
|
||||
return next
|
||||
}
|
||||
return ws.Handlers{
|
||||
OnOpen: func(ctx context.Context, s *ws.Session) error {
|
||||
o.opened(s)
|
||||
if next.OnOpen != nil {
|
||||
return next.OnOpen(ctx, s)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
OnMessage: next.OnMessage,
|
||||
OnClose: func(s *ws.Session, err error) {
|
||||
defer func() {
|
||||
p := recover()
|
||||
o.closed(s, err, p != nil)
|
||||
if p != nil {
|
||||
panic(p)
|
||||
}
|
||||
}()
|
||||
if next.OnClose != nil {
|
||||
next.OnClose(s, err)
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (o *Observer) opened(s *ws.Session) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
o.mu.Lock()
|
||||
_, exists := o.active[s]
|
||||
if !exists {
|
||||
o.active[s] = observation{started: time.Now()}
|
||||
o.totals.Opened++
|
||||
}
|
||||
o.mu.Unlock()
|
||||
if !exists && o.log != nil {
|
||||
o.log.Info("ws.open", "sessionId", s.ID(), "phase", "open")
|
||||
}
|
||||
}
|
||||
|
||||
func (o *Observer) closed(s *ws.Session, err error, panicked bool) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
stats := s.Stats()
|
||||
o.mu.Lock()
|
||||
observed, exists := o.active[s]
|
||||
delete(o.active, s)
|
||||
// A connection can close before its dispatch loop reaches OnOpen. Count
|
||||
// that observed physical connection too, with zero measured duration.
|
||||
if !exists {
|
||||
o.totals.Opened++
|
||||
}
|
||||
o.totals.Closed++
|
||||
if panicked {
|
||||
o.totals.CloseCallbackPanics++
|
||||
}
|
||||
addMessages(&o.totals, stats)
|
||||
o.mu.Unlock()
|
||||
duration := time.Duration(0)
|
||||
if !observed.started.IsZero() {
|
||||
duration = time.Since(observed.started)
|
||||
}
|
||||
if o.log != nil {
|
||||
code, phase := closeFields(err, panicked)
|
||||
o.log.Info("ws.close", "sessionId", s.ID(), "durationMs", duration.Seconds()*1000, "closeCode", code, "phase", phase)
|
||||
}
|
||||
}
|
||||
|
||||
func closeFields(err error, panicked bool) (int, string) {
|
||||
if panicked {
|
||||
return 0, "on_close_panic"
|
||||
}
|
||||
if err == nil {
|
||||
return 1000, "close"
|
||||
}
|
||||
var closed *ws.CloseError
|
||||
if errors.As(err, &closed) {
|
||||
return closed.Code, "close"
|
||||
}
|
||||
var op *ws.OpError
|
||||
if errors.As(err, &op) {
|
||||
// OpError is exported. Never write an arbitrary user-created Op to logs.
|
||||
switch op.Op {
|
||||
case "read", "write", "on_open", "on_message", "on_close", "ping", "pong", "close", "encode":
|
||||
return 0, op.Op
|
||||
}
|
||||
}
|
||||
if errors.Is(err, context.Canceled) {
|
||||
return 0, "canceled"
|
||||
}
|
||||
if errors.Is(err, context.DeadlineExceeded) {
|
||||
return 0, "timeout"
|
||||
}
|
||||
return 0, "error"
|
||||
}
|
||||
|
||||
func addMessages(out *Snapshot, stats ws.SessionStats) {
|
||||
out.MessagesSent += stats.MessagesSent
|
||||
out.MessagesReceived += stats.MessagesReceived
|
||||
out.BytesSent += stats.BytesSent
|
||||
out.BytesReceived += stats.BytesReceived
|
||||
}
|
||||
|
||||
// Snapshot samples live Session counters and adds immutable closed totals. It
|
||||
// costs O(active sessions), acquires no locks on the message path, and keeps no
|
||||
// per-connection records after OnClose. Message totals count I/O, not peer ACKs.
|
||||
func (o *Observer) Snapshot() Snapshot {
|
||||
if o == nil {
|
||||
return Snapshot{}
|
||||
}
|
||||
o.mu.Lock()
|
||||
defer o.mu.Unlock()
|
||||
out := o.totals
|
||||
out.Active = len(o.active)
|
||||
for session := range o.active {
|
||||
stats := session.Stats()
|
||||
addMessages(&out, stats)
|
||||
out.WriteQueuedMessages += stats.QueuedMessages
|
||||
out.WriteQueuedBytes += stats.QueuedBytes
|
||||
out.ReadQueuedMessages += stats.ReadQueuedMessages
|
||||
out.ReadQueuedBytes += stats.ReadQueuedBytes
|
||||
}
|
||||
o.totals.SampledPeakWriteQueuedMessages = max(o.totals.SampledPeakWriteQueuedMessages, out.WriteQueuedMessages)
|
||||
o.totals.SampledPeakWriteQueuedBytes = max(o.totals.SampledPeakWriteQueuedBytes, out.WriteQueuedBytes)
|
||||
o.totals.SampledPeakReadQueuedMessages = max(o.totals.SampledPeakReadQueuedMessages, out.ReadQueuedMessages)
|
||||
o.totals.SampledPeakReadQueuedBytes = max(o.totals.SampledPeakReadQueuedBytes, out.ReadQueuedBytes)
|
||||
out.SampledPeakWriteQueuedMessages = o.totals.SampledPeakWriteQueuedMessages
|
||||
out.SampledPeakWriteQueuedBytes = o.totals.SampledPeakWriteQueuedBytes
|
||||
out.SampledPeakReadQueuedMessages = o.totals.SampledPeakReadQueuedMessages
|
||||
out.SampledPeakReadQueuedBytes = o.totals.SampledPeakReadQueuedBytes
|
||||
return out
|
||||
}
|
||||
300
ws/observe/observe_test.go
Normal file
300
ws/observe/observe_test.go
Normal file
@@ -0,0 +1,300 @@
|
||||
package observe
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.yuntuo.hk/k2safe/go-omni/logx"
|
||||
"git.yuntuo.hk/k2safe/go-omni/ws"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
func echo(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
for {
|
||||
kind, data, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if err = conn.WriteMessage(kind, data); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
func testClient(t *testing.T) *ws.Client {
|
||||
t.Helper()
|
||||
c, err := ws.NewClient(ws.ClientConfig{ProxyMode: "direct", Session: ws.SessionConfig{CloseTimeout: 50 * time.Millisecond}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := c.Shutdown(ctx); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
})
|
||||
return c
|
||||
}
|
||||
func endpoint(s *httptest.Server) string { return "ws" + strings.TrimPrefix(s.URL, "http") }
|
||||
func wait(t *testing.T, ch <-chan struct{}) {
|
||||
t.Helper()
|
||||
select {
|
||||
case <-ch:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timeout")
|
||||
}
|
||||
}
|
||||
|
||||
func TestObserverConcurrentSessionsAndTotals(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(echo))
|
||||
defer server.Close()
|
||||
client := testClient(t)
|
||||
o := New(nil)
|
||||
const connections = 8
|
||||
const count = 10
|
||||
var workers sync.WaitGroup
|
||||
stopSamples := make(chan struct{})
|
||||
samplesDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(samplesDone)
|
||||
for {
|
||||
select {
|
||||
case <-stopSamples:
|
||||
return
|
||||
default:
|
||||
_ = o.Snapshot()
|
||||
}
|
||||
}
|
||||
}()
|
||||
for i := 0; i < connections; i++ {
|
||||
workers.Add(1)
|
||||
go func() {
|
||||
defer workers.Done()
|
||||
opened := make(chan struct{})
|
||||
received := make(chan struct{}, count)
|
||||
s, err := client.Dial(context.Background(), endpoint(server), o.Wrap(ws.Handlers{OnOpen: func(context.Context, *ws.Session) error { close(opened); return nil }, OnMessage: func(context.Context, *ws.Session, ws.Message) error { received <- struct{}{}; return nil }}))
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
wait(t, opened)
|
||||
for j := 0; j < count; j++ {
|
||||
if err = s.Send(context.Background(), ws.Message{Type: ws.Text, Data: []byte("test")}); err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
wait(t, received)
|
||||
}
|
||||
_ = s.Close(1000, "")
|
||||
wait(t, s.Done())
|
||||
}()
|
||||
}
|
||||
workers.Wait()
|
||||
close(stopSamples)
|
||||
wait(t, samplesDone)
|
||||
got := o.Snapshot()
|
||||
if got.Opened != connections || got.Closed != connections || got.Active != 0 || got.MessagesSent != connections*count || got.MessagesReceived != connections*count || got.BytesSent != 4*connections*count || got.BytesReceived != 4*connections*count {
|
||||
t.Fatalf("snapshot: %+v", got)
|
||||
}
|
||||
if again := o.Snapshot(); again != got {
|
||||
t.Fatalf("closed totals changed: %+v -> %+v", got, again)
|
||||
}
|
||||
}
|
||||
|
||||
func TestObserverCloseCallbackBoundaryAndPanic(t *testing.T) {
|
||||
for _, panics := range []bool{false, true} {
|
||||
t.Run(fmt.Sprint(panics), func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(echo))
|
||||
defer server.Close()
|
||||
client := testClient(t)
|
||||
o := New(nil)
|
||||
opened := make(chan struct{})
|
||||
closing := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
s, err := client.Dial(context.Background(), endpoint(server), o.Wrap(ws.Handlers{OnOpen: func(context.Context, *ws.Session) error { close(opened); return nil }, OnClose: func(*ws.Session, error) {
|
||||
close(closing)
|
||||
<-release
|
||||
if panics {
|
||||
panic("SECRET CALLBACK PANIC")
|
||||
}
|
||||
}}))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wait(t, opened)
|
||||
_ = s.Close(1000, "")
|
||||
wait(t, closing)
|
||||
if got := o.Snapshot(); got.Active != 1 || got.Closed != 0 {
|
||||
t.Fatalf("closed before callback finished: %+v", got)
|
||||
}
|
||||
select {
|
||||
case <-s.Done():
|
||||
t.Fatal("Done closed before user callback")
|
||||
default:
|
||||
}
|
||||
close(release)
|
||||
wait(t, s.Done())
|
||||
got := o.Snapshot()
|
||||
if got.Active != 0 || got.Closed != 1 {
|
||||
t.Fatalf("callback did not release observation: %+v", got)
|
||||
}
|
||||
if panics && got.CloseCallbackPanics != 1 {
|
||||
t.Fatalf("panic not counted: %+v", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestObserverPayloadFreeLogsAndIndependentDurations(t *testing.T) {
|
||||
var output bytes.Buffer
|
||||
logger, err := logx.New(logx.Config{Writer: &output, Format: "json", Level: "info"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer logger.Close()
|
||||
o := New(logger)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
if _, _, err = conn.ReadMessage(); err == nil {
|
||||
_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(1008, "SECRET CLOSE REASON"), time.Now().Add(time.Second))
|
||||
_, _, _ = conn.ReadMessage()
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
client := testClient(t)
|
||||
var sessions []*ws.Session
|
||||
for i := 0; i < 2; i++ {
|
||||
opened := make(chan struct{})
|
||||
s, err := client.Dial(context.Background(), endpoint(server)+"/?token=SECRET-URL", o.Wrap(ws.Handlers{OnOpen: func(context.Context, *ws.Session) error { close(opened); return nil }}))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wait(t, opened)
|
||||
sessions = append(sessions, s)
|
||||
if i == 0 {
|
||||
time.Sleep(40 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
for _, s := range sessions {
|
||||
if err = s.Send(context.Background(), ws.Message{Type: ws.Text, Data: []byte("SECRET MESSAGE")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wait(t, s.Done())
|
||||
}
|
||||
if strings.Contains(output.String(), "SECRET") {
|
||||
t.Fatalf("log leaked source data: %s", output.String())
|
||||
}
|
||||
var durations []float64
|
||||
for _, line := range strings.Split(strings.TrimSpace(output.String()), "\n") {
|
||||
var row map[string]any
|
||||
if err = json.Unmarshal([]byte(line), &row); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for key := range row {
|
||||
switch key {
|
||||
case "time", "level", "msg", "sessionId", "phase", "durationMs", "closeCode":
|
||||
default:
|
||||
t.Fatalf("unexpected log field %q", key)
|
||||
}
|
||||
}
|
||||
if row["msg"] == "ws.close" {
|
||||
if row["closeCode"] != float64(1008) {
|
||||
t.Fatalf("close code: %#v", row)
|
||||
}
|
||||
durations = append(durations, row["durationMs"].(float64))
|
||||
}
|
||||
}
|
||||
if len(durations) != 2 || durations[0] < 30 {
|
||||
t.Fatalf("timers not per session: %v", durations)
|
||||
}
|
||||
}
|
||||
|
||||
func TestObserverQueueSnapshotsAndSampledPeaks(t *testing.T) {
|
||||
startSending := make(chan struct{})
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
<-startSending
|
||||
for i := 0; i < 5; i++ {
|
||||
if conn.WriteMessage(websocket.TextMessage, []byte("queue")) != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
for {
|
||||
if _, _, err = conn.ReadMessage(); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
client := testClient(t)
|
||||
o := New(nil)
|
||||
opened := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
s, err := client.Dial(context.Background(), endpoint(server), o.Wrap(ws.Handlers{OnOpen: func(context.Context, *ws.Session) error { close(opened); return nil }, OnMessage: func(context.Context, *ws.Session, ws.Message) error { <-release; return nil }}))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wait(t, opened)
|
||||
close(startSending)
|
||||
deadline := time.After(2 * time.Second)
|
||||
var sampled Snapshot
|
||||
for {
|
||||
sampled = o.Snapshot()
|
||||
if sampled.MessagesReceived == 5 {
|
||||
break
|
||||
}
|
||||
select {
|
||||
case <-deadline:
|
||||
close(release)
|
||||
t.Fatal("receive timeout")
|
||||
default:
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
if sampled.ReadQueuedMessages != 5 || sampled.ReadQueuedBytes != 25 || sampled.SampledPeakReadQueuedMessages != 5 {
|
||||
t.Fatalf("queue snapshot: %+v", sampled)
|
||||
}
|
||||
close(release)
|
||||
_ = s.Close(1000, "")
|
||||
wait(t, s.Done())
|
||||
final := o.Snapshot()
|
||||
if final.ReadQueuedMessages != 0 || final.SampledPeakReadQueuedMessages != 5 {
|
||||
t.Fatalf("final: %+v", final)
|
||||
}
|
||||
}
|
||||
|
||||
type opaqueError struct{}
|
||||
|
||||
func (opaqueError) Error() string { panic("must not expand error") }
|
||||
func TestObserverCloseFieldsNeverExpandsErrors(t *testing.T) {
|
||||
for _, err := range []error{opaqueError{}, &ws.OpError{Op: "SECRET STAGE", Cause: opaqueError{}}, &ws.CloseError{Code: 1008, Cause: opaqueError{}}} {
|
||||
_, phase := closeFields(err, false)
|
||||
if strings.Contains(phase, "SECRET") {
|
||||
t.Fatal("untrusted phase")
|
||||
}
|
||||
}
|
||||
if code, phase := closeFields(opaqueError{}, true); code != 0 || phase != "on_close_panic" {
|
||||
t.Fatalf("panic fields: %d %s", code, phase)
|
||||
}
|
||||
}
|
||||
44
ws/omni/handler.go
Normal file
44
ws/omni/handler.go
Normal file
@@ -0,0 +1,44 @@
|
||||
// Package omni connects managed WebSocket sessions to a GoOmni application.
|
||||
package omni
|
||||
|
||||
import (
|
||||
"errors"
|
||||
|
||||
"git.yuntuo.hk/k2safe/go-omni/app"
|
||||
"git.yuntuo.hk/k2safe/go-omni/config"
|
||||
"git.yuntuo.hk/k2safe/go-omni/httpx"
|
||||
"git.yuntuo.hk/k2safe/go-omni/ws"
|
||||
)
|
||||
|
||||
// Handler preserves the HTTP middleware chain and records the actual upgrade
|
||||
// status. The server owns all session callbacks after this handler returns.
|
||||
// Register Server.Stop/Drain on App or svc Runtime before accepting connections.
|
||||
func Handler(server *ws.Server, handlers ws.Handlers) httpx.Handler {
|
||||
if server == nil {
|
||||
panic("ws/omni: nil server")
|
||||
}
|
||||
return func(c *httpx.Context) error {
|
||||
_, err := server.Accept(c.Raw(), c.Request(), handlers)
|
||||
if err == nil {
|
||||
c.MarkUpgraded()
|
||||
}
|
||||
// Accept has already written any handshake failure response.
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// Mount registers a GET endpoint and its shutdown barriers on a standalone App.
|
||||
// Use Handler with svc.OnStop/OnDrain when the service Runtime owns dependencies.
|
||||
// Mount once per Server. Configure the App before calling it.
|
||||
func Mount(a *app.App, path string, server *ws.Server, handlers ws.Handlers) error {
|
||||
if a == nil || server == nil {
|
||||
return errors.New("ws/omni: app and server required")
|
||||
}
|
||||
if a.Config().Transport != config.TransportNetHTTP && a.Config().Transport != "" {
|
||||
return ws.ErrUnsupportedTransport
|
||||
}
|
||||
a.GET(path, Handler(server, handlers))
|
||||
a.OnStop("ws:"+path, server.Stop)
|
||||
a.OnDrain("ws:"+path, server.Drain)
|
||||
return nil
|
||||
}
|
||||
92
ws/omni/handler_test.go
Normal file
92
ws/omni/handler_test.go
Normal file
@@ -0,0 +1,92 @@
|
||||
package omni_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
omni "git.yuntuo.hk/k2safe/go-omni"
|
||||
"git.yuntuo.hk/k2safe/go-omni/httpx"
|
||||
"git.yuntuo.hk/k2safe/go-omni/lifecycle"
|
||||
"git.yuntuo.hk/k2safe/go-omni/ws"
|
||||
wsomni "git.yuntuo.hk/k2safe/go-omni/ws/omni"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
func TestMountLifetimeAndStatus(t *testing.T) {
|
||||
a := omni.New(omni.WithLogger(false))
|
||||
var observed atomic.Int64
|
||||
a.Use(func(next httpx.Handler) httpx.Handler {
|
||||
return func(c *httpx.Context) error { err := next(c); observed.Store(int64(c.Status())); return err }
|
||||
})
|
||||
s, err := ws.NewServer(ws.ServerConfig{Session: ws.SessionConfig{CloseTimeout: 30 * time.Millisecond}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = wsomni.Mount(a, "/ws", s, ws.Handlers{OnMessage: func(ctx context.Context, s *ws.Session, m ws.Message) error { return s.Send(ctx, m) }}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h := httptest.NewServer(a)
|
||||
defer h.Close()
|
||||
defer a.Close()
|
||||
c, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(h.URL, "http")+"/ws", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer c.Close()
|
||||
for i := 0; i < 10; i++ {
|
||||
if err = c.WriteMessage(websocket.TextMessage, []byte("hello")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c.SetReadDeadline(time.Now().Add(time.Second))
|
||||
_, v, e := c.ReadMessage()
|
||||
if e != nil || string(v) != "hello" {
|
||||
t.Fatal(e)
|
||||
}
|
||||
}
|
||||
if observed.Load() != 101 {
|
||||
t.Fatal("upgrade logged as", observed.Load())
|
||||
}
|
||||
var cleanup atomic.Bool
|
||||
a.OnShutdown(func() error {
|
||||
if s.Stats().ActiveConnections != 0 {
|
||||
t.Error("resource closed before sessions")
|
||||
}
|
||||
cleanup.Store(true)
|
||||
return nil
|
||||
})
|
||||
if err = a.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !cleanup.Load() {
|
||||
t.Fatal("missing cleanup")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAppRetainsResourcesForBlockedComponent(t *testing.T) {
|
||||
a := omni.New(omni.WithLogger(false))
|
||||
closed := false
|
||||
a.OnStop("blocked", func(context.Context) error { return nil })
|
||||
a.OnDrain("blocked", func(context.Context) error { return context.DeadlineExceeded })
|
||||
a.OnShutdown(func() error { closed = true; return nil })
|
||||
if err := a.Close(); !errors.Is(err, lifecycle.ErrResourcesRetained) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if closed {
|
||||
t.Fatal("unsafe cleanup")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMountRejectsNative(t *testing.T) {
|
||||
a := omni.New(omni.WithTransport(omni.TransportNative))
|
||||
defer a.Close()
|
||||
s, _ := ws.NewServer(ws.ServerConfig{})
|
||||
defer s.Shutdown(context.Background())
|
||||
if err := wsomni.Mount(a, "/ws", s, ws.Handlers{}); !errors.Is(err, ws.ErrUnsupportedTransport) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
98
ws/rpc/README.md
Normal file
98
ws/rpc/README.md
Normal file
@@ -0,0 +1,98 @@
|
||||
# WebSocket 双向调用与订阅
|
||||
|
||||
`ws/rpc` 是可选的 JSON 消息协议。普通 `ws` 的文本/二进制消息不依赖此协议。双端必须协商 `rpc.Subprotocol`（`go-omni.rpc.v1`），客户端和服务端都可发起调用、处理调用和发送事件。它不承诺离线消息、断线重放、事务回滚或 exactly-once。
|
||||
|
||||
## 调用
|
||||
|
||||
```go
|
||||
peer, err := rpc.New(rpc.Config{})
|
||||
if err != nil { return err }
|
||||
err = peer.Handle("echo", func(ctx context.Context, params json.RawMessage) (any, error) {
|
||||
return params, nil
|
||||
})
|
||||
if err != nil { return err }
|
||||
client, err := ws.NewClient(ws.ClientConfig{
|
||||
Subprotocols: []string{rpc.Subprotocol}, RequireSubprotocol: true,
|
||||
})
|
||||
if err != nil { return err }
|
||||
session, err := client.Dial(ctx, endpoint, peer.Handlers())
|
||||
if err != nil { return err }
|
||||
if err := peer.WaitReady(ctx); err != nil { return err }
|
||||
var result map[string]any
|
||||
err = peer.Call(ctx, "echo", map[string]any{"message": "hello"}, &result)
|
||||
// Notify 确认本地写入完成，不等待对端处理，也不返回业务执行结果。
|
||||
err = peer.Notify(ctx, "changed", map[string]any{"revision": 2})
|
||||
// 退出时使用有界关闭 context，并检查返回值。
|
||||
err = client.Shutdown(shutdownCtx)
|
||||
_ = session
|
||||
```
|
||||
|
||||
双向调用允许 Handler 内发起另一个 `Call`。接收路径立即处理回复，把请求分派到有上限的任务中，避免串行消息回调形成双向等待。路由在连接建立前通过 `Handle` 注册，重复方法或连接后注册返回错误。方法最多 128 字节，使用 ASCII 字母、数字及 `._-/:`。
|
||||
|
||||
服务端用 `ws.ServerConfig.Prepare` 为每次握手创建独立 Peer 并返回 `peer.Handlers()`；客户端自动重连用 `ws.ClientConfig.Prepare` 为每次连接创建独立 Peer。旧 Peer 永不绑定新连接。将新 Peer 交给应用后，先 `WaitReady`，再发起调用或恢复订阅。**不要在 `OnOpen` 中同步 Call/Restore**，因为 Session 在 OnOpen 返回后才分派回复。
|
||||
|
||||
完整真实 TCP 示例和退出流程见 [`examples/websocketRPC`](../../examples/websocketRPC)。
|
||||
|
||||
## 协议和错误
|
||||
|
||||
```json
|
||||
{"v":1,"type":"call","id":"1","method":"echo","params":{"message":"hello"},"timeoutMs":30000}
|
||||
{"v":1,"type":"reply","id":"1","result":{"message":"hello"}}
|
||||
{"v":1,"type":"reply","id":"2","error":{"code":"DENIED","message":"access denied"}}
|
||||
{"v":1,"type":"event","method":"changed","params":{"revision":2}}
|
||||
{"v":1,"type":"cancel","id":"1"}
|
||||
```
|
||||
|
||||
一条 WS 文本消息对应一个 envelope。协议版本、类型、字段组合、方法名和请求 ID 都校验；未知字段、重复 JSON 键（包括嵌套 params/result）、超过 32 层的 JSON、无效 UTF-8、二进制消息及超限消息被拒绝。收到重复的在途请求 ID 视为协议错误；迟到或重复回复丢弃，不创建新的 pending。未知调用回复 `METHOD_NOT_FOUND`；未知事件为协议错误，会终止连接。
|
||||
|
||||
`Call` 的执行预算取调用方 context 与 `Config.CallTimeout` 的较小值，预算毫秒数发送给对端，接收方再次按自己的配置裁剪。取消时尽力发送 cancel；取消消息入队失败不会无限等待，接收方仍有自己的超时。超时不表示远端业务已撤销。`Call` 与 `Notify` 都不会自动重试。
|
||||
|
||||
仅显式返回 `*rpc.Error{Code, Message}` 才会向对端发送自定义文本。普通错误、编码错误和 panic 使用通用 `INTERNAL`，不发送实现细节。超时/取消分别映射为 `DEADLINE_EXCEEDED` / `CANCELED`。事件处理无响应，可通过 `Config.OnEventError` 在本地观察失败；回调纳入排空等待。对端的公开错误仍是外部输入，日志不要自动输出其正文。
|
||||
|
||||
## 有界执行和退出
|
||||
|
||||
| 配置 | 默认 | 含义 |
|
||||
| --- | --- | --- |
|
||||
| `MaxMessageBytes` | 1 MiB | 完整 JSON envelope 的大小上限，应与底层 WS 消息限额匹配 |
|
||||
| `MaxPending` | 128 | 本端等待响应的调用上限 |
|
||||
| `MaxConcurrent` | 32 | 同时执行的入站调用及事件总数 |
|
||||
| `CallTimeout` | 30s | 入站业务及出站 Call 的最大预算 |
|
||||
| `WriteTimeout` | 5s | 事件和响应的写入期限 |
|
||||
|
||||
超出 pending 返回 `ErrBusy`；对端超出入站执行容量时，调用收到 `BUSY`，事件导致连接关闭，避免静默丢失。handler 必须响应 context；Go 不能强杀不退出的业务 goroutine。`Stop` 停止接收协议工作并取消 pending/incoming；`Wait` 等待所有入站 handler 和事件错误回调实际退出；`Shutdown` 额外发起 Session 关闭，但 Session 的最终关闭结果仍通过其 `Wait` 或所属 Client/Server.Shutdown 检查。
|
||||
|
||||
`Peer.Handlers().OnClose` 将上述任务接入 Session 完成屏障；若自定义包装 Handlers，必须保留这一回调。不要从本 Peer 的 Handler 内调用 `Wait/Shutdown`，可以调用 `Stop` 请求退出。关闭超时后，按框架生命周期规则保留仍被业务任务使用的依赖。
|
||||
|
||||
## 订阅和重连恢复
|
||||
|
||||
`Subscriptions` 管理本地**期望订阅**；服务端使用普通 RPC 路由处理加入/退出，执行自己的授权并可结合 `ws/hub`。默认方法为 `rpc.subscribe` / `rpc.unsubscribe`，可配置；此模块不预设用户、租户或主题权限。
|
||||
|
||||
```go
|
||||
subscriptions, err := rpc.NewSubscriptions(rpc.SubscriptionConfig{})
|
||||
if err != nil { return err }
|
||||
// nil Peer 仅保存期望状态，适合首次连接前配置。
|
||||
err = subscriptions.Subscribe(ctx, nil, "updates", map[string]any{"filter": "all"})
|
||||
// 每次连接创建新 Peer，WaitReady 后从 OnOpen 之外显式恢复。
|
||||
restored, err := subscriptions.Restore(ctx, newPeer)
|
||||
// 在线新增/移除同时发出 RPC 并等待确认。
|
||||
err = subscriptions.Subscribe(ctx, newPeer, "metrics", nil)
|
||||
err = subscriptions.Unsubscribe(ctx, newPeer, "updates")
|
||||
keys, err := subscriptions.Keys(ctx)
|
||||
_, _ = restored, keys
|
||||
```
|
||||
|
||||
订阅请求形如 `{"key":"updates","params":{...}}`；取消请求形如 `{"key":"updates"}`。服务端应将 key 的授权绑定到已验证的连接身份，每次恢复重新验证。相同连接中同一 key 的订阅应幂等；同 key 参数改变表示更新期望条件，由服务端定义其通用处理规则。
|
||||
|
||||
本地 Subscribe 在发送前保存期望状态，即使远端拒绝仍保留，应用必须检查错误并决定移除或以后重试；Unsubscribe 在发送前移除本地期望状态，若确认失败，当前连接仍可能收到事件，直至远端移除或连接关闭。Restore 按 key 排序，遇首个错误停止并返回已确认数量。操作串行且等待可通过 context 取消，防止并发恢复和移除颠倒状态。
|
||||
|
||||
默认最多 256 个订阅，每个 params 最多 16 KiB。key 是最多 256 字节的 UTF-8 不透明字符串，禁止 CR/LF/NUL。恢复只恢复订阅关系，断线期间事件不会补发。若业务需要历史事件，使用业务游标和持久化消息机制另行实现。
|
||||
|
||||
## 验证
|
||||
|
||||
```sh
|
||||
go test -race ./ws/rpc
|
||||
go test ./ws/rpc -run '^$' -fuzz FuzzProtocolDecode -fuzztime=10s
|
||||
go test ./ws/rpc -run '^$' -bench . -benchmem
|
||||
```
|
||||
|
||||
单元测试覆盖双向嵌套调用、并发上限、取消、迟到回复、重复键/层数/长度、错误脱敏、业务退出屏障、订阅恢复和真实 TCP 握手/关闭。`BenchmarkPeerCall` 是内存传输上的协议开销基准，不代表真实网络吞吐；`BenchmarkProtocolDecode` 测协议解析与严格输入校验。
|
||||
67
ws/rpc/config.go
Normal file
67
ws/rpc/config.go
Normal file
@@ -0,0 +1,67 @@
|
||||
// Package rpc implements an optional, versioned JSON protocol over ws sessions.
|
||||
// Both endpoints may issue calls and events. It does not provide durable delivery.
|
||||
package rpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
)
|
||||
|
||||
const Subprotocol = "go-omni.rpc.v1"
|
||||
|
||||
var (
|
||||
ErrClosed = errors.New("rpc: closed")
|
||||
ErrNotConnected = errors.New("rpc: not connected")
|
||||
ErrBusy = errors.New("rpc: capacity reached")
|
||||
ErrProtocol = errors.New("rpc: invalid protocol message")
|
||||
)
|
||||
|
||||
// Config bounds allocations, concurrency, and the lifetime of incoming work.
|
||||
// Zero fields select defaults. A canceled handler must return before Wait succeeds.
|
||||
type Config struct {
|
||||
MaxMessageBytes int `json:"maxMessageBytes"`
|
||||
MaxPending int `json:"maxPending"`
|
||||
MaxConcurrent int `json:"maxConcurrent"`
|
||||
CallTimeout time.Duration `json:"callTimeout"`
|
||||
WriteTimeout time.Duration `json:"writeTimeout"`
|
||||
// OnEventError observes local notification failures; they have no reply on
|
||||
// the wire. This callback runs within the handler completion barrier.
|
||||
OnEventError func(context.Context, error) `json:"-"`
|
||||
}
|
||||
|
||||
func DefaultConfig() Config {
|
||||
return Config{MaxMessageBytes: 1 << 20, MaxPending: 128, MaxConcurrent: 32, CallTimeout: 30 * time.Second, WriteTimeout: 5 * time.Second}
|
||||
}
|
||||
func (c Config) Validate() error { _, err := normalize(c); return err }
|
||||
func normalize(c Config) (Config, error) {
|
||||
if c.MaxMessageBytes < 0 || c.MaxMessageBytes > 16<<20 || c.MaxPending < 0 || c.MaxPending > 65536 || c.MaxConcurrent < 0 || c.MaxConcurrent > 4096 || c.CallTimeout < 0 || c.CallTimeout > 24*time.Hour || c.WriteTimeout < 0 || c.WriteTimeout > time.Minute {
|
||||
return c, errors.New("rpc: configuration exceeds supported bounds")
|
||||
}
|
||||
d := DefaultConfig()
|
||||
if c.MaxMessageBytes == 0 {
|
||||
c.MaxMessageBytes = d.MaxMessageBytes
|
||||
}
|
||||
if c.MaxPending == 0 {
|
||||
c.MaxPending = d.MaxPending
|
||||
}
|
||||
if c.MaxConcurrent == 0 {
|
||||
c.MaxConcurrent = d.MaxConcurrent
|
||||
}
|
||||
if c.CallTimeout == 0 {
|
||||
c.CallTimeout = d.CallTimeout
|
||||
}
|
||||
if c.WriteTimeout == 0 {
|
||||
c.WriteTimeout = d.WriteTimeout
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// Error is an explicitly public error. Only this type's code/message may be
|
||||
// sent to the other endpoint; all other handler errors become INTERNAL.
|
||||
type Error struct {
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
func (e *Error) Error() string { return "rpc: " + e.Code }
|
||||
424
ws/rpc/peer.go
Normal file
424
ws/rpc/peer.go
Normal file
@@ -0,0 +1,424 @@
|
||||
package rpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"git.yuntuo.hk/k2safe/go-omni/ws"
|
||||
)
|
||||
|
||||
type Handler func(context.Context, json.RawMessage) (any, error)
|
||||
|
||||
type transport interface {
|
||||
Context() context.Context
|
||||
Send(context.Context, ws.Message) error
|
||||
TrySend(ws.Message) error
|
||||
Close(int, string) error
|
||||
Abort()
|
||||
}
|
||||
|
||||
// Peer belongs to exactly one physical session. Construct a new Peer on every
|
||||
// reconnect. Register routes before passing Handlers to a client/server.
|
||||
type Peer struct {
|
||||
mu sync.Mutex
|
||||
cfg Config
|
||||
conn transport
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
routes map[string]Handler
|
||||
pending map[string]chan envelope
|
||||
incoming map[string]context.CancelFunc
|
||||
seq uint64
|
||||
active int
|
||||
closed bool
|
||||
done chan struct{}
|
||||
ready chan struct{}
|
||||
doneOnce sync.Once
|
||||
}
|
||||
|
||||
func New(cfg Config) (*Peer, error) {
|
||||
cfg, err := normalize(cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
return &Peer{cfg: cfg, ctx: ctx, cancel: cancel, routes: make(map[string]Handler), pending: make(map[string]chan envelope), incoming: make(map[string]context.CancelFunc), done: make(chan struct{}), ready: make(chan struct{})}, nil
|
||||
}
|
||||
|
||||
func (p *Peer) Handle(method string, handler Handler) error {
|
||||
if !validName(method, 128) || handler == nil {
|
||||
return errors.New("rpc: valid method and handler required")
|
||||
}
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if p.closed {
|
||||
return ErrClosed
|
||||
}
|
||||
if p.conn != nil {
|
||||
return errors.New("rpc: register methods before binding")
|
||||
}
|
||||
if _, ok := p.routes[method]; ok {
|
||||
return errors.New("rpc: duplicate method")
|
||||
}
|
||||
p.routes[method] = handler
|
||||
return nil
|
||||
}
|
||||
|
||||
// Handlers integrates protocol work into the Session completion barrier.
|
||||
// OnClose waits for all accepted handlers, including handlers ignoring cancel.
|
||||
// Owners must use a bounded Session.Wait/Shutdown and retain dependencies if it
|
||||
// times out. Do not call Wait or Shutdown from a registered handler itself.
|
||||
func (p *Peer) Handlers() ws.Handlers {
|
||||
return ws.Handlers{
|
||||
OnOpen: func(_ context.Context, s *ws.Session) error {
|
||||
if s.Subprotocol() != Subprotocol {
|
||||
return ErrProtocol
|
||||
}
|
||||
return p.bind(s)
|
||||
},
|
||||
OnMessage: func(_ context.Context, _ *ws.Session, m ws.Message) error { return p.receive(m) },
|
||||
OnClose: func(_ *ws.Session, _ error) { p.Stop(); _ = p.Wait(context.Background()) },
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Peer) bind(c transport) error {
|
||||
p.mu.Lock()
|
||||
if p.closed {
|
||||
p.mu.Unlock()
|
||||
return ErrClosed
|
||||
}
|
||||
if p.conn != nil {
|
||||
p.mu.Unlock()
|
||||
return errors.New("rpc: peer already bound")
|
||||
}
|
||||
p.cancel()
|
||||
p.ctx, p.cancel = context.WithCancel(c.Context())
|
||||
p.conn = c
|
||||
close(p.ready)
|
||||
ctx := p.ctx
|
||||
p.mu.Unlock()
|
||||
context.AfterFunc(ctx, p.Stop)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop cancels pending calls and incoming handlers without waiting for them.
|
||||
// It does not own the underlying Session; Shutdown additionally closes it.
|
||||
func (p *Peer) Stop() {
|
||||
p.mu.Lock()
|
||||
if !p.closed {
|
||||
p.closed = true
|
||||
p.cancel()
|
||||
for _, cancel := range p.incoming {
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
if p.active == 0 {
|
||||
p.doneOnce.Do(func() { close(p.done) })
|
||||
}
|
||||
p.mu.Unlock()
|
||||
}
|
||||
func (p *Peer) Done() <-chan struct{} { return p.done }
|
||||
|
||||
// WaitReady waits for OnOpen to bind the transport. Call and Restore should run
|
||||
// outside OnOpen: replies are delivered only after that callback returns.
|
||||
func (p *Peer) WaitReady(ctx context.Context) error {
|
||||
select {
|
||||
case <-p.ready:
|
||||
p.mu.Lock()
|
||||
closed := p.closed
|
||||
p.mu.Unlock()
|
||||
if closed {
|
||||
return ErrClosed
|
||||
}
|
||||
return nil
|
||||
case <-p.done:
|
||||
return ErrClosed
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
func (p *Peer) Wait(ctx context.Context) error {
|
||||
select {
|
||||
case <-p.done:
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
select {
|
||||
case <-p.done:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
func (p *Peer) Shutdown(ctx context.Context) error {
|
||||
p.Stop()
|
||||
p.mu.Lock()
|
||||
c := p.conn
|
||||
p.mu.Unlock()
|
||||
if c != nil {
|
||||
_ = c.Close(1000, "")
|
||||
}
|
||||
return p.Wait(ctx)
|
||||
}
|
||||
|
||||
// Call waits for the peer's reply, bounded by ctx and Config.CallTimeout.
|
||||
// Cancellation is best effort on the wire; a timeout never proves rollback.
|
||||
func (p *Peer) Call(ctx context.Context, method string, input, out any) error {
|
||||
if !validName(method, 128) {
|
||||
return errors.New("rpc: invalid method")
|
||||
}
|
||||
params, err := json.Marshal(input)
|
||||
if err != nil {
|
||||
return fmt.Errorf("rpc: encode input: %w", err)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, p.cfg.CallTimeout)
|
||||
defer cancel()
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
p.mu.Lock()
|
||||
if p.closed {
|
||||
p.mu.Unlock()
|
||||
return ErrClosed
|
||||
}
|
||||
if p.conn == nil {
|
||||
p.mu.Unlock()
|
||||
return ErrNotConnected
|
||||
}
|
||||
if len(p.pending) >= p.cfg.MaxPending {
|
||||
p.mu.Unlock()
|
||||
return ErrBusy
|
||||
}
|
||||
p.seq++
|
||||
id := strconv.FormatUint(p.seq, 36)
|
||||
ch := make(chan envelope, 1)
|
||||
p.pending[id] = ch
|
||||
c, work := p.conn, p.ctx
|
||||
p.mu.Unlock()
|
||||
defer func() { p.mu.Lock(); delete(p.pending, id); p.mu.Unlock() }()
|
||||
deadline, _ := ctx.Deadline()
|
||||
budget := max(int64(1), time.Until(deadline).Milliseconds())
|
||||
if err = p.send(ctx, c, envelope{Version: 1, Type: "call", ID: id, Method: method, Params: params, TimeoutMS: budget}); err != nil {
|
||||
return err
|
||||
}
|
||||
select {
|
||||
case reply := <-ch:
|
||||
if reply.Error != nil {
|
||||
return reply.Error
|
||||
}
|
||||
if out == nil {
|
||||
return nil
|
||||
}
|
||||
return json.Unmarshal(reply.Result, out)
|
||||
case <-ctx.Done():
|
||||
// Never block a caller waiting for room to send cancellation. The remote
|
||||
// handler is independently bounded by the transmitted deadline.
|
||||
_ = p.try(c, envelope{Version: 1, Type: "cancel", ID: id})
|
||||
return ctx.Err()
|
||||
case <-work.Done():
|
||||
return ErrClosed
|
||||
}
|
||||
}
|
||||
|
||||
// Notify confirms a local write, not remote processing. Notifications are
|
||||
// handled with the same concurrency and timeout limits as incoming calls.
|
||||
func (p *Peer) Notify(ctx context.Context, method string, input any) error {
|
||||
if !validName(method, 128) {
|
||||
return errors.New("rpc: invalid method")
|
||||
}
|
||||
params, err := json.Marshal(input)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
p.mu.Lock()
|
||||
c, closed := p.conn, p.closed
|
||||
p.mu.Unlock()
|
||||
if closed {
|
||||
return ErrClosed
|
||||
}
|
||||
if c == nil {
|
||||
return ErrNotConnected
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, p.cfg.WriteTimeout)
|
||||
defer cancel()
|
||||
return p.send(ctx, c, envelope{Version: 1, Type: "event", Method: method, Params: params})
|
||||
}
|
||||
|
||||
func (p *Peer) encode(e envelope) (ws.Message, error) {
|
||||
data, err := json.Marshal(e)
|
||||
if err != nil {
|
||||
return ws.Message{}, err
|
||||
}
|
||||
if len(data) > p.cfg.MaxMessageBytes {
|
||||
return ws.Message{}, ws.ErrMessageTooLarge
|
||||
}
|
||||
return ws.Message{Type: ws.Text, Data: data}, nil
|
||||
}
|
||||
func (p *Peer) send(ctx context.Context, c transport, e envelope) error {
|
||||
m, err := p.encode(e)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return c.Send(ctx, m)
|
||||
}
|
||||
func (p *Peer) try(c transport, e envelope) error {
|
||||
m, err := p.encode(e)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return c.TrySend(m)
|
||||
}
|
||||
|
||||
func (p *Peer) receive(m ws.Message) error {
|
||||
if m.Type != ws.Text {
|
||||
return ErrProtocol
|
||||
}
|
||||
e, err := decode(m.Data, p.cfg.MaxMessageBytes)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
p.mu.Lock()
|
||||
if p.closed {
|
||||
p.mu.Unlock()
|
||||
return ErrClosed
|
||||
}
|
||||
if p.conn == nil {
|
||||
p.mu.Unlock()
|
||||
return ErrNotConnected
|
||||
}
|
||||
switch e.Type {
|
||||
case "reply":
|
||||
if ch := p.pending[e.ID]; ch != nil {
|
||||
select {
|
||||
case ch <- e:
|
||||
default:
|
||||
}
|
||||
}
|
||||
p.mu.Unlock()
|
||||
return nil
|
||||
case "cancel":
|
||||
if cancel := p.incoming[e.ID]; cancel != nil {
|
||||
cancel()
|
||||
}
|
||||
p.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
c := p.conn
|
||||
if e.Type == "call" {
|
||||
if _, exists := p.incoming[e.ID]; exists {
|
||||
p.mu.Unlock()
|
||||
return ErrProtocol
|
||||
}
|
||||
}
|
||||
handler := p.routes[e.Method]
|
||||
if handler == nil {
|
||||
p.mu.Unlock()
|
||||
if e.Type == "event" {
|
||||
return ErrProtocol
|
||||
}
|
||||
return p.try(c, envelope{Version: 1, Type: "reply", ID: e.ID, Error: &Error{"METHOD_NOT_FOUND", "method not found"}})
|
||||
}
|
||||
if p.active >= p.cfg.MaxConcurrent {
|
||||
p.mu.Unlock()
|
||||
if e.Type == "event" {
|
||||
return ErrBusy
|
||||
}
|
||||
return p.try(c, envelope{Version: 1, Type: "reply", ID: e.ID, Error: &Error{"BUSY", "peer capacity reached"}})
|
||||
}
|
||||
timeout := p.cfg.CallTimeout
|
||||
if e.TimeoutMS > 0 {
|
||||
timeout = min(timeout, time.Duration(e.TimeoutMS)*time.Millisecond)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(p.ctx, timeout)
|
||||
if e.Type == "call" {
|
||||
p.incoming[e.ID] = cancel
|
||||
}
|
||||
p.active++
|
||||
p.mu.Unlock()
|
||||
go p.execute(ctx, cancel, c, e, handler)
|
||||
return nil
|
||||
}
|
||||
|
||||
func invoke(ctx context.Context, h Handler, params json.RawMessage) (value any, err error) {
|
||||
defer func() {
|
||||
if recover() != nil {
|
||||
value = nil
|
||||
err = errors.New("rpc: handler panic")
|
||||
}
|
||||
}()
|
||||
return h(ctx, params)
|
||||
}
|
||||
func publicError(err error) (result *Error) {
|
||||
result = &Error{"INTERNAL", "request failed"}
|
||||
defer func() {
|
||||
if recover() != nil {
|
||||
result = &Error{"INTERNAL", "request failed"}
|
||||
}
|
||||
}()
|
||||
if errors.Is(err, context.DeadlineExceeded) {
|
||||
return &Error{"DEADLINE_EXCEEDED", "request deadline exceeded"}
|
||||
}
|
||||
if errors.Is(err, context.Canceled) {
|
||||
return &Error{"CANCELED", "request canceled"}
|
||||
}
|
||||
var exposed *Error
|
||||
if errors.As(err, &exposed) && validError(exposed) {
|
||||
return &Error{exposed.Code, exposed.Message}
|
||||
}
|
||||
return &Error{"INTERNAL", "request failed"}
|
||||
}
|
||||
|
||||
func marshalResult(value any) (raw json.RawMessage, err error) {
|
||||
defer func() {
|
||||
if recover() != nil {
|
||||
raw = nil
|
||||
err = errors.New("rpc: result encoding failed")
|
||||
}
|
||||
}()
|
||||
return json.Marshal(value)
|
||||
}
|
||||
func (p *Peer) execute(ctx context.Context, cancel context.CancelFunc, c transport, e envelope, h Handler) {
|
||||
defer func() {
|
||||
cancel()
|
||||
p.mu.Lock()
|
||||
delete(p.incoming, e.ID)
|
||||
p.active--
|
||||
if p.closed && p.active == 0 {
|
||||
p.doneOnce.Do(func() { close(p.done) })
|
||||
}
|
||||
p.mu.Unlock()
|
||||
}()
|
||||
value, err := invoke(ctx, h, e.Params)
|
||||
if e.Type == "event" {
|
||||
if err != nil && p.cfg.OnEventError != nil {
|
||||
func() { defer func() { _ = recover() }(); p.cfg.OnEventError(ctx, err) }()
|
||||
}
|
||||
return
|
||||
}
|
||||
if err == nil {
|
||||
err = ctx.Err()
|
||||
}
|
||||
reply := envelope{Version: 1, Type: "reply", ID: e.ID}
|
||||
if err == nil {
|
||||
reply.Result, err = marshalResult(value)
|
||||
}
|
||||
if err != nil {
|
||||
reply.Result = nil
|
||||
reply.Error = publicError(err)
|
||||
}
|
||||
writeCtx, writeCancel := context.WithTimeout(p.ctx, p.cfg.WriteTimeout)
|
||||
defer writeCancel()
|
||||
if err = p.send(writeCtx, c, reply); errors.Is(err, ws.ErrMessageTooLarge) {
|
||||
reply.Result = nil
|
||||
reply.Error = &Error{"RESULT_TOO_LARGE", "result exceeds message limit"}
|
||||
err = p.send(writeCtx, c, reply)
|
||||
}
|
||||
if err != nil {
|
||||
c.Abort()
|
||||
}
|
||||
}
|
||||
409
ws/rpc/peer_test.go
Normal file
409
ws/rpc/peer_test.go
Normal file
@@ -0,0 +1,409 @@
|
||||
package rpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.yuntuo.hk/k2safe/go-omni/ws"
|
||||
)
|
||||
|
||||
type memoryTransport struct {
|
||||
ctx context.Context
|
||||
other *Peer
|
||||
}
|
||||
|
||||
func (m *memoryTransport) Context() context.Context { return m.ctx }
|
||||
func (m *memoryTransport) Send(ctx context.Context, msg ws.Message) error {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
return m.other.receive(msg)
|
||||
}
|
||||
func (m *memoryTransport) TrySend(msg ws.Message) error { return m.other.receive(msg) }
|
||||
func (m *memoryTransport) Close(int, string) error { return nil }
|
||||
func (m *memoryTransport) Abort() { m.other.Stop() }
|
||||
|
||||
func pair(t testing.TB, cfg Config, setup func(*Peer, *Peer)) (*Peer, *Peer) {
|
||||
t.Helper()
|
||||
a, err := New(cfg)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
b, err := New(cfg)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if setup != nil {
|
||||
setup(a, b)
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
if err = a.bind(&memoryTransport{ctx, b}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = b.bind(&memoryTransport{ctx, a}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
a.Stop()
|
||||
b.Stop()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := a.Wait(ctx); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
if err := b.Wait(ctx); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
})
|
||||
return a, b
|
||||
}
|
||||
func handle(t testing.TB, p *Peer, name string, h Handler) {
|
||||
t.Helper()
|
||||
if err := p.Handle(name, h); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBidirectionalNestedCalls(t *testing.T) {
|
||||
a, _ := pair(t, Config{}, func(a, b *Peer) {
|
||||
handle(t, a, "double", func(_ context.Context, v json.RawMessage) (any, error) {
|
||||
var n int
|
||||
_ = json.Unmarshal(v, &n)
|
||||
return n * 2, nil
|
||||
})
|
||||
handle(t, b, "nested", func(ctx context.Context, v json.RawMessage) (any, error) {
|
||||
var n int
|
||||
_ = json.Unmarshal(v, &n)
|
||||
var doubled int
|
||||
err := b.Call(ctx, "double", n, &doubled)
|
||||
return doubled + 1, err
|
||||
})
|
||||
})
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 20; i++ {
|
||||
wg.Add(1)
|
||||
go func(n int) {
|
||||
defer wg.Done()
|
||||
var result int
|
||||
if err := a.Call(context.Background(), "nested", n, &result); err != nil || result != n*2+1 {
|
||||
t.Errorf("result %d, error %v", result, err)
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestCancellationBusyAndCompletionBarrier(t *testing.T) {
|
||||
entered := make(chan struct{})
|
||||
exit := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
a, b := pair(t, Config{MaxConcurrent: 1}, func(a, b *Peer) {
|
||||
handle(t, b, "block", func(ctx context.Context, _ json.RawMessage) (any, error) {
|
||||
close(entered)
|
||||
<-ctx.Done()
|
||||
close(exit)
|
||||
<-release
|
||||
return nil, ctx.Err()
|
||||
})
|
||||
})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
result := make(chan error, 1)
|
||||
go func() { result <- a.Call(ctx, "block", nil, nil) }()
|
||||
<-entered
|
||||
if err := a.Call(context.Background(), "block", nil, nil); err == nil {
|
||||
t.Fatal("capacity not enforced")
|
||||
} else {
|
||||
var pub *Error
|
||||
if !errors.As(err, &pub) || pub.Code != "BUSY" {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := <-result; !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
<-exit
|
||||
b.Stop()
|
||||
wait, c := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer c()
|
||||
if err := b.Wait(wait); !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatal("handler not joined", err)
|
||||
}
|
||||
close(release)
|
||||
if err := b.Wait(context.Background()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestErrorsDoNotExposeImplementation(t *testing.T) {
|
||||
a, _ := pair(t, Config{}, func(a, b *Peer) {
|
||||
handle(t, b, "private", func(context.Context, json.RawMessage) (any, error) { return nil, errors.New("password=secret") })
|
||||
handle(t, b, "panic", func(context.Context, json.RawMessage) (any, error) { panic("token=private") })
|
||||
handle(t, b, "public", func(context.Context, json.RawMessage) (any, error) {
|
||||
return nil, &Error{"DENIED", "public explanation"}
|
||||
})
|
||||
})
|
||||
for _, name := range []string{"private", "panic", "missing", "public"} {
|
||||
err := a.Call(context.Background(), name, nil, nil)
|
||||
var pub *Error
|
||||
if !errors.As(err, &pub) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.Contains(pub.Message, "secret") || strings.Contains(pub.Message, "private") {
|
||||
t.Fatal("leaked error", pub.Message)
|
||||
}
|
||||
if name == "public" && pub.Message != "public explanation" {
|
||||
t.Fatal(pub)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestReplyLateDuplicateAndProtocolLimits(t *testing.T) {
|
||||
p, _ := New(Config{})
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
q, _ := New(Config{})
|
||||
_ = p.bind(&memoryTransport{ctx, q})
|
||||
defer p.Stop()
|
||||
// Late and duplicated replies must not block the reader or create allocations.
|
||||
for i := 0; i < 10; i++ {
|
||||
if err := p.receive(ws.Message{Type: ws.Text, Data: []byte(`{"v":1,"type":"reply","id":"old","result":null}`)}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
invalid := []string{
|
||||
`{"v":2,"type":"event","method":"x","params":null}`,
|
||||
`{"v":1,"type":"call","id":"x","id":"y","method":"x","params":null}`,
|
||||
`{"v":1,"type":"event","method":"x","params":{"role":1,"role":2}}`,
|
||||
`{"v":1,"type":"reply","id":"a","result":null,"error":{"code":"BAD","message":"x"}}`,
|
||||
`{"v":1,"type":"cancel","id":"a","params":null}`,
|
||||
`{"v":1,"type":"call","id":"x","method":"x","params":null,"extra":true}`,
|
||||
`{"v":1,"type":"event","method":"x","params":` + strings.Repeat("[", 34) + "0" + strings.Repeat("]", 34) + `}`,
|
||||
}
|
||||
for _, data := range invalid {
|
||||
if _, err := decode([]byte(data), 1<<20); !errors.Is(err, ErrProtocol) {
|
||||
t.Fatalf("accepted %s", data)
|
||||
}
|
||||
}
|
||||
if _, err := decode([]byte(`{"v":1,"type":"event","method":"x","params":null}`), 4); err == nil {
|
||||
t.Fatal("size ignored")
|
||||
}
|
||||
if err := p.receive(ws.Message{Type: ws.Binary, Data: []byte("x")}); !errors.Is(err, ErrProtocol) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPendingLimitAndDisconnect(t *testing.T) {
|
||||
entered := make(chan struct{})
|
||||
a, b := pair(t, Config{MaxPending: 1}, func(a, b *Peer) {
|
||||
handle(t, b, "wait", func(ctx context.Context, _ json.RawMessage) (any, error) {
|
||||
close(entered)
|
||||
<-ctx.Done()
|
||||
return nil, ctx.Err()
|
||||
})
|
||||
})
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- a.Call(context.Background(), "wait", nil, nil) }()
|
||||
<-entered
|
||||
if err := a.Call(context.Background(), "wait", nil, nil); !errors.Is(err, ErrBusy) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
a.Stop()
|
||||
b.Stop()
|
||||
if err := <-done; !errors.Is(err, ErrClosed) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := a.bind(&memoryTransport{context.Background(), b}); !errors.Is(err, ErrClosed) {
|
||||
t.Fatal("peer rebound", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubscriptionsRestoreAndRemoval(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
var keys []string
|
||||
setup := func(a, b *Peer) {
|
||||
handle(t, b, "rpc.subscribe", func(_ context.Context, raw json.RawMessage) (any, error) {
|
||||
var sub Subscription
|
||||
if err := json.Unmarshal(raw, &sub); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
mu.Lock()
|
||||
keys = append(keys, sub.Key)
|
||||
mu.Unlock()
|
||||
return nil, nil
|
||||
})
|
||||
handle(t, b, "rpc.unsubscribe", func(context.Context, json.RawMessage) (any, error) { return nil, nil })
|
||||
}
|
||||
a, _ := pair(t, Config{}, setup)
|
||||
s, err := NewSubscriptions(SubscriptionConfig{MaxSubscriptions: 2})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = s.Subscribe(context.Background(), nil, "b", map[string]int{"n": 1})
|
||||
_ = s.Subscribe(context.Background(), nil, "a", nil)
|
||||
if err = s.Subscribe(context.Background(), nil, "c", nil); !errors.Is(err, ErrBusy) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n, err := s.Restore(context.Background(), a); n != 2 || err != nil {
|
||||
t.Fatal(n, err)
|
||||
}
|
||||
if err = s.Unsubscribe(context.Background(), a, "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
next, _ := pair(t, Config{}, setup)
|
||||
if n, err := s.Restore(context.Background(), next); n != 1 || err != nil {
|
||||
t.Fatal(n, err)
|
||||
}
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if strings.Join(keys, ",") != "a,b,b" {
|
||||
t.Fatal(keys)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRealTCPRPCAndShutdown(t *testing.T) {
|
||||
opened := make(chan *Peer, 1)
|
||||
server, err := ws.NewServer(ws.ServerConfig{Subprotocols: []string{Subprotocol}, RequireSubprotocol: true, Prepare: func(context.Context, *http.Request) (ws.Handlers, error) {
|
||||
p, _ := New(Config{})
|
||||
handle(t, p, "echo", func(_ context.Context, raw json.RawMessage) (any, error) { return raw, nil })
|
||||
opened <- p
|
||||
return p.Handlers(), nil
|
||||
}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
httpServer := httptest.NewServer(server.Handler(ws.Handlers{}))
|
||||
defer httpServer.Close()
|
||||
client, err := ws.NewClient(ws.ClientConfig{Subprotocols: []string{Subprotocol}, RequireSubprotocol: true, ProxyMode: "direct"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
_ = client.Shutdown(ctx)
|
||||
_ = server.Shutdown(ctx)
|
||||
}()
|
||||
p, _ := New(Config{})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
defer cancel()
|
||||
session, err := client.Dial(ctx, "ws"+strings.TrimPrefix(httpServer.URL, "http"), p.Handlers())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = p.WaitReady(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
remote := <-opened
|
||||
var result map[string]int
|
||||
if err = p.Call(ctx, "echo", map[string]int{"n": 42}, &result); err != nil || result["n"] != 42 {
|
||||
t.Fatal(result, err)
|
||||
}
|
||||
if err = p.Shutdown(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = session.Wait(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = remote.Wait(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkPeerCall(b *testing.B) {
|
||||
a, _ := pair(b, Config{}, func(a, p *Peer) {
|
||||
handle(b, p, "echo", func(_ context.Context, raw json.RawMessage) (any, error) { return raw, nil })
|
||||
})
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var out int
|
||||
if err := a.Call(context.Background(), "echo", 42, &out); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
func BenchmarkProtocolDecode(b *testing.B) {
|
||||
data := []byte(`{"v":1,"type":"call","id":"abc","method":"status.get","params":{"detail":true},"timeoutMs":1000}`)
|
||||
b.ReportAllocs()
|
||||
b.SetBytes(int64(len(data)))
|
||||
for i := 0; i < b.N; i++ {
|
||||
if _, err := decode(data, 1<<20); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
func FuzzProtocolDecode(f *testing.F) {
|
||||
for _, s := range []string{`{"v":1,"type":"event","method":"x","params":null}`, `{"v":1,"type":"reply","id":"x","result":42}`, `null`, `{"v":1,"v":2}`} {
|
||||
f.Add([]byte(s))
|
||||
}
|
||||
f.Fuzz(func(t *testing.T, data []byte) { _, _ = decode(data, 4096) })
|
||||
}
|
||||
|
||||
type panickingResult struct{}
|
||||
|
||||
func (panickingResult) MarshalJSON() ([]byte, error) { panic("secret encoder state") }
|
||||
func TestResultEncodingPanicAndOversize(t *testing.T) {
|
||||
a, _ := pair(t, Config{MaxMessageBytes: 256}, func(a, b *Peer) {
|
||||
handle(t, b, "panicResult", func(context.Context, json.RawMessage) (any, error) { return panickingResult{}, nil })
|
||||
handle(t, b, "large", func(context.Context, json.RawMessage) (any, error) { return strings.Repeat("x", 1024), nil })
|
||||
})
|
||||
for _, method := range []string{"panicResult", "large"} {
|
||||
err := a.Call(context.Background(), method, nil, nil)
|
||||
var e *Error
|
||||
if !errors.As(err, &e) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if method == "panicResult" && e.Code != "INTERNAL" {
|
||||
t.Fatal(e)
|
||||
}
|
||||
if method == "large" && e.Code != "RESULT_TOO_LARGE" {
|
||||
t.Fatal(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
func TestEventErrorObservation(t *testing.T) {
|
||||
observed := make(chan error, 1)
|
||||
a, _ := pair(t, Config{OnEventError: func(_ context.Context, err error) { observed <- err; panic("observer panic") }}, func(a, b *Peer) {
|
||||
handle(t, b, "fail", func(context.Context, json.RawMessage) (any, error) { return nil, errors.New("local failure") })
|
||||
})
|
||||
if err := a.Notify(context.Background(), "fail", nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
select {
|
||||
case err := <-observed:
|
||||
if err.Error() != "local failure" {
|
||||
t.Fatal(err)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("event error not observed")
|
||||
}
|
||||
}
|
||||
func TestDuplicateIncomingIDRejected(t *testing.T) {
|
||||
entered := make(chan struct{})
|
||||
a, b := pair(t, Config{}, func(a, b *Peer) {
|
||||
handle(t, b, "wait", func(ctx context.Context, _ json.RawMessage) (any, error) {
|
||||
close(entered)
|
||||
<-ctx.Done()
|
||||
return nil, ctx.Err()
|
||||
})
|
||||
})
|
||||
raw := ws.Message{Type: ws.Text, Data: []byte(`{"v":1,"type":"call","id":"same","method":"wait","params":null}`)}
|
||||
if err := b.receive(raw); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
<-entered
|
||||
if err := b.receive(raw); !errors.Is(err, ErrProtocol) {
|
||||
t.Fatal("duplicate active ID accepted", err)
|
||||
}
|
||||
a.Stop()
|
||||
b.Stop()
|
||||
}
|
||||
129
ws/rpc/protocol.go
Normal file
129
ws/rpc/protocol.go
Normal file
@@ -0,0 +1,129 @@
|
||||
package rpc
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
type envelope struct {
|
||||
Version int `json:"v"`
|
||||
Type string `json:"type"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Method string `json:"method,omitempty"`
|
||||
Params json.RawMessage `json:"params,omitempty"`
|
||||
Result json.RawMessage `json:"result,omitempty"`
|
||||
Error *Error `json:"error,omitempty"`
|
||||
TimeoutMS int64 `json:"timeoutMs,omitempty"`
|
||||
}
|
||||
|
||||
func validName(s string, max int) bool {
|
||||
if s == "" || len(s) > max {
|
||||
return false
|
||||
}
|
||||
for _, c := range s {
|
||||
if !(c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' || strings.ContainsRune("._-/:", c)) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
func validError(e *Error) bool {
|
||||
return e != nil && validName(e.Code, 64) && len(e.Message) <= 1024 && utf8.ValidString(e.Message)
|
||||
}
|
||||
|
||||
func decode(data []byte, limit int) (envelope, error) {
|
||||
var e envelope
|
||||
if len(data) > limit || !utf8.Valid(data) {
|
||||
return e, ErrProtocol
|
||||
}
|
||||
// Reject duplicate keys at every depth instead of accepting last-key-wins.
|
||||
scan := json.NewDecoder(bytes.NewReader(data))
|
||||
scan.UseNumber()
|
||||
if err := uniqueValue(scan, 0); err != nil {
|
||||
return e, ErrProtocol
|
||||
}
|
||||
if _, err := scan.Token(); err != io.EOF {
|
||||
return e, ErrProtocol
|
||||
}
|
||||
d := json.NewDecoder(bytes.NewReader(data))
|
||||
d.DisallowUnknownFields()
|
||||
if err := d.Decode(&e); err != nil {
|
||||
return e, ErrProtocol
|
||||
}
|
||||
if e.Version != 1 || e.TimeoutMS < 0 || e.TimeoutMS > int64((24*60*60)*1000) {
|
||||
return e, ErrProtocol
|
||||
}
|
||||
switch e.Type {
|
||||
case "call":
|
||||
if !validName(e.ID, 64) || !validName(e.Method, 128) || len(e.Params) == 0 || len(e.Result) != 0 || e.Error != nil {
|
||||
return e, ErrProtocol
|
||||
}
|
||||
case "event":
|
||||
if e.ID != "" || !validName(e.Method, 128) || len(e.Params) == 0 || len(e.Result) != 0 || e.Error != nil || e.TimeoutMS != 0 {
|
||||
return e, ErrProtocol
|
||||
}
|
||||
case "reply":
|
||||
if !validName(e.ID, 64) || e.Method != "" || len(e.Params) != 0 || e.TimeoutMS != 0 || (len(e.Result) == 0) == (e.Error == nil) || (e.Error != nil && !validError(e.Error)) {
|
||||
return e, ErrProtocol
|
||||
}
|
||||
case "cancel":
|
||||
if !validName(e.ID, 64) || e.Method != "" || len(e.Params) != 0 || len(e.Result) != 0 || e.Error != nil || e.TimeoutMS != 0 {
|
||||
return e, ErrProtocol
|
||||
}
|
||||
default:
|
||||
return e, ErrProtocol
|
||||
}
|
||||
return e, nil
|
||||
}
|
||||
|
||||
func uniqueValue(d *json.Decoder, depth int) error {
|
||||
if depth > 32 {
|
||||
return ErrProtocol
|
||||
}
|
||||
t, err := d.Token()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
delim, ok := t.(json.Delim)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
switch delim {
|
||||
case '{':
|
||||
seen := make(map[string]bool)
|
||||
for d.More() {
|
||||
k, err := d.Token()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
key, ok := k.(string)
|
||||
if !ok || seen[key] {
|
||||
return ErrProtocol
|
||||
}
|
||||
seen[key] = true
|
||||
if err := uniqueValue(d, depth+1); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
t, err = d.Token()
|
||||
if err != nil || t != json.Delim('}') {
|
||||
return ErrProtocol
|
||||
}
|
||||
case '[':
|
||||
for d.More() {
|
||||
if err := uniqueValue(d, depth+1); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
t, err = d.Token()
|
||||
if err != nil || t != json.Delim(']') {
|
||||
return ErrProtocol
|
||||
}
|
||||
default:
|
||||
return ErrProtocol
|
||||
}
|
||||
return nil
|
||||
}
|
||||
159
ws/rpc/subscriptions.go
Normal file
159
ws/rpc/subscriptions.go
Normal file
@@ -0,0 +1,159 @@
|
||||
package rpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"sort"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// Subscription is the generic request sent to the configured subscribe method.
|
||||
// The server must authenticate and authorize Key and Params itself.
|
||||
type Subscription struct {
|
||||
Key string `json:"key"`
|
||||
Params json.RawMessage `json:"params"`
|
||||
}
|
||||
|
||||
type SubscriptionConfig struct {
|
||||
MaxSubscriptions int `json:"maxSubscriptions"`
|
||||
MaxParamsBytes int `json:"maxParamsBytes"`
|
||||
SubscribeMethod string `json:"subscribeMethod"`
|
||||
UnsubscribeMethod string `json:"unsubscribeMethod"`
|
||||
}
|
||||
|
||||
// Subscriptions stores desired subscriptions, never received messages. Reuse it
|
||||
// across sessions and explicitly Restore after each new Peer is bound. Server
|
||||
// subscribe handlers must be idempotent for a key within a session; authorization
|
||||
// is checked anew on every restore. No missed-event replay is implied.
|
||||
type Subscriptions struct {
|
||||
cfg SubscriptionConfig
|
||||
gate chan struct{}
|
||||
desired map[string]json.RawMessage
|
||||
}
|
||||
|
||||
func validSubscriptionKey(key string) bool {
|
||||
return key != "" && len(key) <= 256 && utf8.ValidString(key) && !strings.ContainsAny(key, "\r\n\x00")
|
||||
}
|
||||
|
||||
func NewSubscriptions(cfg SubscriptionConfig) (*Subscriptions, error) {
|
||||
if cfg.MaxSubscriptions < 0 || cfg.MaxSubscriptions > 65536 || cfg.MaxParamsBytes < 0 || cfg.MaxParamsBytes > 1<<20 {
|
||||
return nil, errors.New("rpc: invalid subscription limits")
|
||||
}
|
||||
if cfg.MaxSubscriptions == 0 {
|
||||
cfg.MaxSubscriptions = 256
|
||||
}
|
||||
if cfg.MaxParamsBytes == 0 {
|
||||
cfg.MaxParamsBytes = 16 << 10
|
||||
}
|
||||
if cfg.SubscribeMethod == "" {
|
||||
cfg.SubscribeMethod = "rpc.subscribe"
|
||||
}
|
||||
if cfg.UnsubscribeMethod == "" {
|
||||
cfg.UnsubscribeMethod = "rpc.unsubscribe"
|
||||
}
|
||||
if !validName(cfg.SubscribeMethod, 128) || !validName(cfg.UnsubscribeMethod, 128) || cfg.SubscribeMethod == cfg.UnsubscribeMethod {
|
||||
return nil, errors.New("rpc: distinct subscription methods required")
|
||||
}
|
||||
return &Subscriptions{cfg: cfg, gate: make(chan struct{}, 1), desired: make(map[string]json.RawMessage)}, nil
|
||||
}
|
||||
func (s *Subscriptions) lock(ctx context.Context) error {
|
||||
select {
|
||||
case s.gate <- struct{}{}:
|
||||
if err := ctx.Err(); err != nil {
|
||||
<-s.gate
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
func (s *Subscriptions) unlock() { <-s.gate }
|
||||
|
||||
// Subscribe records the desired state before attempting the call. On failure,
|
||||
// it remains desired for a later Restore; the returned error must be inspected.
|
||||
// A nil peer stores the subscription for the next explicit Restore.
|
||||
func (s *Subscriptions) Subscribe(ctx context.Context, p *Peer, key string, params any) error {
|
||||
if !validSubscriptionKey(key) {
|
||||
return errors.New("rpc: invalid subscription key")
|
||||
}
|
||||
raw, err := json.Marshal(params)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(raw) > s.cfg.MaxParamsBytes {
|
||||
return errors.New("rpc: subscription params exceed limit")
|
||||
}
|
||||
if err = s.lock(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.unlock()
|
||||
if _, exists := s.desired[key]; !exists && len(s.desired) >= s.cfg.MaxSubscriptions {
|
||||
return ErrBusy
|
||||
}
|
||||
s.desired[key] = raw
|
||||
if p == nil {
|
||||
return nil
|
||||
}
|
||||
return p.Call(ctx, s.cfg.SubscribeMethod, Subscription{key, raw}, nil)
|
||||
}
|
||||
|
||||
// Unsubscribe removes desired state even if the remote acknowledgement fails.
|
||||
// The current connection may still receive events until remote removal or close.
|
||||
func (s *Subscriptions) Unsubscribe(ctx context.Context, p *Peer, key string) error {
|
||||
if !validSubscriptionKey(key) {
|
||||
return errors.New("rpc: invalid subscription key")
|
||||
}
|
||||
if err := s.lock(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
defer s.unlock()
|
||||
delete(s.desired, key)
|
||||
if p == nil {
|
||||
return nil
|
||||
}
|
||||
return p.Call(ctx, s.cfg.UnsubscribeMethod, struct {
|
||||
Key string `json:"key"`
|
||||
}{key}, nil)
|
||||
}
|
||||
|
||||
// Restore reauthorizes all desired subscriptions in deterministic order. It
|
||||
// stops at the first error and reports how many were acknowledged. A caller can
|
||||
// retry on the same peer only if its subscription handler is idempotent.
|
||||
func (s *Subscriptions) Restore(ctx context.Context, p *Peer) (int, error) {
|
||||
if p == nil {
|
||||
return 0, ErrNotConnected
|
||||
}
|
||||
if err := s.lock(ctx); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer s.unlock()
|
||||
keys := make([]string, 0, len(s.desired))
|
||||
for key := range s.desired {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
for i, key := range keys {
|
||||
if err := p.Call(ctx, s.cfg.SubscribeMethod, Subscription{key, s.desired[key]}, nil); err != nil {
|
||||
return i, err
|
||||
}
|
||||
}
|
||||
return len(keys), nil
|
||||
}
|
||||
|
||||
// Keys returns a sorted snapshot. In-flight restoration serializes desired
|
||||
// changes; waiting for that operation always respects ctx.
|
||||
func (s *Subscriptions) Keys(ctx context.Context) ([]string, error) {
|
||||
if err := s.lock(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer s.unlock()
|
||||
keys := make([]string, 0, len(s.desired))
|
||||
for key := range s.desired {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys, nil
|
||||
}
|
||||
405
ws/server.go
Normal file
405
ws/server.go
Normal file
@@ -0,0 +1,405 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
// ServerConfig is independent of HTTP body limits and HTTP request deadlines.
|
||||
// Context owns accepted sessions; it must outlive HTTP handlers. Prepare runs
|
||||
// before upgrading and may authenticate or create per-connection handlers.
|
||||
type ServerConfig struct {
|
||||
Session SessionConfig `json:"session"`
|
||||
MaxConnections int `json:"maxConnections"`
|
||||
HandshakeTimeout time.Duration `json:"handshakeTimeout"`
|
||||
AllowedOrigins []string `json:"allowedOrigins"`
|
||||
RequireOrigin bool `json:"requireOrigin"`
|
||||
Subprotocols []string `json:"subprotocols"`
|
||||
RequireSubprotocol bool `json:"requireSubprotocol"`
|
||||
EnableCompression bool `json:"enableCompression"`
|
||||
HandshakeRate float64 `json:"handshakeRate"` // Per second; zero disables rate limiting.
|
||||
HandshakeBurst int `json:"handshakeBurst"`
|
||||
Context context.Context `json:"-"`
|
||||
Prepare func(context.Context, *http.Request) (Handlers, error) `json:"-"`
|
||||
}
|
||||
|
||||
func DefaultServerConfig() ServerConfig {
|
||||
return ServerConfig{Session: DefaultSessionConfig(), MaxConnections: 4096, HandshakeTimeout: 10 * time.Second}
|
||||
}
|
||||
|
||||
func (cfg ServerConfig) Validate() error { _, err := normalizeServerConfig(cfg); return err }
|
||||
|
||||
func normalizeServerConfig(cfg ServerConfig) (ServerConfig, error) {
|
||||
var err error
|
||||
cfg.Session, err = normalizeSessionConfig(cfg.Session)
|
||||
if err != nil {
|
||||
return cfg, err
|
||||
}
|
||||
if cfg.MaxConnections < 0 || cfg.MaxConnections > 1<<20 || cfg.HandshakeTimeout < 0 || cfg.HandshakeRate < 0 || cfg.HandshakeRate > 1e9 || cfg.HandshakeRate != cfg.HandshakeRate || cfg.HandshakeBurst < 0 || cfg.HandshakeBurst > 1<<20 {
|
||||
return cfg, errors.New("ws: invalid server limits")
|
||||
}
|
||||
if cfg.MaxConnections == 0 {
|
||||
cfg.MaxConnections = 4096
|
||||
}
|
||||
if cfg.HandshakeTimeout == 0 {
|
||||
cfg.HandshakeTimeout = 10 * time.Second
|
||||
}
|
||||
if cfg.HandshakeRate > 0 && cfg.HandshakeBurst == 0 {
|
||||
cfg.HandshakeBurst = 1
|
||||
}
|
||||
if cfg.HandshakeRate == 0 && cfg.HandshakeBurst != 0 {
|
||||
return cfg, errors.New("ws: handshake burst requires a rate")
|
||||
}
|
||||
if len(cfg.AllowedOrigins) > 1024 {
|
||||
return cfg, errors.New("ws: too many allowed origins")
|
||||
}
|
||||
cfg.AllowedOrigins = append([]string(nil), cfg.AllowedOrigins...)
|
||||
for i, v := range cfg.AllowedOrigins {
|
||||
o, ok := canonicalOrigin(v)
|
||||
if !ok {
|
||||
return cfg, errors.New("ws: invalid allowed origin")
|
||||
}
|
||||
cfg.AllowedOrigins[i] = o
|
||||
}
|
||||
if len(cfg.Subprotocols) > 64 {
|
||||
return cfg, errors.New("ws: too many subprotocols")
|
||||
}
|
||||
cfg.Subprotocols = append([]string(nil), cfg.Subprotocols...)
|
||||
seen := make(map[string]bool)
|
||||
for _, v := range cfg.Subprotocols {
|
||||
if !protocolToken(v) || seen[v] {
|
||||
return cfg, errors.New("ws: invalid or duplicate subprotocol")
|
||||
}
|
||||
seen[v] = true
|
||||
}
|
||||
if cfg.RequireSubprotocol && len(cfg.Subprotocols) == 0 {
|
||||
return cfg, errors.New("ws: required subprotocol list is empty")
|
||||
}
|
||||
if cfg.Context == nil {
|
||||
cfg.Context = context.Background()
|
||||
}
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
func protocolToken(s string) bool {
|
||||
if len(s) == 0 || len(s) > 256 {
|
||||
return false
|
||||
}
|
||||
for _, r := range s {
|
||||
if r <= 32 || r >= 127 || strings.ContainsRune("()<>@,;:\\\"/[]?={}", r) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func canonicalOrigin(raw string) (string, bool) {
|
||||
if len(raw) > 4096 || strings.ContainsAny(raw, " \t\r\n,\\?#") {
|
||||
return "", false
|
||||
}
|
||||
u, err := url.Parse(raw)
|
||||
if err != nil || u == nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.Path != "" || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || u.Opaque != "" || strings.Contains(u.Host, "*") || strings.HasSuffix(u.Host, ":") {
|
||||
return "", false
|
||||
}
|
||||
host := strings.ToLower(u.Host)
|
||||
// net/url accepts some unbracketed colon hosts. An Origin must contain a
|
||||
// serialized DNS host or a bracketed, valid IPv6 address, never that ambiguity.
|
||||
name := u.Hostname()
|
||||
if strings.Contains(name, ":") {
|
||||
if !strings.HasPrefix(u.Host, "[") || net.ParseIP(name) == nil {
|
||||
return "", false
|
||||
}
|
||||
} else {
|
||||
if strings.ContainsAny(u.Host, "[]") {
|
||||
return "", false
|
||||
}
|
||||
for _, r := range name {
|
||||
if !(r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '.' || r == '-' || r == '_') {
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
}
|
||||
if port := u.Port(); port != "" {
|
||||
n, err := strconv.Atoi(port)
|
||||
if err != nil || n < 1 || n > 65535 {
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
if u.Scheme == "http" && u.Port() == "80" {
|
||||
host = strings.TrimSuffix(host, ":80")
|
||||
}
|
||||
if u.Scheme == "https" && u.Port() == "443" {
|
||||
host = strings.TrimSuffix(host, ":443")
|
||||
}
|
||||
return u.Scheme + "://" + host, true
|
||||
}
|
||||
|
||||
// RejectError lets Prepare select a public HTTP 4xx status. Error details are
|
||||
// deliberately never written into the handshake response.
|
||||
type RejectError struct{ StatusCode int }
|
||||
|
||||
func (e *RejectError) Error() string { return "ws: handshake rejected" }
|
||||
|
||||
type ServerStats struct {
|
||||
ActiveConnections, Handshakes int
|
||||
Accepted, Rejected uint64
|
||||
}
|
||||
|
||||
// Server owns all upgraded sessions, including their callbacks. Use Stop and
|
||||
// Drain with lifecycle.Runtime, or Shutdown before releasing shared resources.
|
||||
type Server struct {
|
||||
cfg ServerConfig
|
||||
mu sync.Mutex
|
||||
stopped bool
|
||||
active int // includes reserved handshakes; protected by the same admission lock
|
||||
handshakes int
|
||||
sessions map[string]*Session
|
||||
done chan struct{}
|
||||
doneClosed bool
|
||||
accepted, rejected uint64
|
||||
tokens float64
|
||||
refill time.Time
|
||||
stopParent func() bool
|
||||
}
|
||||
|
||||
func NewServer(cfg ServerConfig) (*Server, error) {
|
||||
cfg, err := normalizeServerConfig(cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
s := &Server{cfg: cfg, sessions: make(map[string]*Session), done: make(chan struct{}), tokens: float64(cfg.HandshakeBurst), refill: time.Now()}
|
||||
s.stopParent = context.AfterFunc(cfg.Context, func() { _ = s.Stop(context.Background()) })
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (s *Server) Handler(h Handlers) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = s.Accept(w, r, h) })
|
||||
}
|
||||
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { _, _ = s.Accept(w, r, Handlers{}) }
|
||||
|
||||
func (s *Server) originAllowed(r *http.Request) bool {
|
||||
values := r.Header.Values("Origin")
|
||||
if len(values) == 0 {
|
||||
return !s.cfg.RequireOrigin
|
||||
}
|
||||
if len(values) != 1 {
|
||||
return false
|
||||
}
|
||||
origin, ok := canonicalOrigin(values[0])
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
if len(s.cfg.AllowedOrigins) > 0 {
|
||||
for _, v := range s.cfg.AllowedOrigins {
|
||||
if v == origin {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
scheme := "http"
|
||||
if r.TLS != nil {
|
||||
scheme = "https"
|
||||
}
|
||||
self, ok := canonicalOrigin(scheme + "://" + r.Host)
|
||||
return ok && origin == self
|
||||
}
|
||||
|
||||
func (s *Server) admit() int {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.stopped || s.cfg.Context.Err() != nil {
|
||||
s.rejected++
|
||||
return http.StatusServiceUnavailable
|
||||
}
|
||||
if s.active >= s.cfg.MaxConnections {
|
||||
s.rejected++
|
||||
return http.StatusServiceUnavailable
|
||||
}
|
||||
if s.cfg.HandshakeRate > 0 {
|
||||
now := time.Now()
|
||||
s.tokens = min(float64(s.cfg.HandshakeBurst), s.tokens+now.Sub(s.refill).Seconds()*s.cfg.HandshakeRate)
|
||||
s.refill = now
|
||||
if s.tokens < 1 {
|
||||
s.rejected++
|
||||
return http.StatusTooManyRequests
|
||||
}
|
||||
s.tokens--
|
||||
}
|
||||
s.active++
|
||||
s.handshakes++
|
||||
return 0
|
||||
}
|
||||
|
||||
func (s *Server) finishLocked() {
|
||||
if s.stopped && s.active == 0 && !s.doneClosed {
|
||||
s.doneClosed = true
|
||||
close(s.done)
|
||||
}
|
||||
}
|
||||
|
||||
// Accept writes the handshake response (including failures). It returns after
|
||||
// transferring the connection to Server; callers must not reuse the writer or
|
||||
// retain request-scoped state in callbacks. Session uses ServerConfig.Context.
|
||||
func (s *Server) Accept(w http.ResponseWriter, r *http.Request, h Handlers) (session *Session, retErr error) {
|
||||
if status := s.admit(); status != 0 {
|
||||
http.Error(w, http.StatusText(status), status)
|
||||
return nil, &HandshakeError{StatusCode: status}
|
||||
}
|
||||
transferred := false
|
||||
defer func() {
|
||||
if !transferred {
|
||||
s.mu.Lock()
|
||||
s.active--
|
||||
s.handshakes--
|
||||
s.rejected++
|
||||
s.finishLocked()
|
||||
s.mu.Unlock()
|
||||
}
|
||||
}()
|
||||
fail := func(status int, cause error) (*Session, error) {
|
||||
http.Error(w, http.StatusText(status), status)
|
||||
return nil, &HandshakeError{StatusCode: status, Cause: cause}
|
||||
}
|
||||
if r.Method != http.MethodGet || !websocket.IsWebSocketUpgrade(r) {
|
||||
return fail(http.StatusBadRequest, nil)
|
||||
}
|
||||
if !s.originAllowed(r) {
|
||||
return fail(http.StatusForbidden, nil)
|
||||
}
|
||||
if _, ok := w.(http.Hijacker); !ok {
|
||||
return fail(http.StatusNotImplemented, ErrUnsupportedTransport)
|
||||
}
|
||||
if s.cfg.RequireSubprotocol {
|
||||
matched := false
|
||||
for _, offered := range websocket.Subprotocols(r) {
|
||||
for _, supported := range s.cfg.Subprotocols {
|
||||
if offered == supported {
|
||||
matched = true
|
||||
}
|
||||
}
|
||||
}
|
||||
if !matched {
|
||||
return fail(http.StatusBadRequest, nil)
|
||||
}
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(r.Context(), s.cfg.HandshakeTimeout)
|
||||
defer cancel()
|
||||
if s.cfg.Prepare != nil {
|
||||
var err error
|
||||
h, err = prepareSafely(s.cfg.Prepare, ctx, r)
|
||||
if err != nil {
|
||||
status := http.StatusUnauthorized
|
||||
var reject *RejectError
|
||||
if errors.As(err, &reject) && reject != nil && reject.StatusCode >= 400 && reject.StatusCode < 500 {
|
||||
status = reject.StatusCode
|
||||
}
|
||||
return fail(status, nil)
|
||||
}
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return fail(http.StatusRequestTimeout, ctx.Err())
|
||||
}
|
||||
deadline, _ := ctx.Deadline()
|
||||
remaining := time.Until(deadline)
|
||||
if remaining <= 0 {
|
||||
return fail(http.StatusRequestTimeout, context.DeadlineExceeded)
|
||||
}
|
||||
upgrader := websocket.Upgrader{HandshakeTimeout: remaining, Subprotocols: s.cfg.Subprotocols, EnableCompression: s.cfg.EnableCompression, CheckOrigin: func(*http.Request) bool { return true }, Error: func(w http.ResponseWriter, _ *http.Request, status int, _ error) {
|
||||
http.Error(w, http.StatusText(status), status)
|
||||
}}
|
||||
conn, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return nil, &HandshakeError{Cause: errors.New("upgrade failed")}
|
||||
}
|
||||
session = newSession(conn, s.cfg.Session, h, s.cfg.Context)
|
||||
s.mu.Lock()
|
||||
s.handshakes--
|
||||
s.sessions[session.ID()] = session
|
||||
s.accepted++
|
||||
stopped := s.stopped
|
||||
transferred = true
|
||||
s.mu.Unlock()
|
||||
session.start()
|
||||
if stopped {
|
||||
_ = session.Close(websocket.CloseGoingAway, "")
|
||||
}
|
||||
go func() {
|
||||
<-session.Done()
|
||||
s.mu.Lock()
|
||||
delete(s.sessions, session.ID())
|
||||
s.active--
|
||||
s.finishLocked()
|
||||
s.mu.Unlock()
|
||||
}()
|
||||
return session, nil
|
||||
}
|
||||
|
||||
func prepareSafely(fn func(context.Context, *http.Request) (Handlers, error), ctx context.Context, r *http.Request) (h Handlers, err error) {
|
||||
defer func() {
|
||||
if recover() != nil {
|
||||
err = errors.New("ws: preparation failed")
|
||||
}
|
||||
}()
|
||||
return fn(ctx, r)
|
||||
}
|
||||
|
||||
// Stop closes admission and begins bounded graceful closing of existing
|
||||
// sessions. It never waits for callbacks and is safe to call from a callback.
|
||||
func (s *Server) Stop(context.Context) error {
|
||||
s.mu.Lock()
|
||||
s.stopped = true
|
||||
all := make([]*Session, 0, len(s.sessions))
|
||||
for _, v := range s.sessions {
|
||||
all = append(all, v)
|
||||
}
|
||||
s.finishLocked()
|
||||
s.mu.Unlock()
|
||||
for _, v := range all {
|
||||
_ = v.Close(websocket.CloseGoingAway, "")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Drain returns nil only when all handshakes, I/O and callbacks have exited.
|
||||
// On timeout it aborts sockets, but does not claim blocked user callbacks ended.
|
||||
func (s *Server) Drain(ctx context.Context) error {
|
||||
_ = s.Stop(ctx)
|
||||
select {
|
||||
case <-s.done:
|
||||
s.stopParent()
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
select {
|
||||
case <-s.done:
|
||||
s.stopParent()
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
s.mu.Lock()
|
||||
all := make([]*Session, 0, len(s.sessions))
|
||||
for _, v := range s.sessions {
|
||||
all = append(all, v)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
for _, v := range all {
|
||||
v.Abort()
|
||||
}
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
func (s *Server) Shutdown(ctx context.Context) error { return s.Drain(ctx) }
|
||||
func (s *Server) Stats() ServerStats {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return ServerStats{ActiveConnections: len(s.sessions), Handshakes: s.handshakes, Accepted: s.accepted, Rejected: s.rejected}
|
||||
}
|
||||
326
ws/server_test.go
Normal file
326
ws/server_test.go
Normal file
@@ -0,0 +1,326 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
func managedServer(t testing.TB, cfg ServerConfig, h Handlers) (*Server, *httptest.Server, string) {
|
||||
t.Helper()
|
||||
s, err := NewServer(cfg)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
httpServer := httptest.NewServer(s.Handler(h))
|
||||
t.Cleanup(func() {
|
||||
httpServer.Close()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := s.Shutdown(ctx); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
})
|
||||
return s, httpServer, "ws" + strings.TrimPrefix(httpServer.URL, "http")
|
||||
}
|
||||
|
||||
func TestServerOriginAndProtocolPolicy(t *testing.T) {
|
||||
_, host, endpoint := managedServer(t, ServerConfig{Subprotocols: []string{"test.v1"}, RequireSubprotocol: true}, Handlers{})
|
||||
for _, tt := range []struct {
|
||||
name, origin string
|
||||
code int
|
||||
}{
|
||||
{"service", "", 101}, {"same", host.URL, 101}, {"foreign", "https://foreign.invalid", 403}, {"null", "null", 403}, {"path", host.URL + "/", 403}, {"comma", host.URL + ", https://evil.invalid", 403},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
d := websocket.Dialer{Subprotocols: []string{"test.v1"}}
|
||||
h := http.Header{}
|
||||
if tt.origin != "" {
|
||||
h.Set("Origin", tt.origin)
|
||||
}
|
||||
c, r, err := d.Dial(endpoint, h)
|
||||
if c != nil {
|
||||
defer c.Close()
|
||||
}
|
||||
if r == nil || r.StatusCode != tt.code {
|
||||
t.Fatalf("status=%v err=%v", r, err)
|
||||
}
|
||||
if tt.code == 101 && (err != nil || c.Subprotocol() != "test.v1") {
|
||||
t.Fatal("negotiation failed", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
_, r, err := websocket.DefaultDialer.Dial(endpoint, nil)
|
||||
if err == nil || r.StatusCode != 400 {
|
||||
t.Fatal("missing subprotocol accepted")
|
||||
}
|
||||
d := websocket.Dialer{Subprotocols: []string{"test.v1"}}
|
||||
_, r, err = d.Dial(endpoint, http.Header{"Origin": []string{host.URL, host.URL}})
|
||||
if err == nil || r.StatusCode != 403 {
|
||||
t.Fatal("duplicate origin accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerAllowedOriginsAndAuthentication(t *testing.T) {
|
||||
s, _, endpoint := managedServer(t, ServerConfig{AllowedOrigins: []string{"https://allowed.invalid"}, RequireOrigin: true, Prepare: func(ctx context.Context, r *http.Request) (Handlers, error) {
|
||||
if r.Header.Get("Authorization") != "Bearer local-test" {
|
||||
return Handlers{}, &RejectError{StatusCode: 403}
|
||||
}
|
||||
return Handlers{OnMessage: func(ctx context.Context, s *Session, m Message) error { return s.Send(ctx, m) }}, nil
|
||||
}}, Handlers{})
|
||||
for _, h := range []http.Header{nil, {"Origin": []string{"https://allowed.invalid"}}, {"Authorization": []string{"Bearer local-test"}}, {"Origin": []string{"https://wrong.invalid"}, "Authorization": []string{"Bearer local-test"}}} {
|
||||
c, r, e := websocket.DefaultDialer.Dial(endpoint, h)
|
||||
if c != nil {
|
||||
c.Close()
|
||||
}
|
||||
if e == nil || r.StatusCode != 403 {
|
||||
t.Fatal("unauthorized handshake accepted")
|
||||
}
|
||||
}
|
||||
c, _, err := websocket.DefaultDialer.Dial(endpoint, http.Header{"Origin": []string{"https://allowed.invalid"}, "Authorization": []string{"Bearer local-test"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer c.Close()
|
||||
if err = c.WriteMessage(websocket.TextMessage, []byte("echo")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c.SetReadDeadline(time.Now().Add(time.Second))
|
||||
_, body, err := c.ReadMessage()
|
||||
if err != nil || string(body) != "echo" {
|
||||
t.Fatal(string(body), err)
|
||||
}
|
||||
if s.Stats().Accepted != 1 {
|
||||
t.Fatal(s.Stats())
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerStopDuringHandshake(t *testing.T) {
|
||||
entered, release := make(chan struct{}), make(chan struct{})
|
||||
s, _, endpoint := managedServer(t, ServerConfig{Prepare: func(context.Context, *http.Request) (Handlers, error) {
|
||||
close(entered)
|
||||
<-release
|
||||
return Handlers{}, nil
|
||||
}}, Handlers{})
|
||||
dialDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(dialDone)
|
||||
c, _, _ := websocket.DefaultDialer.Dial(endpoint, nil)
|
||||
if c != nil {
|
||||
defer c.Close()
|
||||
c.ReadMessage()
|
||||
}
|
||||
}()
|
||||
<-entered
|
||||
_ = s.Stop(context.Background())
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
|
||||
defer cancel()
|
||||
if err := s.Drain(ctx); !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatal("handshake was not tracked", err)
|
||||
}
|
||||
close(release)
|
||||
select {
|
||||
case <-dialDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("accepted connection escaped stop")
|
||||
}
|
||||
ctx2, cancel2 := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel2()
|
||||
if err := s.Drain(ctx2); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerLimitsAndShutdownCallbacks(t *testing.T) {
|
||||
entered, release := make(chan struct{}), make(chan struct{})
|
||||
s, _, endpoint := managedServer(t, ServerConfig{MaxConnections: 1, Session: SessionConfig{CloseTimeout: 20 * time.Millisecond}}, Handlers{OnMessage: func(context.Context, *Session, Message) error { close(entered); <-release; return nil }})
|
||||
c, _, err := websocket.DefaultDialer.Dial(endpoint, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer c.Close()
|
||||
_, r, err := websocket.DefaultDialer.Dial(endpoint, nil)
|
||||
if err == nil || r.StatusCode != 503 {
|
||||
t.Fatal("connection cap failed")
|
||||
}
|
||||
if err = c.WriteMessage(websocket.TextMessage, []byte("hold")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
<-entered
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
||||
defer cancel()
|
||||
if err = s.Shutdown(ctx); !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatal("blocked callback incorrectly drained", err)
|
||||
}
|
||||
close(release)
|
||||
ctx2, cancel2 := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel2()
|
||||
if err = s.Drain(ctx2); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, r, err = websocket.DefaultDialer.Dial(endpoint, nil)
|
||||
if err == nil || r.StatusCode != 503 {
|
||||
t.Fatal("server admitted after shutdown")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerHandshakeRateAndPanic(t *testing.T) {
|
||||
_, _, endpoint := managedServer(t, ServerConfig{HandshakeRate: 0.001, HandshakeBurst: 1, Prepare: func(context.Context, *http.Request) (Handlers, error) { panic("SECRET") }}, Handlers{})
|
||||
_, r, e := websocket.DefaultDialer.Dial(endpoint, nil)
|
||||
if e == nil || r.StatusCode != 401 {
|
||||
t.Fatal("prepare panic not contained", e)
|
||||
}
|
||||
_, r, e = websocket.DefaultDialer.Dial(endpoint, nil)
|
||||
if e == nil || r.StatusCode != 429 {
|
||||
t.Fatal("rate cap not enforced", e)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerUnsupportedWriter(t *testing.T) {
|
||||
s, err := NewServer(ServerConfig{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
r := httptest.NewRequest("GET", "http://localhost/ws", nil)
|
||||
r.Header.Set("Upgrade", "websocket")
|
||||
r.Header.Set("Connection", "Upgrade")
|
||||
w := httptest.NewRecorder()
|
||||
_, err = s.Accept(w, r, Handlers{})
|
||||
if !errors.Is(err, ErrUnsupportedTransport) || w.Code != 501 {
|
||||
t.Fatal(w.Code, err)
|
||||
}
|
||||
if err = s.Shutdown(context.Background()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerConcurrentShutdown(t *testing.T) {
|
||||
s, _, endpoint := managedServer(t, ServerConfig{Session: SessionConfig{CloseTimeout: 10 * time.Millisecond}}, Handlers{})
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 40; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
c, _, _ := websocket.DefaultDialer.Dial(endpoint, nil)
|
||||
if c != nil {
|
||||
defer c.Close()
|
||||
c.SetReadDeadline(time.Now().Add(time.Second))
|
||||
c.ReadMessage()
|
||||
}
|
||||
}()
|
||||
}
|
||||
_ = s.Stop(context.Background())
|
||||
wg.Wait()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := s.Drain(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if stats := s.Stats(); stats.ActiveConnections != 0 || stats.Handshakes != 0 {
|
||||
t.Fatal(stats)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerConfigValidation(t *testing.T) {
|
||||
for _, cfg := range []ServerConfig{{AllowedOrigins: []string{"*"}}, {AllowedOrigins: []string{"https://user:secret@host"}}, {AllowedOrigins: []string{"https://host/path"}}, {HandshakeRate: -1}, {HandshakeBurst: 1}, {RequireSubprotocol: true}, {Subprotocols: []string{"bad\r\n"}}, {MaxConnections: -1}} {
|
||||
if err := cfg.Validate(); err == nil {
|
||||
t.Fatal("bad config accepted")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerParallelConnections(t *testing.T) {
|
||||
s, _, endpoint := managedServer(t, ServerConfig{MaxConnections: 256}, Handlers{OnMessage: func(ctx context.Context, s *Session, m Message) error { return s.Send(ctx, m) }})
|
||||
var workers sync.WaitGroup
|
||||
for i := 0; i < 128; i++ {
|
||||
workers.Add(1)
|
||||
go func(id int) {
|
||||
defer workers.Done()
|
||||
c, _, err := websocket.DefaultDialer.Dial(endpoint, nil)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
defer c.Close()
|
||||
c.SetReadDeadline(time.Now().Add(15 * time.Second))
|
||||
c.SetWriteDeadline(time.Now().Add(15 * time.Second))
|
||||
for n := 0; n < 50; n++ {
|
||||
value := fmt.Sprintf("%d:%d", id, n)
|
||||
if err := c.WriteMessage(websocket.TextMessage, []byte(value)); err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
_, got, err := c.ReadMessage()
|
||||
if err != nil || string(got) != value {
|
||||
t.Errorf("session isolation: %s %v", got, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
workers.Wait()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := s.Shutdown(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if s.Stats().Accepted != 128 {
|
||||
t.Fatal(s.Stats())
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerTypedNilRejection(t *testing.T) {
|
||||
_, _, endpoint := managedServer(t, ServerConfig{Prepare: func(context.Context, *http.Request) (Handlers, error) { var err *RejectError; return Handlers{}, err }}, Handlers{})
|
||||
_, r, err := websocket.DefaultDialer.Dial(endpoint, nil)
|
||||
if err == nil || r == nil || r.StatusCode != 401 {
|
||||
t.Fatal(r, err)
|
||||
}
|
||||
}
|
||||
|
||||
func FuzzOrigin(f *testing.F) {
|
||||
for _, v := range []string{"https://example.com", "http://localhost:8080", "null", "https://[::1]:443", "https://host:invalid", "https://evil@safe"} {
|
||||
f.Add(v)
|
||||
}
|
||||
f.Fuzz(func(t *testing.T, v string) {
|
||||
a, ok := canonicalOrigin(v)
|
||||
if ok {
|
||||
b, ok2 := canonicalOrigin(a)
|
||||
if !ok2 || a != b {
|
||||
t.Fatal("not canonical")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func BenchmarkServerRoundTrip(b *testing.B) {
|
||||
for _, size := range []int{64, 1024, 64 << 10} {
|
||||
b.Run(fmt.Sprint(size), func(b *testing.B) {
|
||||
_, _, endpoint := managedServer(b, ServerConfig{}, Handlers{OnMessage: func(ctx context.Context, s *Session, m Message) error { return s.Send(ctx, m) }})
|
||||
c, _, err := websocket.DefaultDialer.Dial(endpoint, nil)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
defer c.Close()
|
||||
payload := make([]byte, size)
|
||||
b.SetBytes(int64(size) * 2)
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if err = c.WriteMessage(websocket.BinaryMessage, payload); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
if _, _, err = c.ReadMessage(); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
568
ws/session.go
Normal file
568
ws/session.go
Normal file
@@ -0,0 +1,568 @@
|
||||
// Package ws provides bounded, managed WebSocket connections.
|
||||
package ws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strconv"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
type MessageType int
|
||||
|
||||
const (
|
||||
Text MessageType = websocket.TextMessage
|
||||
Binary MessageType = websocket.BinaryMessage
|
||||
)
|
||||
|
||||
type Message struct {
|
||||
Type MessageType
|
||||
Data []byte
|
||||
}
|
||||
type Handler func(context.Context, *Session, Message) error
|
||||
|
||||
// Handlers execute serially for each session. OnClose runs after all I/O and
|
||||
// message callbacks have stopped; it must not call Wait on its own session.
|
||||
type Handlers struct {
|
||||
OnOpen func(context.Context, *Session) error
|
||||
OnMessage Handler
|
||||
OnClose func(*Session, error)
|
||||
}
|
||||
|
||||
type SessionStats struct {
|
||||
MessagesSent, MessagesReceived, BytesSent, BytesReceived uint64
|
||||
QueuedMessages int
|
||||
QueuedBytes int64
|
||||
ReadQueuedMessages int
|
||||
ReadQueuedBytes int64
|
||||
}
|
||||
|
||||
type outgoing struct {
|
||||
message Message
|
||||
ctx context.Context
|
||||
result chan error
|
||||
}
|
||||
type controlFrame struct {
|
||||
kind int
|
||||
data []byte
|
||||
}
|
||||
|
||||
var nextSessionID atomic.Uint64
|
||||
|
||||
// Session owns exactly one reader and one writer. Sends are ordered by queue
|
||||
// admission. Message callbacks never run on the reader goroutine.
|
||||
type Session struct {
|
||||
conn *websocket.Conn
|
||||
cfg SessionConfig
|
||||
handlers Handlers
|
||||
id string
|
||||
parent context.Context
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
done, stopped, closingCh, peerClosed, closeSent chan struct{}
|
||||
startOnce, stopOnce, peerOnce sync.Once
|
||||
wg sync.WaitGroup
|
||||
mu sync.Mutex
|
||||
closing bool
|
||||
closeCode int
|
||||
closeReason string
|
||||
err error
|
||||
stats SessionStats
|
||||
out chan outgoing
|
||||
in chan Message
|
||||
controls chan controlFrame
|
||||
}
|
||||
|
||||
// newSession requires a validated config. Ownership transfers before start so
|
||||
// callers can register the session without racing completion.
|
||||
func newSession(conn *websocket.Conn, cfg SessionConfig, handlers Handlers, parent context.Context) *Session {
|
||||
if parent == nil {
|
||||
parent = context.Background()
|
||||
}
|
||||
ctx, cancel := context.WithCancel(parent)
|
||||
cfg = cfg.defaults()
|
||||
return &Session{conn: conn, cfg: cfg, handlers: handlers, id: "ws-" + strconv.FormatUint(nextSessionID.Add(1), 10),
|
||||
parent: parent, ctx: ctx, cancel: cancel, done: make(chan struct{}), stopped: make(chan struct{}),
|
||||
closingCh: make(chan struct{}), peerClosed: make(chan struct{}), closeSent: make(chan struct{}),
|
||||
out: make(chan outgoing, cfg.WriteQueueSize), in: make(chan Message, cfg.ReadQueueSize), controls: make(chan controlFrame, 8)}
|
||||
}
|
||||
|
||||
func (s *Session) ID() string { return s.id }
|
||||
func (s *Session) Context() context.Context { return s.ctx }
|
||||
func (s *Session) Done() <-chan struct{} { return s.done }
|
||||
func (s *Session) Err() error { s.mu.Lock(); defer s.mu.Unlock(); return s.err }
|
||||
func (s *Session) Subprotocol() string { return s.conn.Subprotocol() }
|
||||
func (s *Session) Stats() SessionStats { s.mu.Lock(); defer s.mu.Unlock(); return s.stats }
|
||||
|
||||
func (s *Session) start() {
|
||||
s.startOnce.Do(func() {
|
||||
s.wg.Add(4)
|
||||
go func() { defer s.wg.Done(); s.readLoop() }()
|
||||
go func() { defer s.wg.Done(); s.writeLoop() }()
|
||||
go func() { defer s.wg.Done(); s.dispatchLoop() }()
|
||||
go func() { defer s.wg.Done(); s.monitor() }()
|
||||
go func() {
|
||||
s.wg.Wait()
|
||||
for {
|
||||
select {
|
||||
case item := <-s.out:
|
||||
s.complete(item, ErrClosed)
|
||||
default:
|
||||
goto drained
|
||||
}
|
||||
}
|
||||
drained:
|
||||
// Release any discarded inbound payloads before publishing completion.
|
||||
for {
|
||||
select {
|
||||
case msg := <-s.in:
|
||||
s.releaseRead(msg)
|
||||
default:
|
||||
goto notified
|
||||
}
|
||||
}
|
||||
notified:
|
||||
if s.handlers.OnClose != nil {
|
||||
func() {
|
||||
defer func() {
|
||||
if p := recover(); p != nil {
|
||||
s.recordError(&OpError{Op: "on_close", Cause: fmt.Errorf("callback panic: %v", p)})
|
||||
}
|
||||
}()
|
||||
s.handlers.OnClose(s, s.Err())
|
||||
}()
|
||||
}
|
||||
close(s.done)
|
||||
}()
|
||||
})
|
||||
}
|
||||
|
||||
// Send waits for the actual write. A cancellation before writing skips the
|
||||
// message. Cancellation during writing terminates the connection; delivery may
|
||||
// already have occurred and the caller must not automatically replay it.
|
||||
func (s *Session) Send(ctx context.Context, msg Message) error {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
item, err := s.enqueue(ctx, msg, true)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
select {
|
||||
case err := <-item.result:
|
||||
return err
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// TrySend only acknowledges queue admission, not delivery. The payload is copied.
|
||||
func (s *Session) TrySend(msg Message) error {
|
||||
_, err := s.enqueue(context.Background(), msg, false)
|
||||
return err
|
||||
}
|
||||
func (s *Session) SendJSON(ctx context.Context, value any) error {
|
||||
data, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
return &OpError{Op: "encode", Cause: err}
|
||||
}
|
||||
return s.Send(ctx, Message{Type: Text, Data: data})
|
||||
}
|
||||
|
||||
func (s *Session) enqueue(ctx context.Context, msg Message, wait bool) (outgoing, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return outgoing{}, err
|
||||
}
|
||||
if msg.Type != Text && msg.Type != Binary {
|
||||
return outgoing{}, ErrInvalidMessage
|
||||
}
|
||||
if int64(len(msg.Data)) > s.cfg.MaxMessageBytes {
|
||||
return outgoing{}, ErrMessageTooLarge
|
||||
}
|
||||
if msg.Type == Text && !utf8.Valid(msg.Data) {
|
||||
return outgoing{}, ErrInvalidMessage
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.closing {
|
||||
return outgoing{}, ErrClosed
|
||||
}
|
||||
if s.stats.QueuedMessages >= s.cfg.WriteQueueSize || int64(len(msg.Data)) > s.cfg.WriteQueueBytes-s.stats.QueuedBytes {
|
||||
return outgoing{}, ErrQueueFull
|
||||
}
|
||||
item := outgoing{message: Message{Type: msg.Type, Data: append([]byte(nil), msg.Data...)}, ctx: ctx}
|
||||
if wait {
|
||||
item.result = make(chan error, 1)
|
||||
}
|
||||
select {
|
||||
case s.out <- item:
|
||||
s.stats.QueuedMessages++
|
||||
s.stats.QueuedBytes += int64(len(msg.Data))
|
||||
return item, nil
|
||||
default:
|
||||
return outgoing{}, ErrQueueFull
|
||||
}
|
||||
}
|
||||
|
||||
func validCloseCode(code int) bool {
|
||||
switch code {
|
||||
case 1000, 1001, 1002, 1003, 1007, 1008, 1009, 1010, 1011, 1012, 1013, 1014:
|
||||
return true
|
||||
}
|
||||
return code >= 3000 && code <= 4999
|
||||
}
|
||||
|
||||
// Close stops admission and begins a bounded graceful close, draining already
|
||||
// admitted sends. It is idempotent and never waits for a callback to finish.
|
||||
func (s *Session) Close(code int, reason string) error {
|
||||
if !validCloseCode(code) || len(reason) > 123 || !utf8.ValidString(reason) {
|
||||
return ErrInvalidMessage
|
||||
}
|
||||
s.beginClose(code, reason, nil)
|
||||
return nil
|
||||
}
|
||||
func (s *Session) beginClose(code int, reason string, cause error) {
|
||||
s.mu.Lock()
|
||||
if cause != nil && s.err == nil {
|
||||
s.err = cause
|
||||
}
|
||||
if !s.closing {
|
||||
s.closing = true
|
||||
s.closeCode = code
|
||||
s.closeReason = reason
|
||||
close(s.closingCh)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
s.cancel()
|
||||
}
|
||||
func (s *Session) recordError(err error) {
|
||||
s.mu.Lock()
|
||||
if s.err == nil {
|
||||
s.err = err
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
func (s *Session) terminate(err error) {
|
||||
s.beginClose(1001, "", err)
|
||||
s.stopOnce.Do(func() { close(s.stopped); _ = s.conn.Close() })
|
||||
}
|
||||
|
||||
// Abort immediately closes the socket; Wait still waits for callbacks to return.
|
||||
func (s *Session) Abort() { s.terminate(ErrClosed) }
|
||||
func (s *Session) Wait(ctx context.Context) error {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
select {
|
||||
case <-s.done:
|
||||
return s.Err()
|
||||
default:
|
||||
}
|
||||
select {
|
||||
case <-s.done:
|
||||
return s.Err()
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Session) monitor() {
|
||||
select {
|
||||
case <-s.stopped:
|
||||
return
|
||||
case <-s.parent.Done():
|
||||
s.beginClose(1001, "", s.parent.Err())
|
||||
case <-s.closingCh:
|
||||
}
|
||||
timer := time.NewTimer(s.cfg.CloseTimeout)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-s.stopped:
|
||||
return
|
||||
case <-timer.C:
|
||||
s.terminate(&OpError{Op: "close", Cause: context.DeadlineExceeded})
|
||||
return
|
||||
case <-s.closeSent:
|
||||
}
|
||||
select {
|
||||
case <-s.stopped:
|
||||
return
|
||||
case <-timer.C:
|
||||
s.terminate(&OpError{Op: "close", Cause: context.DeadlineExceeded})
|
||||
case <-s.peerClosed:
|
||||
s.terminate(nil)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Session) readLoop() {
|
||||
s.conn.SetReadLimit(s.cfg.MaxMessageBytes)
|
||||
if !s.cfg.DisableHeartbeat {
|
||||
_ = s.conn.SetReadDeadline(time.Now().Add(s.cfg.PongTimeout))
|
||||
s.conn.SetPongHandler(func(string) error { return s.conn.SetReadDeadline(time.Now().Add(s.cfg.PongTimeout)) })
|
||||
}
|
||||
s.conn.SetPingHandler(func(data string) error {
|
||||
select {
|
||||
case s.controls <- controlFrame{websocket.PongMessage, []byte(data)}:
|
||||
return nil
|
||||
default:
|
||||
return ErrQueueFull
|
||||
}
|
||||
})
|
||||
s.conn.SetCloseHandler(func(code int, _ string) error {
|
||||
var cause error
|
||||
s.mu.Lock()
|
||||
locallyClosing := s.closing
|
||||
s.mu.Unlock()
|
||||
// A remote going-away frame can signal a restarting service. Preserve
|
||||
// its code for Client.Run's retry policy; our own close acknowledgement
|
||||
// remains a normal local shutdown.
|
||||
if code != 1000 && code != 1005 && (code != 1001 || !locallyClosing) {
|
||||
cause = &CloseError{Code: code}
|
||||
}
|
||||
reply := code
|
||||
if !validCloseCode(reply) {
|
||||
reply = 1000
|
||||
}
|
||||
s.beginClose(reply, "", cause)
|
||||
s.peerOnce.Do(func() { close(s.peerClosed) })
|
||||
return &websocket.CloseError{Code: code}
|
||||
})
|
||||
for {
|
||||
kind, reader, err := s.conn.NextReader()
|
||||
if err != nil {
|
||||
select {
|
||||
case <-s.peerClosed:
|
||||
return
|
||||
case <-s.stopped:
|
||||
return
|
||||
default:
|
||||
}
|
||||
var ce *websocket.CloseError
|
||||
if errors.As(err, &ce) {
|
||||
s.terminate(&CloseError{Code: ce.Code, Cause: err})
|
||||
} else if errors.Is(err, websocket.ErrReadLimit) {
|
||||
s.terminate(ErrMessageTooLarge)
|
||||
} else {
|
||||
s.terminate(&OpError{Op: "read", Cause: err})
|
||||
}
|
||||
return
|
||||
}
|
||||
data, err := io.ReadAll(io.LimitReader(reader, s.cfg.MaxMessageBytes+1))
|
||||
if err != nil {
|
||||
if errors.Is(err, websocket.ErrReadLimit) {
|
||||
s.beginClose(1009, "", ErrMessageTooLarge)
|
||||
} else {
|
||||
s.terminate(&OpError{Op: "read", Cause: err})
|
||||
}
|
||||
return
|
||||
}
|
||||
if int64(len(data)) > s.cfg.MaxMessageBytes {
|
||||
s.beginClose(1009, "", ErrMessageTooLarge)
|
||||
return
|
||||
}
|
||||
if kind == websocket.TextMessage && !utf8.Valid(data) {
|
||||
s.beginClose(1007, "", ErrInvalidMessage)
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
if s.closing {
|
||||
s.mu.Unlock()
|
||||
continue
|
||||
}
|
||||
if s.stats.ReadQueuedMessages >= s.cfg.ReadQueueSize || int64(len(data)) > s.cfg.ReadQueueBytes-s.stats.ReadQueuedBytes {
|
||||
s.mu.Unlock()
|
||||
s.beginClose(1008, "", ErrQueueFull)
|
||||
return
|
||||
}
|
||||
msg := Message{Type: MessageType(kind), Data: data}
|
||||
select {
|
||||
case s.in <- msg:
|
||||
s.stats.ReadQueuedMessages++
|
||||
s.stats.ReadQueuedBytes += int64(len(data))
|
||||
s.stats.MessagesReceived++
|
||||
s.stats.BytesReceived += uint64(len(data))
|
||||
s.mu.Unlock()
|
||||
default:
|
||||
s.mu.Unlock()
|
||||
s.beginClose(1008, "", ErrQueueFull)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Session) releaseRead(msg Message) {
|
||||
s.mu.Lock()
|
||||
s.stats.ReadQueuedMessages--
|
||||
s.stats.ReadQueuedBytes -= int64(len(msg.Data))
|
||||
s.mu.Unlock()
|
||||
}
|
||||
func (s *Session) callback(fn func() error) (err error) {
|
||||
defer func() {
|
||||
if p := recover(); p != nil {
|
||||
err = fmt.Errorf("callback panic: %v", p)
|
||||
}
|
||||
}()
|
||||
return fn()
|
||||
}
|
||||
func (s *Session) dispatchLoop() {
|
||||
if s.ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
if s.handlers.OnOpen != nil {
|
||||
if err := s.callback(func() error { return s.handlers.OnOpen(s.ctx, s) }); err != nil {
|
||||
s.beginClose(1011, "", &OpError{Op: "on_open", Cause: err})
|
||||
return
|
||||
}
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
case msg := <-s.in:
|
||||
if s.ctx.Err() != nil {
|
||||
s.releaseRead(msg)
|
||||
return
|
||||
}
|
||||
var err error
|
||||
if s.handlers.OnMessage != nil {
|
||||
err = s.callback(func() error { return s.handlers.OnMessage(s.ctx, s, msg) })
|
||||
}
|
||||
s.releaseRead(msg)
|
||||
if err != nil {
|
||||
s.beginClose(1011, "", &OpError{Op: "on_message", Cause: err})
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Session) complete(item outgoing, err error) {
|
||||
s.mu.Lock()
|
||||
s.stats.QueuedMessages--
|
||||
s.stats.QueuedBytes -= int64(len(item.message.Data))
|
||||
if err == nil {
|
||||
s.stats.MessagesSent++
|
||||
s.stats.BytesSent += uint64(len(item.message.Data))
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if item.result != nil {
|
||||
item.result <- err
|
||||
}
|
||||
}
|
||||
func (s *Session) write(item outgoing) error {
|
||||
if err := item.ctx.Err(); err != nil {
|
||||
s.complete(item, err)
|
||||
return nil
|
||||
}
|
||||
deadline := time.Now().Add(s.cfg.WriteTimeout)
|
||||
if d, ok := item.ctx.Deadline(); ok && d.Before(deadline) {
|
||||
deadline = d
|
||||
}
|
||||
if err := s.conn.SetWriteDeadline(deadline); err != nil {
|
||||
s.complete(item, &OpError{Op: "write", Cause: err})
|
||||
return err
|
||||
}
|
||||
// Only the writer touches Gorilla's deadline field. Cancellation interrupts
|
||||
// the underlying socket directly; joining the callback prevents a late expiry
|
||||
// from corrupting the following message's deadline.
|
||||
canceled := make(chan struct{})
|
||||
stop := context.AfterFunc(item.ctx, func() { _ = s.conn.UnderlyingConn().SetWriteDeadline(time.Now()); close(canceled) })
|
||||
err := s.conn.WriteMessage(int(item.message.Type), item.message.Data)
|
||||
if !stop() {
|
||||
<-canceled
|
||||
}
|
||||
if err != nil {
|
||||
if item.ctx.Err() != nil {
|
||||
err = item.ctx.Err()
|
||||
} else if d, ok := item.ctx.Deadline(); ok && !time.Now().Before(d) {
|
||||
// The socket timer can win the race with context's timer.
|
||||
err = context.DeadlineExceeded
|
||||
}
|
||||
wrapped := &OpError{Op: "write", Cause: err}
|
||||
s.complete(item, wrapped)
|
||||
return wrapped
|
||||
}
|
||||
s.complete(item, nil)
|
||||
return nil
|
||||
}
|
||||
func (s *Session) writeControl(frame controlFrame) error {
|
||||
return s.conn.WriteControl(frame.kind, frame.data, time.Now().Add(s.cfg.WriteTimeout))
|
||||
}
|
||||
func (s *Session) writeLoop() {
|
||||
var ticks <-chan time.Time
|
||||
if !s.cfg.DisableHeartbeat {
|
||||
ticker := time.NewTicker(s.cfg.PingInterval)
|
||||
defer ticker.Stop()
|
||||
ticks = ticker.C
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case <-s.stopped:
|
||||
return
|
||||
case <-s.closingCh:
|
||||
s.finishWrites()
|
||||
return
|
||||
default:
|
||||
}
|
||||
select {
|
||||
case <-s.stopped:
|
||||
return
|
||||
case <-s.closingCh:
|
||||
s.finishWrites()
|
||||
return
|
||||
case item := <-s.out:
|
||||
if err := s.write(item); err != nil {
|
||||
s.terminate(err)
|
||||
return
|
||||
}
|
||||
case frame := <-s.controls:
|
||||
if err := s.writeControl(frame); err != nil {
|
||||
s.terminate(&OpError{Op: "control", Cause: err})
|
||||
return
|
||||
}
|
||||
case <-ticks:
|
||||
if err := s.writeControl(controlFrame{kind: websocket.PingMessage}); err != nil {
|
||||
s.terminate(&OpError{Op: "ping", Cause: err})
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
func (s *Session) finishWrites() {
|
||||
for {
|
||||
select {
|
||||
case <-s.stopped:
|
||||
return
|
||||
case item := <-s.out:
|
||||
select {
|
||||
case <-s.peerClosed:
|
||||
s.complete(item, ErrClosed)
|
||||
continue
|
||||
default:
|
||||
}
|
||||
if err := s.write(item); err != nil {
|
||||
s.terminate(err)
|
||||
return
|
||||
}
|
||||
default:
|
||||
goto closeFrame
|
||||
}
|
||||
}
|
||||
closeFrame:
|
||||
s.mu.Lock()
|
||||
code, reason := s.closeCode, s.closeReason
|
||||
s.mu.Unlock()
|
||||
if err := s.writeControl(controlFrame{websocket.CloseMessage, websocket.FormatCloseMessage(code, reason)}); err != nil {
|
||||
s.terminate(&OpError{Op: "close", Cause: err})
|
||||
return
|
||||
}
|
||||
close(s.closeSent)
|
||||
}
|
||||
392
ws/session_test.go
Normal file
392
ws/session_test.go
Normal file
@@ -0,0 +1,392 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
func sessionPair(t testing.TB, cfg SessionConfig, h Handlers, start, compression bool) (*Session, *websocket.Conn) {
|
||||
t.Helper()
|
||||
if err := cfg.Validate(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
accepted := make(chan *Session, 1)
|
||||
up := websocket.Upgrader{EnableCompression: compression}
|
||||
host := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := up.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
s := newSession(conn, cfg, h, context.Background())
|
||||
accepted <- s
|
||||
if start {
|
||||
s.start()
|
||||
}
|
||||
}))
|
||||
dialer := websocket.Dialer{EnableCompression: compression}
|
||||
peer, _, err := dialer.Dial("ws"+strings.TrimPrefix(host.URL, "http"), nil)
|
||||
if err != nil {
|
||||
host.Close()
|
||||
t.Fatal(err)
|
||||
}
|
||||
s := <-accepted
|
||||
t.Cleanup(func() {
|
||||
s.Abort()
|
||||
s.start()
|
||||
_ = peer.Close()
|
||||
host.Close()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
_ = s.Wait(ctx)
|
||||
})
|
||||
return s, peer
|
||||
}
|
||||
|
||||
func sessionTestConfig() SessionConfig {
|
||||
return SessionConfig{DisableHeartbeat: true, CloseTimeout: 100 * time.Millisecond, WriteTimeout: time.Second}
|
||||
}
|
||||
func waitSession(t *testing.T, s *Session) error {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
err := s.Wait(ctx)
|
||||
if errors.Is(err, context.DeadlineExceeded) && ctx.Err() != nil {
|
||||
t.Fatal("session did not finish")
|
||||
}
|
||||
return err
|
||||
}
|
||||
func peerClose(t *testing.T, peer *websocket.Conn) {
|
||||
t.Helper()
|
||||
_ = peer.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(1000, ""), time.Now().Add(time.Second))
|
||||
_ = peer.SetReadDeadline(time.Now().Add(time.Second))
|
||||
_, _, _ = peer.ReadMessage()
|
||||
}
|
||||
|
||||
func TestSessionEchoAndCompletion(t *testing.T) {
|
||||
var opened, closed atomic.Bool
|
||||
s, peer := sessionPair(t, sessionTestConfig(), Handlers{
|
||||
OnOpen: func(context.Context, *Session) error { opened.Store(true); return nil },
|
||||
OnMessage: func(ctx context.Context, s *Session, msg Message) error {
|
||||
if !opened.Load() {
|
||||
return errors.New("message before open")
|
||||
}
|
||||
return s.Send(ctx, msg)
|
||||
},
|
||||
OnClose: func(*Session, error) { closed.Store(true) },
|
||||
}, true, false)
|
||||
for _, kind := range []int{websocket.TextMessage, websocket.BinaryMessage} {
|
||||
if err := peer.WriteMessage(kind, []byte("hello")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, data, err := peer.ReadMessage()
|
||||
if err != nil || got != kind || string(data) != "hello" {
|
||||
t.Fatalf("echo: %d %q %v", got, data, err)
|
||||
}
|
||||
}
|
||||
peerClose(t, peer)
|
||||
if err := waitSession(t, s); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !closed.Load() {
|
||||
t.Fatal("Done preceded OnClose")
|
||||
}
|
||||
stats := s.Stats()
|
||||
if stats.MessagesReceived != 2 || stats.MessagesSent != 2 || stats.QueuedMessages != 0 || stats.ReadQueuedMessages != 0 {
|
||||
t.Fatalf("stats: %+v", stats)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionQueueBoundsCopiesAndGracefulDrain(t *testing.T) {
|
||||
cfg := sessionTestConfig()
|
||||
cfg.WriteQueueSize = 2
|
||||
cfg.WriteQueueBytes = 4
|
||||
s, peer := sessionPair(t, cfg, Handlers{}, false, false)
|
||||
data := []byte("ab")
|
||||
if err := s.TrySend(Message{Type: Text, Data: data}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
data[0] = 'z'
|
||||
if err := s.TrySend(Message{Type: Text, Data: []byte("cd")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.TrySend(Message{Type: Text, Data: []byte("e")}); !errors.Is(err, ErrQueueFull) {
|
||||
t.Fatalf("overflow: %v", err)
|
||||
}
|
||||
if err := s.Close(1000, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.TrySend(Message{Type: Text}); !errors.Is(err, ErrClosed) {
|
||||
t.Fatalf("closing send: %v", err)
|
||||
}
|
||||
s.start()
|
||||
for _, want := range []string{"ab", "cd"} {
|
||||
_, got, err := peer.ReadMessage()
|
||||
if err != nil || string(got) != want {
|
||||
t.Fatalf("drain: %q %v", got, err)
|
||||
}
|
||||
}
|
||||
_, _, err := peer.ReadMessage()
|
||||
var ce *websocket.CloseError
|
||||
if !errors.As(err, &ce) || ce.Code != 1000 {
|
||||
t.Fatalf("close: %v", err)
|
||||
}
|
||||
if err := waitSession(t, s); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionCanceledQueueDoesNotCloseConnection(t *testing.T) {
|
||||
s, peer := sessionPair(t, sessionTestConfig(), Handlers{}, false, false)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
result := make(chan error, 1)
|
||||
go func() { result <- s.Send(ctx, Message{Type: Text, Data: []byte("skip")}) }()
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for s.Stats().QueuedMessages == 0 {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("send never admitted")
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
cancel()
|
||||
if err := <-result; !errors.Is(err, context.Canceled) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.TrySend(Message{Type: Text, Data: []byte("next")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.start()
|
||||
_, data, err := peer.ReadMessage()
|
||||
if err != nil || string(data) != "next" {
|
||||
t.Fatalf("got %q %v", data, err)
|
||||
}
|
||||
peerClose(t, peer)
|
||||
if err := waitSession(t, s); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionConcurrentSendAndJSON(t *testing.T) {
|
||||
s, peer := sessionPair(t, sessionTestConfig(), Handlers{}, true, false)
|
||||
const count = 24
|
||||
results := make(chan error, count)
|
||||
for i := 0; i < count; i++ {
|
||||
go func(i int) { results <- s.SendJSON(context.Background(), i) }(i)
|
||||
}
|
||||
seen := make(map[string]bool)
|
||||
for i := 0; i < count; i++ {
|
||||
_, data, err := peer.ReadMessage()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if seen[string(data)] {
|
||||
t.Fatal("duplicate write")
|
||||
}
|
||||
seen[string(data)] = true
|
||||
}
|
||||
for i := 0; i < count; i++ {
|
||||
if err := <-results; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
peerClose(t, peer)
|
||||
if err := waitSession(t, s); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionBoundedDispatcherDoesNotBlockPing(t *testing.T) {
|
||||
entered := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
var once sync.Once
|
||||
s, peer := sessionPair(t, sessionTestConfig(), Handlers{OnMessage: func(ctx context.Context, _ *Session, _ Message) error {
|
||||
once.Do(func() { close(entered) })
|
||||
select {
|
||||
case <-release:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
return nil
|
||||
}}, true, false)
|
||||
t.Cleanup(func() { close(release) })
|
||||
if err := peer.WriteMessage(websocket.TextMessage, []byte("block callback")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
<-entered
|
||||
pong := make(chan struct{}, 1)
|
||||
peer.SetPongHandler(func(data string) error {
|
||||
if data == "probe" {
|
||||
pong <- struct{}{}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
readDone := make(chan struct{})
|
||||
go func() { defer close(readDone); _, _, _ = peer.ReadMessage() }()
|
||||
if err := peer.WriteControl(websocket.PingMessage, []byte("probe"), time.Now().Add(time.Second)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
select {
|
||||
case <-pong:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("blocked callback starved control frames")
|
||||
}
|
||||
_ = s.Close(1000, "")
|
||||
<-readDone
|
||||
if err := waitSession(t, s); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionInboundQueueFullCancelsCallback(t *testing.T) {
|
||||
cfg := sessionTestConfig()
|
||||
cfg.ReadQueueSize = 1
|
||||
cfg.ReadQueueBytes = 8
|
||||
entered := make(chan struct{})
|
||||
s, peer := sessionPair(t, cfg, Handlers{OnMessage: func(ctx context.Context, _ *Session, _ Message) error { close(entered); <-ctx.Done(); return nil }}, true, false)
|
||||
_ = peer.WriteMessage(websocket.TextMessage, []byte("one"))
|
||||
<-entered
|
||||
_ = peer.WriteMessage(websocket.TextMessage, []byte("two"))
|
||||
_, _, err := peer.ReadMessage()
|
||||
var ce *websocket.CloseError
|
||||
if !errors.As(err, &ce) || ce.Code != 1008 {
|
||||
t.Fatalf("close: %v", err)
|
||||
}
|
||||
if err := waitSession(t, s); !errors.Is(err, ErrQueueFull) {
|
||||
t.Fatalf("session: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionCompressedReadLimitAndInvalidUTF8(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
data []byte
|
||||
code int
|
||||
cause error
|
||||
}{
|
||||
{"compressed", []byte(strings.Repeat("x", 4096)), 1009, ErrMessageTooLarge},
|
||||
{"utf8", []byte{0xff}, 1007, ErrInvalidMessage},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
cfg := sessionTestConfig()
|
||||
cfg.MaxMessageBytes = 128
|
||||
s, peer := sessionPair(t, cfg, Handlers{}, true, true)
|
||||
if err := peer.WriteMessage(websocket.TextMessage, tc.data); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, _, err := peer.ReadMessage()
|
||||
var ce *websocket.CloseError
|
||||
if !errors.As(err, &ce) || ce.Code != tc.code {
|
||||
t.Fatalf("close: %v", err)
|
||||
}
|
||||
if err := waitSession(t, s); !errors.Is(err, tc.cause) {
|
||||
t.Fatalf("limit: %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionCallbackPanicAndCloseBarrier(t *testing.T) {
|
||||
closeEntered := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
s, peer := sessionPair(t, sessionTestConfig(), Handlers{
|
||||
OnMessage: func(context.Context, *Session, Message) error { panic("secret peer data") },
|
||||
OnClose: func(*Session, error) { close(closeEntered); <-release },
|
||||
}, true, false)
|
||||
_ = peer.WriteMessage(websocket.TextMessage, []byte("panic"))
|
||||
_, _, _ = peer.ReadMessage()
|
||||
<-closeEntered
|
||||
select {
|
||||
case <-s.Done():
|
||||
t.Fatal("Done ignored running OnClose")
|
||||
default:
|
||||
}
|
||||
close(release)
|
||||
if err := waitSession(t, s); err == nil || strings.Contains(err.Error(), "secret") {
|
||||
t.Fatalf("panic error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionHeartbeatDetectsUnresponsivePeer(t *testing.T) {
|
||||
cfg := SessionConfig{PingInterval: 10 * time.Millisecond, PongTimeout: 30 * time.Millisecond, CloseTimeout: 30 * time.Millisecond}
|
||||
s, _ := sessionPair(t, cfg, Handlers{}, true, false)
|
||||
if err := waitSession(t, s); err == nil {
|
||||
t.Fatal("peer without reader remained healthy")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionWriteCancellationStopsConnection(t *testing.T) {
|
||||
cfg := sessionTestConfig()
|
||||
cfg.MaxMessageBytes = 16 << 20
|
||||
cfg.WriteQueueBytes = 16 << 20
|
||||
s, _ := sessionPair(t, cfg, Handlers{}, true, false)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
err := s.Send(ctx, Message{Type: Binary, Data: make([]byte, 16<<20)})
|
||||
if !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("send: %v", err)
|
||||
}
|
||||
if err := waitSession(t, s); !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("write cancellation: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionValidationAndRedactedErrors(t *testing.T) {
|
||||
for _, cfg := range []SessionConfig{{MaxMessageBytes: -1}, {ReadQueueSize: -1}, {WriteQueueBytes: -1}, {WriteTimeout: -1}, {PingInterval: time.Second, PongTimeout: time.Second}} {
|
||||
if cfg.Validate() == nil {
|
||||
t.Fatalf("invalid config accepted: %+v", cfg)
|
||||
}
|
||||
}
|
||||
s, _ := sessionPair(t, sessionTestConfig(), Handlers{}, false, false)
|
||||
for _, msg := range []Message{{Type: 9}, {Type: Text, Data: []byte{0xff}}} {
|
||||
if !errors.Is(s.TrySend(msg), ErrInvalidMessage) {
|
||||
t.Fatal("invalid message accepted")
|
||||
}
|
||||
}
|
||||
if !errors.Is(s.TrySend(Message{Type: Binary, Data: make([]byte, (1<<20)+1)}), ErrMessageTooLarge) {
|
||||
t.Fatal("large outbound accepted")
|
||||
}
|
||||
for _, code := range []int{1005, 1006, 1015, 5000} {
|
||||
if s.Close(code, "") == nil {
|
||||
t.Fatalf("invalid close %d", code)
|
||||
}
|
||||
}
|
||||
secret := fmt.Errorf("https://user:password@example.test?token=secret")
|
||||
for _, err := range []error{&OpError{Op: "dial", Cause: secret}, &HandshakeError{StatusCode: 401, Cause: secret}, &CloseError{Code: 1008, Cause: secret}} {
|
||||
if strings.Contains(err.Error(), "secret") || !errors.Is(err, secret) {
|
||||
t.Fatalf("error wrapping: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkSessionSend(b *testing.B) {
|
||||
s, peer := sessionPair(b, sessionTestConfig(), Handlers{}, true, false)
|
||||
readerDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(readerDone)
|
||||
for {
|
||||
if _, _, err := peer.ReadMessage(); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
msg := Message{Type: Binary, Data: make([]byte, 256)}
|
||||
b.ReportAllocs()
|
||||
b.SetBytes(256)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if err := s.Send(context.Background(), msg); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
b.StopTimer()
|
||||
_ = s.Close(1000, "")
|
||||
<-readerDone
|
||||
}
|
||||
2
ws/testdata/fuzz/FuzzOrigin/1701c6390d7ef919
vendored
Normal file
2
ws/testdata/fuzz/FuzzOrigin/1701c6390d7ef919
vendored
Normal file
@@ -0,0 +1,2 @@
|
||||
go test fuzz v1
|
||||
string("https://::443")
|
||||
Reference in New Issue
Block a user