打开两个标签页，连接后发送文本或广播。
+未连接
+ + diff --git a/examples/websocket/main.go b/examples/websocket/main.go new file mode 100644 index 0000000..d2dd4e5 --- /dev/null +++ b/examples/websocket/main.go @@ -0,0 +1,107 @@ +package main + +import ( + "context" + _ "embed" + "errors" + "flag" + "fmt" + "net/http/httptest" + "os" + "strings" + "time" + + omni "git.yuntuo.hk/k2safe/go-omni" + "git.yuntuo.hk/k2safe/go-omni/httpx" + "git.yuntuo.hk/k2safe/go-omni/ws" + "git.yuntuo.hk/k2safe/go-omni/ws/hub" + wsomni "git.yuntuo.hk/k2safe/go-omni/ws/omni" +) + +//go:embed index.html +var page []byte + +func main() { + serve := flag.Bool("serve", false, "serve the browser demo on 127.0.0.1:8088") + flag.Parse() + if err := run(*serve); err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } +} + +func run(serve bool) (runErr error) { + application := omni.New(omni.WithLogger(false)) + defer func() { runErr = errors.Join(runErr, application.Close()) }() + connections, err := hub.New(hub.Config{}) + if err != nil { + return err + } + application.OnShutdown(connections.Close) + server, err := ws.NewServer(ws.ServerConfig{}) + if err != nil { + return err + } + if err = wsomni.Mount(application, "/ws", server, ws.Handlers{ + OnOpen: func(_ context.Context, s *ws.Session) error { + if err := connections.Add(s); err != nil { + return err + } + return connections.Join(s.ID(), "demo") + }, + OnMessage: func(ctx context.Context, s *ws.Session, m ws.Message) error { + if m.Type == ws.Text && string(m.Data) == "broadcast" { + result, err := connections.SendGroup("demo", ws.Message{Type: ws.Text, Data: []byte("hello group")}) + if err != nil { + return err + } + if len(result.Errors) > 0 { + return s.SendJSON(ctx, map[string]int{"queued": result.Queued, "failed": len(result.Errors)}) + } + return nil + } + return s.Send(ctx, m) + }, + }); err != nil { + return err + } + application.GET("/", func(c *httpx.Context) error { + c.Raw().Header().Set("Content-Type", "text/html; charset=utf-8") + _, err := c.Raw().Write(page) + return err + }) + if serve { + fmt.Println("Open http://127.0.0.1:8088 in two browser tabs; Ctrl+C stops cleanly.") + return application.Run("127.0.0.1:8088") + } + httpServer := httptest.NewServer(application) + defer httpServer.Close() + client, err := ws.NewClient(ws.ClientConfig{ProxyMode: "direct"}) + if err != nil { + return err + } + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + runErr = errors.Join(runErr, client.Shutdown(ctx)) + }() + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + replies := make(chan string, 2) + session, err := client.Dial(ctx, "ws"+strings.TrimPrefix(httpServer.URL, "http")+"/ws", ws.Handlers{OnMessage: func(_ context.Context, _ *ws.Session, m ws.Message) error { replies <- string(m.Data); return nil }}) + if err != nil { + return err + } + for _, message := range []string{"hello websocket", "broadcast"} { + if err = session.Send(ctx, ws.Message{Type: ws.Text, Data: []byte(message)}); err != nil { + return err + } + select { + case reply := <-replies: + fmt.Println(reply) + case <-ctx.Done(): + return ctx.Err() + } + } + return nil +} diff --git a/examples/websocketRPC/README.md b/examples/websocketRPC/README.md new file mode 100644 index 0000000..fedd4eb --- /dev/null +++ b/examples/websocketRPC/README.md @@ -0,0 +1,16 @@ +# 双向 RPC 与订阅恢复 + +从仓库根目录运行，无外部服务、数据库或凭证： + +```sh +go run ./examples/websocketRPC +go test -race ./examples/websocketRPC +``` + +示例在本地随机端口启动 HTTP/1.1 WebSocket 服务，客户端与服务端协商 `go-omni.rpc.v1`。每条连接创建独立 Peer：客户端调用 `double`，服务端处理 `identify` 时反向调用客户端的 `client.name`。随后客户端恢复 `updates` 订阅，服务端通过进程内 Hub 广播 RPC 事件；完成取消订阅和关闭后，用新连接再次执行恢复。 + +每代输出包含 `double=42`、反向调用返回值、`restored=1` 和对应事件序号。两次连接之间保留的是本地期望订阅，不是历史事件。示例显式创建第二条连接，方便验证每代 Peer 和订阅恢复的责任边界；自动网络重连使用 `ws.Client.Run` 与 `Prepare` 创建新 Peer，再从 `OnOpen` 外等待 `WaitReady` 并恢复订阅。 + +示例的订阅授权只允许公开的 `updates` 主题。生产项目在握手 `Prepare` 验证身份，每次订阅/恢复时检查具体主题权限。服务器注册订阅路由时通过连接 ID 加入 Hub；JSON 中的 key 不自动构成授权。没有离线补发、自动重试业务调用或持久化投递。 + +关闭 Client/Server 都使用独立、有界 context 并检查错误。RPC 的 OnClose 回调纳入实际业务任务完成屏障，示例的包装保留了该回调。完整协议、配置和退出语义见 [ws/rpc](../../ws/rpc/README.md)。 diff --git a/examples/websocketRPC/main.go b/examples/websocketRPC/main.go new file mode 100644 index 0000000..8155742 --- /dev/null +++ b/examples/websocketRPC/main.go @@ -0,0 +1,200 @@ +// websocketRPC demonstrates bidirectional calls and explicit subscription +// restoration over two successive real TCP connections, without external services. +package main + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "time" + + "git.yuntuo.hk/k2safe/go-omni/ws" + "git.yuntuo.hk/k2safe/go-omni/ws/hub" + "git.yuntuo.hk/k2safe/go-omni/ws/rpc" +) + +func main() { + if err := run(); err != nil { + panic(err) + } +} + +func run() (runErr error) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + groups, err := hub.New(hub.Config{}) + if err != nil { + return err + } + defer groups.Close() + server, err := ws.NewServer(ws.ServerConfig{ + Context: ctx, Subprotocols: []string{rpc.Subprotocol}, RequireSubprotocol: true, + Prepare: func(context.Context, *http.Request) (ws.Handlers, error) { + peer, err := rpc.New(rpc.Config{}) + if err != nil { + return ws.Handlers{}, err + } + var session *ws.Session + if err = peer.Handle("double", func(_ context.Context, raw json.RawMessage) (any, error) { + var n int + if json.Unmarshal(raw, &n) != nil { + return nil, &rpc.Error{Code: "INVALID_INPUT", Message: "integer required"} + } + return n * 2, nil + }); err != nil { + return ws.Handlers{}, err + } + if err = peer.Handle("identify", func(ctx context.Context, _ json.RawMessage) (any, error) { + var name string + err := peer.Call(ctx, "client.name", nil, &name) + return name, err + }); err != nil { + return ws.Handlers{}, err + } + if err = peer.Handle("rpc.subscribe", func(_ context.Context, raw json.RawMessage) (any, error) { + var sub rpc.Subscription + if json.Unmarshal(raw, &sub) != nil { + return nil, &rpc.Error{Code: "INVALID_INPUT", Message: "invalid subscription"} + } + // Demo policy. A real project authenticates in Prepare, then authorizes the + // requested topic against that verified identity on every subscription. + if sub.Key != "updates" { + return nil, &rpc.Error{Code: "DENIED", Message: "topic not allowed"} + } + return nil, groups.Join(session.ID(), sub.Key) + }); err != nil { + return ws.Handlers{}, err + } + if err = peer.Handle("rpc.unsubscribe", func(_ context.Context, raw json.RawMessage) (any, error) { + var sub struct { + Key string `json:"key"` + } + if json.Unmarshal(raw, &sub) != nil { + return nil, &rpc.Error{Code: "INVALID_INPUT", Message: "invalid subscription"} + } + groups.Leave(session.ID(), sub.Key) + return nil, nil + }); err != nil { + return ws.Handlers{}, err + } + handlers := peer.Handlers() + open := handlers.OnOpen + handlers.OnOpen = func(ctx context.Context, s *ws.Session) error { + session = s + if err := open(ctx, s); err != nil { + return err + } + return groups.Add(s) + } + closeHandler := handlers.OnClose + handlers.OnClose = func(s *ws.Session, err error) { closeHandler(s, err); groups.Remove(s.ID()) } + return handlers, nil + }, + }) + if err != nil { + return err + } + httpServer := httptest.NewServer(server.Handler(ws.Handlers{})) + defer httpServer.Close() + defer func() { + stop, done := context.WithTimeout(context.Background(), 2*time.Second) + defer done() + runErr = errors.Join(runErr, server.Shutdown(stop)) + }() + client, err := ws.NewClient(ws.ClientConfig{ProxyMode: "direct", Subprotocols: []string{rpc.Subprotocol}, RequireSubprotocol: true}) + if err != nil { + return err + } + defer func() { + stop, done := context.WithTimeout(context.Background(), 2*time.Second) + defer done() + runErr = errors.Join(runErr, client.Shutdown(stop)) + }() + subscriptions, err := rpc.NewSubscriptions(rpc.SubscriptionConfig{}) + if err != nil { + return err + } + if err = subscriptions.Subscribe(ctx, nil, "updates", nil); err != nil { + return err + } + endpoint := "ws" + strings.TrimPrefix(httpServer.URL, "http") + for generation := 1; generation <= 2; generation++ { + peer, err := rpc.New(rpc.Config{}) + if err != nil { + return err + } + events := make(chan int, 1) + if err = peer.Handle("client.name", func(context.Context, json.RawMessage) (any, error) { return "generic Go client", nil }); err != nil { + return err + } + if err = peer.Handle("updates.changed", func(ctx context.Context, raw json.RawMessage) (any, error) { + var v int + if err := json.Unmarshal(raw, &v); err != nil { + return nil, err + } + select { + case events <- v: + return nil, nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }); 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 doubled int + if err = peer.Call(ctx, "double", 21, &doubled); err != nil { + return err + } + var name string + if err = peer.Call(ctx, "identify", nil, &name); err != nil { + return err + } + restored, err := subscriptions.Restore(ctx, peer) + if err != nil { + return err + } + // The documented envelope also permits native browser clients and Hub fanout. + payload, err := json.Marshal(map[string]any{"v": 1, "type": "event", "method": "updates.changed", "params": generation}) + if err != nil { + return err + } + sent, err := groups.SendGroup("updates", ws.Message{Type: ws.Text, Data: payload}) + if err != nil { + return err + } + if sent.Queued != 1 || len(sent.Errors) > 0 { + return fmt.Errorf("broadcast failed: queued %d of %d", sent.Queued, sent.Attempted) + } + select { + case received := <-events: + fmt.Printf("connection %d: double=%d, remote callback=%q, restored=%d, event=%d\n", generation, doubled, name, restored, received) + case <-ctx.Done(): + return ctx.Err() + } + if err = subscriptions.Unsubscribe(ctx, peer, "updates"); err != nil { + return err + } + // Keep desired state separately for the next physical connection. + if err = subscriptions.Subscribe(ctx, nil, "updates", nil); err != nil { + return err + } + if err = session.Close(1000, ""); err != nil { + return err + } + if err = session.Wait(ctx); err != nil { + return err + } + } + return nil +} diff --git a/examples/websocketRPC/main_test.go b/examples/websocketRPC/main_test.go new file mode 100644 index 0000000..969c15a --- /dev/null +++ b/examples/websocketRPC/main_test.go @@ -0,0 +1,9 @@ +package main + +import "testing" + +func TestBidirectionalRPCAndRestore(t *testing.T) { + if err := run(); err != nil { + t.Fatal(err) + } +} diff --git a/go.mod b/go.mod index 0daf91d..c8143c5 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,7 @@ require ( github.com/go-sql-driver/mysql v1.9.3 github.com/golang-jwt/jwt/v5 v5.3.1 github.com/golang-migrate/migrate/v4 v4.18.1 + github.com/gorilla/websocket v1.5.3 github.com/minio/minio-go/v7 v7.0.83 github.com/redis/go-redis/v9 v9.17.3 go.opentelemetry.io/otel v1.31.0 diff --git a/go.sum b/go.sum index e758aca..679083f 100644 --- a/go.sum +++ b/go.sum @@ -62,6 +62,8 @@ github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= +github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4= github.com/hashicorp/errwrap v1.1.0 h1:OxrOeh75EUXMY8TBjag2fzXGZ40LB6IKw45YeGUDY2I= github.com/hashicorp/errwrap v1.1.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4= diff --git a/httpx/context.go b/httpx/context.go index 7fa6f8c..110330d 100644 --- a/httpx/context.go +++ b/httpx/context.go @@ -275,6 +275,14 @@ func (c *Context) Status() int { // Responded reports whether this handler has already written a response. func (c *Context) Responded() bool { return c.responded } +// MarkUpgraded records a successful WebSocket HTTP/1.1 upgrade whose 101 bytes +// were written directly to a hijacked connection. Call only after a successful +// upgrade; it does not write headers or upgrade a connection itself. +func (c *Context) MarkUpgraded() { + c.responded = true + c.status = http.StatusSwitchingProtocols +} + // WriteError writes a stable public error response. Internal causes are not // serialized. Responses committed through Context methods or Raw are preserved. func (c *Context) WriteError(err error) { diff --git a/lifecycle/README.md b/lifecycle/README.md index a64abd6..a792398 100644 --- a/lifecycle/README.md +++ b/lifecycle/README.md @@ -1,5 +1,11 @@ # lifecycle +## 独立 App 的后台组件 + +手写 `omni.New()` 应用也可以通过 `app.OnStop(name,hook)` / `app.OnDrain(name,hook)` 登记后台组件。HTTP Handler 结束后先运行 Stop、再 Drain，最后才执行逆序 OnShutdown。复用 `config.Lifecycle` 的阶段预算与 `ErrResourcesRetained` 规则，组件未退出时保留共享资源。没有登记后台钩子的旧 App 行为不变。 + +例如 `ws/omni.Mount` 自动登记 WS Server 的 Stop/Drain；生成项目则通常在 svc Runtime 上登记同一对钩子，并通过 `ws/omni.Handler` 注册路由。一个组件只选择一个所有者，不能把依赖释放登记到较早结束的另一个生命周期中。详见 [WS 接入](../ws/README.md#goomni-接入与停机)。 + ## 后台组件分阶段关闭 从 `v0.15.0-preview` 起，生成项目使用以下顺序：HTTP 停止接收并排空 → 全部 OnStop → 全部 OnDrain → 逆序 OnClose → 关闭日志。初始化、创建 App 或监听失败也经过后台收尾。Runtime 不导入 Orbit、数据库、ORM 或 Redis。 diff --git a/reports/websocket-bench-go1.25.txt b/reports/websocket-bench-go1.25.txt new file mode 100644 index 0000000..a06eed7 --- /dev/null +++ b/reports/websocket-bench-go1.25.txt @@ -0,0 +1,52 @@ +goos: darwin +goarch: arm64 +pkg: git.yuntuo.hk/k2safe/go-omni/ws +cpu: Apple M2 Max +BenchmarkClientRoundTrip-12 15706 47299 ns/op 0.19 MB/s 1520 B/op 13 allocs/op +BenchmarkClientRoundTrip-12 24013 26259 ns/op 0.34 MB/s 1520 B/op 13 allocs/op +BenchmarkClientRoundTrip-12 23534 25959 ns/op 0.35 MB/s 1520 B/op 13 allocs/op +BenchmarkServerRoundTrip/64-12 24288 24335 ns/op 5.26 MB/s 1568 B/op 13 allocs/op +BenchmarkServerRoundTrip/64-12 23325 23724 ns/op 5.40 MB/s 1568 B/op 13 allocs/op +BenchmarkServerRoundTrip/64-12 25240 23571 ns/op 5.43 MB/s 1568 B/op 13 allocs/op +BenchmarkServerRoundTrip/1024-12 23334 26435 ns/op 77.47 MB/s 7136 B/op 17 allocs/op +BenchmarkServerRoundTrip/1024-12 22939 28210 ns/op 72.60 MB/s 7136 B/op 17 allocs/op +BenchmarkServerRoundTrip/1024-12 23402 30573 ns/op 66.99 MB/s 7136 B/op 17 allocs/op +BenchmarkServerRoundTrip/65536-12 4402 140895 ns/op 930.28 MB/s 635956 B/op 45 allocs/op +BenchmarkServerRoundTrip/65536-12 4201 144080 ns/op 909.72 MB/s 635955 B/op 45 allocs/op +BenchmarkServerRoundTrip/65536-12 4046 145193 ns/op 902.74 MB/s 635953 B/op 45 allocs/op +BenchmarkSessionSend-12 110679 5370 ns/op 47.67 MB/s 1168 B/op 9 allocs/op +BenchmarkSessionSend-12 110362 5340 ns/op 47.94 MB/s 1168 B/op 9 allocs/op +BenchmarkSessionSend-12 113160 5339 ns/op 47.95 MB/s 1168 B/op 9 allocs/op +PASS +ok git.yuntuo.hk/k2safe/go-omni/ws 12.286s +goos: darwin +goarch: arm64 +pkg: git.yuntuo.hk/k2safe/go-omni/ws/hub +cpu: Apple M2 Max +BenchmarkHubBroadcast/1-12 3537676 152.5 ns/op 256 B/op 2 allocs/op +BenchmarkHubBroadcast/1-12 3975487 151.5 ns/op 256 B/op 2 allocs/op +BenchmarkHubBroadcast/1-12 3925134 154.0 ns/op 256 B/op 2 allocs/op +BenchmarkHubBroadcast/100-12 176622 3444 ns/op 3544 B/op 4 allocs/op +BenchmarkHubBroadcast/100-12 170431 3533 ns/op 3544 B/op 4 allocs/op +BenchmarkHubBroadcast/100-12 166556 3982 ns/op 3544 B/op 4 allocs/op +BenchmarkHubBroadcast/1000-12 8853 64870 ns/op 54656 B/op 6 allocs/op +BenchmarkHubBroadcast/1000-12 8582 65407 ns/op 54657 B/op 6 allocs/op +BenchmarkHubBroadcast/1000-12 8001 64712 ns/op 54657 B/op 6 allocs/op +PASS +ok git.yuntuo.hk/k2safe/go-omni/ws/hub 6.195s +PASS +ok git.yuntuo.hk/k2safe/go-omni/ws/observe 0.280s +PASS +ok git.yuntuo.hk/k2safe/go-omni/ws/omni 0.300s +goos: darwin +goarch: arm64 +pkg: git.yuntuo.hk/k2safe/go-omni/ws/rpc +cpu: Apple M2 Max +BenchmarkPeerCall-12 31053 17732 ns/op 8386 B/op 193 allocs/op +BenchmarkPeerCall-12 33787 17619 ns/op 8386 B/op 193 allocs/op +BenchmarkPeerCall-12 33901 17817 ns/op 8386 B/op 193 allocs/op +BenchmarkProtocolDecode-12 123168 4893 ns/op 19.62 MB/s 3656 B/op 106 allocs/op +BenchmarkProtocolDecode-12 122533 4824 ns/op 19.90 MB/s 3656 B/op 106 allocs/op +BenchmarkProtocolDecode-12 125842 4675 ns/op 20.54 MB/s 3656 B/op 106 allocs/op +PASS +ok git.yuntuo.hk/k2safe/go-omni/ws/rpc 4.531s diff --git a/reports/websocket-regression-go1.25.txt b/reports/websocket-regression-go1.25.txt new file mode 100644 index 0000000..48db07a --- /dev/null +++ b/reports/websocket-regression-go1.25.txt @@ -0,0 +1,86 @@ +ok git.yuntuo.hk/k2safe/go-omni 0.400s +ok git.yuntuo.hk/k2safe/go-omni/app 5.723s +ok git.yuntuo.hk/k2safe/go-omni/audit (cached) +ok git.yuntuo.hk/k2safe/go-omni/auth 1.045s +ok git.yuntuo.hk/k2safe/go-omni/cache (cached) +ok git.yuntuo.hk/k2safe/go-omni/cache/redis (cached) +ok git.yuntuo.hk/k2safe/go-omni/cmd/go-omni 3.720s +ok git.yuntuo.hk/k2safe/go-omni/cmd/go-omni/apigen 31.148s +ok git.yuntuo.hk/k2safe/go-omni/cmd/go-omni/internal/genfs (cached) +ok git.yuntuo.hk/k2safe/go-omni/cmd/go-omni/projectgen (cached) +ok git.yuntuo.hk/k2safe/go-omni/cmd/go-omni/scaffold 2.786s +ok git.yuntuo.hk/k2safe/go-omni/config (cached) +ok git.yuntuo.hk/k2safe/go-omni/contract (cached) +ok git.yuntuo.hk/k2safe/go-omni/copyx (cached) +ok git.yuntuo.hk/k2safe/go-omni/cors (cached) +ok git.yuntuo.hk/k2safe/go-omni/database 2.063s +ok git.yuntuo.hk/k2safe/go-omni/database/gormx (cached) +ok git.yuntuo.hk/k2safe/go-omni/database/migrate (cached) +ok git.yuntuo.hk/k2safe/go-omni/database/mysql (cached) +ok git.yuntuo.hk/k2safe/go-omni/errors (cached) +? git.yuntuo.hk/k2safe/go-omni/examples/basic [no test files] +? git.yuntuo.hk/k2safe/go-omni/examples/compatibility [no test files] +? git.yuntuo.hk/k2safe/go-omni/examples/compatibility/ids [no test files] +ok git.yuntuo.hk/k2safe/go-omni/examples/configuration (cached) +ok git.yuntuo.hk/k2safe/go-omni/examples/copying (cached) +? git.yuntuo.hk/k2safe/go-omni/examples/cors [no test files] +? git.yuntuo.hk/k2safe/go-omni/examples/developerTools [no test files] +ok git.yuntuo.hk/k2safe/go-omni/examples/httpClient (cached) +? git.yuntuo.hk/k2safe/go-omni/examples/httpSafety [no test files] +ok git.yuntuo.hk/k2safe/go-omni/examples/i18n 1.968s +ok git.yuntuo.hk/k2safe/go-omni/examples/modules 2.777s +ok git.yuntuo.hk/k2safe/go-omni/examples/platform 2.707s +ok git.yuntuo.hk/k2safe/go-omni/examples/reference 2.993s +? git.yuntuo.hk/k2safe/go-omni/examples/reference/handler [no test files] +? git.yuntuo.hk/k2safe/go-omni/examples/reference/model [no test files] +ok git.yuntuo.hk/k2safe/go-omni/examples/reference/repository (cached) +? git.yuntuo.hk/k2safe/go-omni/examples/reference/service [no test files] +ok git.yuntuo.hk/k2safe/go-omni/examples/signedRead (cached) +ok git.yuntuo.hk/k2safe/go-omni/examples/telemetryHeaders 2.822s +? git.yuntuo.hk/k2safe/go-omni/examples/transactions [no test files] +? git.yuntuo.hk/k2safe/go-omni/examples/websocket [no test files] +ok git.yuntuo.hk/k2safe/go-omni/examples/websocketRPC 2.839s +ok git.yuntuo.hk/k2safe/go-omni/health 3.382s +ok git.yuntuo.hk/k2safe/go-omni/httpc (cached) +ok git.yuntuo.hk/k2safe/go-omni/httpx 3.869s +ok git.yuntuo.hk/k2safe/go-omni/i18n (cached) +? git.yuntuo.hk/k2safe/go-omni/i18n/browser [no test files] +ok git.yuntuo.hk/k2safe/go-omni/i18n/errori18n (cached) +ok git.yuntuo.hk/k2safe/go-omni/i18n/httpi18n (cached) +ok git.yuntuo.hk/k2safe/go-omni/i18n/i18ntool (cached) +? git.yuntuo.hk/k2safe/go-omni/internal/configdiag [no test files] +? git.yuntuo.hk/k2safe/go-omni/internal/httplimits [no test files] +ok git.yuntuo.hk/k2safe/go-omni/internal/perfcheck (cached) +? git.yuntuo.hk/k2safe/go-omni/internal/storagefixture [no test files] +? git.yuntuo.hk/k2safe/go-omni/internal/testutil [no test files] +ok git.yuntuo.hk/k2safe/go-omni/lifecycle (cached) +ok git.yuntuo.hk/k2safe/go-omni/logx (cached) +ok git.yuntuo.hk/k2safe/go-omni/metrics (cached) +ok git.yuntuo.hk/k2safe/go-omni/middleware (cached) +ok git.yuntuo.hk/k2safe/go-omni/middleware/redis 2.257s +ok git.yuntuo.hk/k2safe/go-omni/openapi 2.425s +ok git.yuntuo.hk/k2safe/go-omni/pagination (cached) +? git.yuntuo.hk/k2safe/go-omni/queue [no test files] +ok git.yuntuo.hk/k2safe/go-omni/queue/redis (cached) +ok git.yuntuo.hk/k2safe/go-omni/resilience (cached) +ok git.yuntuo.hk/k2safe/go-omni/response (cached) +? git.yuntuo.hk/k2safe/go-omni/scripts/corsBrowser [no test files] +ok git.yuntuo.hk/k2safe/go-omni/storage (cached) +ok git.yuntuo.hk/k2safe/go-omni/storage/local (cached) +ok git.yuntuo.hk/k2safe/go-omni/storage/oss (cached) +ok git.yuntuo.hk/k2safe/go-omni/storage/s3 (cached) +? git.yuntuo.hk/k2safe/go-omni/storage/storagetest [no test files] +ok git.yuntuo.hk/k2safe/go-omni/task (cached) +ok git.yuntuo.hk/k2safe/go-omni/telemetry 2.595s +ok git.yuntuo.hk/k2safe/go-omni/transport 2.868s +ok git.yuntuo.hk/k2safe/go-omni/transport/http2 (cached) +? git.yuntuo.hk/k2safe/go-omni/transport/internal/routing [no test files] +ok git.yuntuo.hk/k2safe/go-omni/transport/native 5.567s +ok git.yuntuo.hk/k2safe/go-omni/transport/nethttp 3.181s +? git.yuntuo.hk/k2safe/go-omni/transport/transporttest [no test files] +? git.yuntuo.hk/k2safe/go-omni/validate [no test files] +ok git.yuntuo.hk/k2safe/go-omni/ws 3.875s +ok git.yuntuo.hk/k2safe/go-omni/ws/hub 3.462s +ok git.yuntuo.hk/k2safe/go-omni/ws/observe 3.337s +ok git.yuntuo.hk/k2safe/go-omni/ws/omni 3.019s +ok git.yuntuo.hk/k2safe/go-omni/ws/rpc 3.064s diff --git a/reports/websocket-validation-2026-10-29.md b/reports/websocket-validation-2026-10-29.md new file mode 100644 index 0000000..2c1efa7 --- /dev/null +++ b/reports/websocket-validation-2026-10-29.md @@ -0,0 +1,85 @@ +# WebSocket 验收记录（2026-10-29） + +本次增加通用 WS Client/Server/Session、Hub、可选 RPC/订阅恢复、观测、App 生命周期与生成项目扩展。对应版本 `v0.31.0-preview`，本文记录源码验证；远程发布验收单独记录。 + +## 结论 + +- 功能、真实 TCP/TLS、并发竞态、静态检查、Go 1.22 兼容性和可运行 Go 示例通过。 +- 最终全库回归输出见 [websocket-regression-go1.25.txt](websocket-regression-go1.25.txt)；Go 1.27.3 下 `go test ./ws/...` 也通过。生成示例只同步 appGen 与清单，生成漂移检查及消费项目编译通过；新增生成消费回归验证了自定义路由和失败清理。 +- Origin 模糊测试发现并修复无括号 IPv6 归一化歧义，失败样本保留在 `ws/testdata/fuzz/FuzzOrigin`；修复后 10 秒约 552,325 次执行通过。 +- 新增核心包在隔离 Go 1.27.3 工具链、govulncheck v1.9.0 下扫描未发现已知漏洞。 +- **包含 ws/omni 的完整接入安全扫描未通过**：仓库原有 `golang.org/x/net v0.33.0` 报告 4 项 HTTP/2 可达漏洞。不能将本次功能验收解释为整个框架已通过安全发布门槛。 +- 浏览器自动化访问 `http://127.0.0.1:8088` 被浏览器权限策略拒绝，因此浏览器 UI/浏览器协议实测未完成；不以 Go 客户端测试代替这一项。示例源文件可供手动运行，临时服务已停止。 + +## 功能与并发验证 + +使用 Apple M2 Max / darwin arm64，主测试工具链 Go 1.25.0。独立 Go 1.22.12 用于最低版本兼容检查，没有替换系统 Go。 + +执行过： + +```sh +go test ./... +go test -race -count=5 -timeout=120s ./ws/... ./examples/websocketRPC ./app ./httpx +go test -race -coverprofile=/tmp/go-omni-websocket-coverage.out ./ws/... +go vet ./ws/... ./app ./httpx ./examples/websocket ./examples/websocketRPC +GOTOOLCHAIN=go1.22.12 go test ./ws/... ./app ./httpx +go test ./ws -run '^$' -fuzz FuzzOrigin -fuzztime=10s -parallel=4 +go test ./ws/rpc -run '^$' -fuzz FuzzProtocolDecode -fuzztime=10s +go run ./examples/websocket +go run ./examples/websocketRPC +``` + +真实网络覆盖：回显、文本/二进制、128 个并行连接各 50 次往返、TLS/mTLS、HTTP CONNECT 代理认证、握手取消/超时、重连时刷新认证头、重试预算、快速重连与连接上限、双向嵌套 RPC、断线后两代 Session 的订阅恢复。 + +最终审查另修复了远端 `1001 going away` 被当成本地主动关闭、导致 Run 不重连的问题，并验证本地主动关闭仍停止重试；补充累计拨号/失败/重连指标和配置上界。修复后 `go test -race -count=3 ./ws/...` 通过。 + +安全与资源覆盖：同源/白名单/无 Origin/重复 Origin/null、拒绝错误子协议、认证拒绝与 panic、连接/握手速率限额、非法 UTF-8、压缩膨胀限额、发送数据拷贝、接收队列满、慢消费者、阻塞业务时仍处理 Ping/Pong、取消写入、断网、反复关闭、握手与停止并发、回调未退出时保留共享资源。 + +RPC 覆盖：双向请求、取消、迟到回复、重复在途 ID、未知字段、重复 JSON 键、嵌套深度、消息长度、pending/执行并发上限、公共与内部错误区分、panic 脱敏、订阅期望状态与恢复。RPC fuzz 10 秒约 71,268 次执行通过。 + +覆盖率快照： + +| 包 | 语句覆盖率 | +| --- | ---: | +| ws | 84.7% | +| ws/hub | 86.5% | +| ws/observe | 88.4% | +| ws/omni | 86.7% | +| ws/rpc | 80.7% | + +覆盖率只描述这组测试执行到的代码，不等同于正确性或安全证明。测试不包含生产长时间浸泡、跨地区网络容量、HTTP/2 extended CONNECT 或 Autobahn 全套认证。 + +## 性能基准 + +```sh +go test -run '^$' -bench . -benchmem -benchtime=500ms -count=3 ./ws/... +``` + +原始输出：[websocket-bench-go1.25.txt](websocket-bench-go1.25.txt)。环境为 Apple M2 Max、Go 1.25.0、GOMAXPROCS=12、loopback、无 TLS、压缩关闭。三次样本的范围如下；首轮客户端样本与其它验证有时间重叠，保留原始波动，不选取单个最好值作为承诺。 + +| 基准 | 时间/操作 | 分配 | +| --- | ---: | ---: | +| Client 回声往返（9 B） | 26.0–47.3 μs | 1520 B / 13 allocs | +| Server 回声往返（64 B） | 23.6–24.3 μs | 1568 B / 13 allocs | +| Server 回声往返（1 KiB） | 26.4–30.6 μs | 7136 B / 17 allocs | +| Server 回声往返（64 KiB） | 140.9–145.2 μs | 约 636 KB / 45 allocs | +| Session 发送（256 B） | 5.34–5.37 μs | 1168 B / 9 allocs | +| Hub 广播到 1000 个模拟连接 | 64.7–65.4 μs | 约 54.7 KB / 6 allocs | +| 内存传输 RPC Call | 17.6–17.8 μs | 8386 B / 193 allocs | +| 严格 RPC envelope 解析 | 4.68–4.89 μs | 3656 B / 106 allocs | + +网络往返基准包括双方处理及本机 socket 成本；64 KiB 场景的分配包含双方完整读入及安全发送拷贝，后续大消息优化需在保持消息所有权的前提下进行。Hub 基准仅测注册表和分发，RPC 内存基准不包含网络。这里没有从单连接 loopback 数字外推生产 QPS 或最大连接数。 + +内存预算需按连接数乘以消息/队列上限估算，并额外计入解码中消息、协议缓冲和业务持有的 payload。默认 4096 连接上限并不意味着任何机器都能承受所有连接同时用满双向 4 MiB 队列。部署时需根据机器容量设置 MaxConnections、消息与队列字节数、RPC 执行并发；压缩及 TLS 应另做容量测量。 + +## 漏洞扫描与剩余风险 + +扫描包含源码可达性分析，不是人工证明每条告警均可利用。没有隐藏、忽略或豁免扫描结果。 + +1. Go 1.25.0 + govulncheck v1.1.4 扫描 `./ws/...`：29 项标准库 symbol findings，以及 4 项 x/net symbol findings；Gorilla WebSocket 没有报告项。工具链版本是现有本机版本。扫描结果摘要保留在 [JSON](websocket-vuln-go1.25.json)。 +2. 隔离 Go 1.27.3 + govulncheck v1.9.0 扫描 `./ws/...`：标准库告警消除，剩余 4 项 x/net HTTP/2 findings，详见 [完整输出](websocket-vuln-patched-go.txt)。 +3. 同一修补工具链扫描 `./ws ./ws/hub ./ws/rpc ./ws/observe`：`No vulnerabilities found.`，见 [输出](websocket-vuln-core-patched-go.txt)。这不覆盖 ws/omni 引入的原有 App/HTTP2 依赖图。 + +剩余报告：GO-2026-6617（HPACK 并发）、GO-2026-6612（流控退款）、GO-2026-6611（重复窗口变更 CPU）、GO-2026-6603（Trailer 内存）。扫描器给出的修复版本为 `golang.org/x/net v0.60.0`；它的 [官方 go.mod](https://raw.githubusercontent.com/golang/net/v0.60.0/go.mod) 要求 Go 1.26，并升级 x/crypto、x/sys、x/text。此变更会改变框架当前公开的 Go 1.22 基线与多个既有适配器，未在 WS 功能提交中隐式实施。需要作为明确的框架安全升级进行全库与消费项目回归，生产发布前解决。 + +生产应使用受支持且已修补的 Go 工具链。Go 1.22 测试只验证源码兼容，不代表该旧工具链具备当前安全修复。对应标准库证书问题的官方说明见 [GO-2026-6809](https://pkg.go.dev/vuln/GO-2026-6809)；HTTP/2 问题见 [GO-2026-6617](https://pkg.go.dev/vuln/GO-2026-6617)。 diff --git a/reports/websocket-vuln-core-patched-go.txt b/reports/websocket-vuln-core-patched-go.txt new file mode 100644 index 0000000..e758922 --- /dev/null +++ b/reports/websocket-vuln-core-patched-go.txt @@ -0,0 +1 @@ +No vulnerabilities found. diff --git a/reports/websocket-vuln-go1.25.json b/reports/websocket-vuln-go1.25.json new file mode 100644 index 0000000..3ed3149 --- /dev/null +++ b/reports/websocket-vuln-go1.25.json @@ -0,0 +1,839 @@ +{ + "command": "go run golang.org/x/vuln/cmd/govulncheck@v1.1.4 -format=json ./ws/...", + "config": { + "protocol_version": "v1.0.0", + "scanner_name": "govulncheck", + "scanner_version": "v1.1.4", + "db": "https://vuln.go.dev", + "db_last_modified": "2026-10-27T21:16:18Z", + "go_version": "go1.25.0", + "scan_level": "symbol", + "scan_mode": "source" + }, + "findings": [ + { + "id": "GO-2025-3487", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.35.0", + "symbolReachable": false, + "summary": "Potential denial of service in golang.org/x/crypto" + }, + { + "id": "GO-2025-4116", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.43.0", + "symbolReachable": false, + "summary": "Potential denial of service in golang.org/x/crypto/ssh/agent" + }, + { + "id": "GO-2025-4134", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.45.0", + "symbolReachable": false, + "summary": "Unbounded memory consumption in golang.org/x/crypto/ssh" + }, + { + "id": "GO-2025-4135", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.45.0", + "symbolReachable": false, + "summary": "Malformed constraint may cause denial of service in golang.org/x/crypto/ssh/agent" + }, + { + "id": "GO-2026-5005", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.52.0", + "symbolReachable": false, + "summary": "Invoking key constraints not enforced in golang.org/x/crypto/ssh/agent" + }, + { + "id": "GO-2026-5006", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.52.0", + "symbolReachable": false, + "summary": "Invoking agent constraints dropped when forwarding keys in golang.org/x/crypto/ssh/agent" + }, + { + "id": "GO-2026-5013", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.52.0", + "symbolReachable": false, + "summary": "Invoking byte arithmetic causes underflow and panic in golang.org/x/crypto/ssh" + }, + { + "id": "GO-2026-5014", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.52.0", + "symbolReachable": false, + "summary": "Invoking bypass of certificate restrictions in golang.org/x/crypto/ssh" + }, + { + "id": "GO-2026-5015", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.52.0", + "symbolReachable": false, + "summary": "Invoking server panic during CheckHostKey/Authenticate in golang.org/x/crypto/ssh" + }, + { + "id": "GO-2026-5016", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.52.0", + "symbolReachable": false, + "summary": "Invoking memory leak when rejecting channels can lead to DoS in golang.org/x/crypto/ssh" + }, + { + "id": "GO-2026-5017", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.52.0", + "symbolReachable": false, + "summary": "Invoking client can cause server deadlock on unexpected responses in golang.org/x/crypto/ssh" + }, + { + "id": "GO-2026-5018", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.52.0", + "symbolReachable": false, + "summary": "Invoking pathological RSA/DSA parameters may cause DoS in golang.org/x/crypto/ssh" + }, + { + "id": "GO-2026-5019", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.52.0", + "symbolReachable": false, + "summary": "Invoking bypass of FIDO/U2F security keys physical interaction in golang.org/x/crypto/ssh" + }, + { + "id": "GO-2026-5020", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.52.0", + "symbolReachable": false, + "summary": "Invoking infinite loop on large channel writes in golang.org/x/crypto/ssh" + }, + { + "id": "GO-2026-5021", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.52.0", + "symbolReachable": false, + "summary": "Invoking auth bypass via unenforced @revoked status in golang.org/x/crypto/ssh/knownhosts" + }, + { + "id": "GO-2026-5023", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.52.0", + "symbolReachable": false, + "summary": "Invoking VerifiedPublicKeyCallback permissions skip enforcement in golang.org/x/crypto/ssh" + }, + { + "id": "GO-2026-5033", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.52.0", + "symbolReachable": false, + "summary": "Invoking pathological inputs can lead to client panic in golang.org/x/crypto/ssh/agent" + }, + { + "id": "GO-2026-5932", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": null, + "symbolReachable": false, + "summary": "The golang.org/x/crypto/openpgp package is unmaintained, unsafe by design, and has known security issues" + }, + { + "id": "GO-2026-6303", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.55.0", + "symbolReachable": false, + "summary": "Source-address critical option not enforced for non-public-key auth callbacks in golang.org/x/crypto/ssh" + }, + { + "id": "GO-2026-6354", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.56.0", + "symbolReachable": false, + "summary": "Prevent DoS on deadlocked undecided channel in golang.org/x/crypto/ssh" + }, + { + "id": "GO-2026-6355", + "module": "golang.org/x/crypto", + "found": "v0.31.0", + "fixed": "v0.56.0", + "symbolReachable": false, + "summary": "Prevent DoS on deadlocked established channel in golang.org/x/crypto/ssh" + }, + { + "id": "GO-2025-3955", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.1", + "symbolReachable": false, + "summary": "CrossOriginProtection insecure bypass patterns not limited to exact matches in net/http" + }, + { + "id": "GO-2025-4006", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.2", + "symbolReachable": false, + "summary": "Excessive CPU consumption in ParseAddress in net/mail" + }, + { + "id": "GO-2025-4007", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.3", + "symbolReachable": true, + "summary": "Quadratic complexity when checking name constraints in crypto/x509" + }, + { + "id": "GO-2025-4008", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.2", + "symbolReachable": true, + "summary": "ALPN negotiation error contains attacker controlled information in crypto/tls" + }, + { + "id": "GO-2025-4009", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.2", + "symbolReachable": false, + "summary": "Quadratic complexity when parsing some invalid inputs in encoding/pem" + }, + { + "id": "GO-2025-4010", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.2", + "symbolReachable": true, + "summary": "Insufficient validation of bracketed IPv6 hostnames in net/url" + }, + { + "id": "GO-2025-4011", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.2", + "symbolReachable": true, + "summary": "Parsing DER payload can cause memory exhaustion in encoding/asn1" + }, + { + "id": "GO-2025-4012", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.2", + "symbolReachable": true, + "summary": "Lack of limit when parsing cookies can cause memory exhaustion in net/http" + }, + { + "id": "GO-2025-4013", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.2", + "symbolReachable": true, + "summary": "Panic when validating certificates with DSA public keys in crypto/x509" + }, + { + "id": "GO-2025-4014", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.2", + "symbolReachable": false, + "summary": "Unbounded allocation when parsing GNU sparse map in archive/tar" + }, + { + "id": "GO-2025-4015", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.2", + "symbolReachable": false, + "summary": "Excessive CPU consumption in Reader.ReadResponse in net/textproto" + }, + { + "id": "GO-2025-4155", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.5", + "symbolReachable": true, + "summary": "Excessive resource consumption when printing error string for host certificate validation in crypto/x509" + }, + { + "id": "GO-2025-4175", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.5", + "symbolReachable": true, + "summary": "Improper application of excluded DNS name constraints when verifying wildcard names in crypto/x509" + }, + { + "id": "GO-2026-4337", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.7", + "symbolReachable": true, + "summary": "Unexpected session resumption in crypto/tls" + }, + { + "id": "GO-2026-4340", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.6", + "symbolReachable": true, + "summary": "Handshake messages may be processed at the incorrect encryption level in crypto/tls" + }, + { + "id": "GO-2026-4341", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.6", + "symbolReachable": false, + "summary": "Memory exhaustion in query parameter parsing in net/url" + }, + { + "id": "GO-2026-4342", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.6", + "symbolReachable": false, + "summary": "Excessive CPU consumption when building archive index in archive/zip" + }, + { + "id": "GO-2026-4601", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.8", + "symbolReachable": true, + "summary": "Incorrect parsing of IPv6 host literals in net/url" + }, + { + "id": "GO-2026-4602", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.8", + "symbolReachable": true, + "summary": "FileInfo can escape from a Root in os" + }, + { + "id": "GO-2026-4603", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.8", + "symbolReachable": false, + "summary": "URLs in meta content attribute actions are not escaped in html/template" + }, + { + "id": "GO-2026-4864", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.9", + "symbolReachable": false, + "summary": "TOCTOU permits root escape on Linux via Root.Chmod in os in internal/syscall/unix" + }, + { + "id": "GO-2026-4865", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.9", + "symbolReachable": false, + "summary": "JsBraceDepth Context Tracking Bugs (XSS) in html/template" + }, + { + "id": "GO-2026-4869", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.9", + "symbolReachable": false, + "summary": "Unbounded allocation for old GNU sparse in archive/tar" + }, + { + "id": "GO-2026-4870", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.9", + "symbolReachable": true, + "summary": "Unauthenticated TLS 1.3 KeyUpdate record can cause persistent connection retention and DoS in crypto/tls" + }, + { + "id": "GO-2026-4918", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.10", + "symbolReachable": false, + "summary": "Infinite loop in HTTP/2 transport when given bad SETTINGS_MAX_FRAME_SIZE in net/http/internal/http2 in golang.org/x/net" + }, + { + "id": "GO-2026-4946", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.9", + "symbolReachable": true, + "summary": "Inefficient policy validation in crypto/x509" + }, + { + "id": "GO-2026-4947", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.9", + "symbolReachable": true, + "summary": "Unexpected work during chain building in crypto/x509" + }, + { + "id": "GO-2026-4970", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.12", + "symbolReachable": false, + "summary": "Root escape via symlink plus trailing slash in os" + }, + { + "id": "GO-2026-4971", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.10", + "symbolReachable": true, + "summary": "Panic in Dial and LookupPort when handling NUL byte on Windows in net" + }, + { + "id": "GO-2026-4976", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.10", + "symbolReachable": false, + "summary": "ReverseProxy forwards queries with more than urlmaxqueryparams parameters in net/http/httputil" + }, + { + "id": "GO-2026-4977", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.10", + "symbolReachable": false, + "summary": "Quadratic string concatenation in consumePhrase in net/mail" + }, + { + "id": "GO-2026-4980", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.10", + "symbolReachable": false, + "summary": "Escaper bypass leads to XSS in html/template" + }, + { + "id": "GO-2026-4981", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.10", + "symbolReachable": false, + "summary": "Crash when handling long CNAME response in net" + }, + { + "id": "GO-2026-4982", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.10", + "symbolReachable": false, + "summary": "Bypass of meta content URL escaping causes XSS in html/template" + }, + { + "id": "GO-2026-4986", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.10", + "symbolReachable": false, + "summary": "Quadratic string concatentation in consumeComment in net/mail" + }, + { + "id": "GO-2026-5026", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.13", + "symbolReachable": false, + "summary": "Invoking failure to reject ASCII-only Punycode-encoded labels in golang.org/x/net/idna" + }, + { + "id": "GO-2026-5037", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.11", + "symbolReachable": true, + "summary": "Inefficient candidate hostname parsing in crypto/x509" + }, + { + "id": "GO-2026-5038", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.11", + "symbolReachable": false, + "summary": "Quadratic complexity in WordDecoder.DecodeHeader in mime" + }, + { + "id": "GO-2026-5039", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.11", + "symbolReachable": true, + "summary": "Arbitrary inputs are included in errors without any escaping in net/textproto" + }, + { + "id": "GO-2026-5856", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.12", + "symbolReachable": true, + "summary": "Invoking Encrypted Client Hello privacy leak in crypto/tls" + }, + { + "id": "GO-2026-5972", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.13", + "symbolReachable": true, + "summary": "Enforce maximum recursion depth in encoding/asn1" + }, + { + "id": "GO-2026-6088", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.13", + "symbolReachable": false, + "summary": "Add recursion depth guard during decode in encoding/xml" + }, + { + "id": "GO-2026-6089", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.13", + "symbolReachable": false, + "summary": "Apply ReadHeaderTimeout when doing unencrypted HTTP/2 check in net/http" + }, + { + "id": "GO-2026-6090", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.13", + "symbolReachable": true, + "summary": "Limit handshake messages we are willing to accept post-handshake in crypto/tls" + }, + { + "id": "GO-2026-6091", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.13", + "symbolReachable": false, + "summary": "Fix Javascript regexp context tracking in html/template" + }, + { + "id": "GO-2026-6218", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.25.13", + "symbolReachable": false, + "summary": "Avoid quadratic complexity in resolvePath in net/url" + }, + { + "id": "GO-2026-6599", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.9", + "symbolReachable": false, + "summary": "Reset context tracking on consecutive template expressions in html/template" + }, + { + "id": "GO-2026-6600", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.9", + "symbolReachable": false, + "summary": "Recognize yield as regexp preceder keyword in html/template" + }, + { + "id": "GO-2026-6603", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.9", + "symbolReachable": true, + "summary": "HTTP/2 server memory exhaustion due to Trailer headers in net/http" + }, + { + "id": "GO-2026-6604", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.9", + "symbolReachable": false, + "summary": "Root.Mkdir(All) can follow junctions out of the root on Windows in os" + }, + { + "id": "GO-2026-6605", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.9", + "symbolReachable": false, + "summary": "HTTP/1 client connection desynchronization after CONNECT rejection in net/http" + }, + { + "id": "GO-2026-6607", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.9", + "symbolReachable": true, + "summary": "Reject malformed ECH outer extension references in crypto/tls" + }, + { + "id": "GO-2026-6608", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.9", + "symbolReachable": true, + "summary": "Memory limit bypass when parsing MIME headers in net/textproto, mime/multipart" + }, + { + "id": "GO-2026-6609", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.9", + "symbolReachable": false, + "summary": "Lack of limit on size of parsed Range headers in net/http" + }, + { + "id": "GO-2026-6610", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.9", + "symbolReachable": false, + "summary": "HTTP/2 transport accepts malformed framing-related headers in net/http" + }, + { + "id": "GO-2026-6611", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.9", + "symbolReachable": true, + "summary": "Excessive CPU consumption from repeated initial window changes in net/http" + }, + { + "id": "GO-2026-6612", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.9", + "symbolReachable": true, + "summary": "Double flow control refund on HTTP/2 server streams in net/http" + }, + { + "id": "GO-2026-6613", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.9", + "symbolReachable": true, + "summary": "HTTP/1 server connection desynchronization after 2xx CONNECT response in net/http" + }, + { + "id": "GO-2026-6617", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.9", + "symbolReachable": true, + "summary": "HTTP/2 server crash due to HPACK encoder race in net/http" + }, + { + "id": "GO-2026-6809", + "module": "stdlib", + "found": "v1.25.0", + "fixed": "v1.26.10", + "symbolReachable": true, + "summary": "Distrust compromised intermediate certificates in crypto/x509" + }, + { + "id": "GO-2025-3503", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.36.0", + "symbolReachable": false, + "summary": "HTTP Proxy bypass using IPv6 Zone IDs in golang.org/x/net" + }, + { + "id": "GO-2025-3595", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.38.0", + "symbolReachable": false, + "summary": "Incorrect Neutralization of Input During Web Page Generation in x/net in golang.org/x/net" + }, + { + "id": "GO-2026-4440", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.45.0", + "symbolReachable": false, + "summary": "Quadratic parsing complexity in golang.org/x/net/html" + }, + { + "id": "GO-2026-4441", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.45.0", + "symbolReachable": false, + "summary": "Infinite parsing loop in golang.org/x/net" + }, + { + "id": "GO-2026-4918", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.53.0", + "symbolReachable": false, + "summary": "Infinite loop in HTTP/2 transport when given bad SETTINGS_MAX_FRAME_SIZE in net/http/internal/http2 in golang.org/x/net" + }, + { + "id": "GO-2026-5025", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.55.0", + "symbolReachable": false, + "summary": "Invoking incorrect handling of namespaced elements in foreign content in golang.org/x/net/html" + }, + { + "id": "GO-2026-5026", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.55.0", + "symbolReachable": false, + "summary": "Invoking failure to reject ASCII-only Punycode-encoded labels in golang.org/x/net/idna" + }, + { + "id": "GO-2026-5027", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.55.0", + "symbolReachable": false, + "summary": "Invoking incorrect handling of HTML elements in foreign content in golang.org/x/net/html" + }, + { + "id": "GO-2026-5028", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.55.0", + "symbolReachable": false, + "summary": "Invoking denial of service when parsing arbitrary HTML in golang.org/x/net/html" + }, + { + "id": "GO-2026-5029", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.55.0", + "symbolReachable": false, + "summary": "Invoking incorrect handling of character references in DOCTYPE nodes in golang.org/x/net/html" + }, + { + "id": "GO-2026-5030", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.55.0", + "symbolReachable": false, + "summary": "Invoking duplicate attributes can cause XSS in golang.org/x/net/html" + }, + { + "id": "GO-2026-5942", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.56.0", + "symbolReachable": false, + "summary": "Parsing an invalid SVCB or HTTPS RR can panic in golang.org/x/net/dns/dnsmessage" + }, + { + "id": "GO-2026-6603", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.60.0", + "symbolReachable": true, + "summary": "HTTP/2 server memory exhaustion due to Trailer headers in net/http" + }, + { + "id": "GO-2026-6610", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.60.0", + "symbolReachable": false, + "summary": "HTTP/2 transport accepts malformed framing-related headers in net/http" + }, + { + "id": "GO-2026-6611", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.60.0", + "symbolReachable": true, + "summary": "Excessive CPU consumption from repeated initial window changes in net/http" + }, + { + "id": "GO-2026-6612", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.60.0", + "symbolReachable": true, + "summary": "Double flow control refund on HTTP/2 server streams in net/http" + }, + { + "id": "GO-2026-6617", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.60.0", + "symbolReachable": true, + "summary": "HTTP/2 server crash due to HPACK encoder race in net/http" + }, + { + "id": "GO-2026-6860", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.62.0", + "symbolReachable": false, + "summary": "Avoid quadratic complexity when coalescing text nodes in html in golang.org/x/net" + }, + { + "id": "GO-2026-6866", + "module": "golang.org/x/net", + "found": "v0.33.0", + "fixed": "v0.62.0", + "symbolReachable": false, + "summary": "Check all ancestor namespace transitions when rendering raw text in html in golang.org/x/net" + }, + { + "id": "GO-2026-5024", + "module": "golang.org/x/sys", + "found": "v0.28.0", + "fixed": "v0.44.0", + "symbolReachable": false, + "summary": "Invoking integer overflow in NewNTUnicodeString in golang.org/x/sys/windows" + }, + { + "id": "GO-2026-5970", + "module": "golang.org/x/text", + "found": "v0.21.0", + "fixed": "v0.39.0", + "symbolReachable": false, + "summary": "Infinite loop on invalid input in golang.org/x/text" + }, + { + "id": "GO-2026-6629", + "module": "golang.org/x/text", + "found": "v0.21.0", + "fixed": "v0.41.0", + "symbolReachable": false, + "summary": "Panic parsing crafted input in x/text/secure/precis in golang.org/x/text" + } + ] +} diff --git a/reports/websocket-vuln-patched-go.txt b/reports/websocket-vuln-patched-go.txt new file mode 100644 index 0000000..55c0165 --- /dev/null +++ b/reports/websocket-vuln-patched-go.txt @@ -0,0 +1,83 @@ +=== Symbol Results === + +Vulnerability #1: GO-2026-6617 + HTTP/2 server crash due to HPACK encoder race in net/http + More info: https://pkg.go.dev/vuln/GO-2026-6617 + Module: golang.org/x/net + Found in: golang.org/x/net@v0.33.0 + Fixed in: golang.org/x/net@v0.60.0 + Example traces found: + #1: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.ErrCode.String + #2: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.FrameType.String + #3: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.Setting.String + #4: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.SettingID.String + #5: ws/server.go:320:31: ws.Server.Accept calls websocket.Upgrader.Upgrade, which eventually calls http2.bufferedWriterTimeoutWriter.Write + #6: ws/server.go:320:31: ws.Server.Accept calls websocket.Upgrader.Upgrade, which eventually calls http2.chunkWriter.Write + #7: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.duplicatePseudoHeaderError.Error + #8: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.headerFieldNameError.Error + #9: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.headerFieldValueError.Error + #10: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.pseudoHeaderError.Error + #11: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.writeData.String + +Vulnerability #2: GO-2026-6612 + Double flow control refund on HTTP/2 server streams in net/http + More info: https://pkg.go.dev/vuln/GO-2026-6612 + Module: golang.org/x/net + Found in: golang.org/x/net@v0.33.0 + Fixed in: golang.org/x/net@v0.60.0 + Example traces found: + #1: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.ErrCode.String + #2: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.FrameType.String + #3: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.Setting.String + #4: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.SettingID.String + #5: ws/server.go:320:31: ws.Server.Accept calls websocket.Upgrader.Upgrade, which eventually calls http2.bufferedWriterTimeoutWriter.Write + #6: ws/server.go:320:31: ws.Server.Accept calls websocket.Upgrader.Upgrade, which eventually calls http2.chunkWriter.Write + #7: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.duplicatePseudoHeaderError.Error + #8: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.headerFieldNameError.Error + #9: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.headerFieldValueError.Error + #10: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.pseudoHeaderError.Error + #11: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.writeData.String + +Vulnerability #3: GO-2026-6611 + Excessive CPU consumption from repeated initial window changes in net/http + More info: https://pkg.go.dev/vuln/GO-2026-6611 + Module: golang.org/x/net + Found in: golang.org/x/net@v0.33.0 + Fixed in: golang.org/x/net@v0.60.0 + Example traces found: + #1: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.ErrCode.String + #2: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.FrameType.String + #3: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.Setting.String + #4: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.SettingID.String + #5: ws/server.go:320:31: ws.Server.Accept calls websocket.Upgrader.Upgrade, which eventually calls http2.bufferedWriterTimeoutWriter.Write + #6: ws/server.go:320:31: ws.Server.Accept calls websocket.Upgrader.Upgrade, which eventually calls http2.chunkWriter.Write + #7: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.duplicatePseudoHeaderError.Error + #8: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.headerFieldNameError.Error + #9: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.headerFieldValueError.Error + #10: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.pseudoHeaderError.Error + #11: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.writeData.String + +Vulnerability #4: GO-2026-6603 + HTTP/2 server memory exhaustion due to Trailer headers in net/http + More info: https://pkg.go.dev/vuln/GO-2026-6603 + Module: golang.org/x/net + Found in: golang.org/x/net@v0.33.0 + Fixed in: golang.org/x/net@v0.60.0 + Example traces found: + #1: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.ErrCode.String + #2: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.FrameType.String + #3: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.Setting.String + #4: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.SettingID.String + #5: ws/server.go:320:31: ws.Server.Accept calls websocket.Upgrader.Upgrade, which eventually calls http2.bufferedWriterTimeoutWriter.Write + #6: ws/server.go:320:31: ws.Server.Accept calls websocket.Upgrader.Upgrade, which eventually calls http2.chunkWriter.Write + #7: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.duplicatePseudoHeaderError.Error + #8: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.headerFieldNameError.Error + #9: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.headerFieldValueError.Error + #10: ws/rpc/peer.go:234:29: rpc.Peer.Notify calls json.Marshal, which eventually calls http2.pseudoHeaderError.Error + #11: ws/errors.go:36:20: ws.HandshakeError.Error calls fmt.Sprintf, which eventually calls http2.writeData.String + +Your code is affected by 4 vulnerabilities from 1 module. +This scan also found 14 vulnerabilities in packages you import and 25 +vulnerabilities in modules you require, but your code doesn't appear to call +these vulnerabilities. +Use '-show verbose' for more details. diff --git a/scripts/verifyWebsocketRelease.py b/scripts/verifyWebsocketRelease.py new file mode 100644 index 0000000..7b2deab --- /dev/null +++ b/scripts/verifyWebsocketRelease.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Verify a fixed WebSocket release and generated routes from an empty module cache.""" +import argparse +import hashlib +import json +import os +from pathlib import Path +import re +import tempfile +import time + +from verifyErrorInheritanceRelease import run + +ROOT = Path(__file__).resolve().parents[1] +MODULE = "git.yuntuo.hk/k2safe/go-omni" + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--version", required=True) + parser.add_argument("--report", type=Path, default=ROOT / "reports/websocket-remote.json") + args = parser.parse_args() + if not re.fullmatch(r"v\d+\.\d+\.\d+(?:-[0-9A-Za-z.-]+)?", args.version): + parser.error("a fixed release version is required") + env = {k: v for k, v in os.environ.items() if not k.startswith("GO_OMNI_")} + env.update(GOWORK="off", GOFLAGS="", GOPRIVATE="git.yuntuo.hk", GONOSUMDB="git.yuntuo.hk") + started = time.monotonic() + with tempfile.TemporaryDirectory(prefix="goOmniWebsocketRelease-") as directory: + temp = Path(directory) + (temp / "bin").mkdir() + env.update(GOBIN=str(temp / "bin"), GOMODCACHE=str(temp / "modcache")) + print(f"Installing remote CLI {args.version} in an empty cache", flush=True) + run(["go", "install", MODULE + "/cmd/go-omni@" + args.version], temp, env) + cli = temp / "bin/go-omni" + starter = temp / "starter" + # Verify the installed default, without --framework-version or a local replace. + run([cli, "new", "--module", "example.com/websocketRelease", starter], temp, env) + run(["go", "mod", "tidy"], starter, env) + dependency = json.loads(run(["go", "list", "-m", "-json", MODULE], starter, env)) + assert dependency["Version"] == args.version and dependency.get("Sum") and not dependency.get("Replace") + downloaded = Path(dependency["Dir"]) + # Add the optional WS dependency only after checking the CLI's default. + run(["go", "get", MODULE + "/ws@" + args.version], starter, env) + print("Testing downloaded WS, Hub, RPC, observation and App integration with race detection", flush=True) + output = run(["go", "test", "-mod=mod", "-race", "-count=1", "-json", "-timeout=180s", + MODULE + "/ws/...", MODULE + "/examples/websocketRPC"], starter, env) + events = [json.loads(line) for line in output.splitlines() if line.startswith("{")] + packages = sorted({e["Package"] for e in events if e.get("Action") == "pass" and "Test" not in e}) + expected = [MODULE + "/ws" + suffix for suffix in ("", "/hub", "/observe", "/omni", "/rpc")] + assert set(expected + [MODULE + "/examples/websocketRPC"]) == set(packages), output + tests = sorted(e["Package"] + ":" + e["Test"] for e in events + if e.get("Action") == "pass" and "Test" in e and "/" not in e["Test"]) + echo = run(["go", "run", "-mod=mod", MODULE + "/examples/websocket"], starter, env) + assert "hello websocket" in echo and "hello group" in echo, echo + rpc = run(["go", "run", "-mod=mod", MODULE + "/examples/websocketRPC"], starter, env) + for generation in (1, 2): + assert f'connection {generation}: double=42, remote callback="generic Go client", restored=1, event={generation}' in rpc, rpc + env.update(GO_OMNI_SCAFFOLD_INTEGRATION="1", GO_OMNI_RELEASE_CLI=str(cli), GO_OMNI_RELEASE_VERSION=args.version) + print("Testing generated custom routes and failure cleanup against the fixed remote release", flush=True) + output = run(["go", "test", "-race", "-count=1", "-json", "-timeout=180s", "./cmd/go-omni/apigen", + "-run", "^TestGeneratedCustomRoutesConsumer$"], downloaded, env) + events = [json.loads(line) for line in output.splitlines() if line.startswith("{")] + assert any(e.get("Action") == "pass" and e.get("Test") == "TestGeneratedCustomRoutesConsumer" for e in events), output + evidence = "".join(e.get("Output", "") for e in events) + assert f'checksum {dependency["Sum"]}; no local replace' in evidence, evidence + final_dependency = json.loads(run(["go", "list", "-m", "-json", MODULE], starter, env)) + assert final_dependency["Version"] == args.version and final_dependency["Sum"] == dependency["Sum"] + assert not final_dependency.get("Replace") + report = { + "passed": True, "frameworkVersion": args.version, "frameworkSum": dependency["Sum"], + "toolchain": run(["go", "version"], temp, env).strip(), + "freshModuleCache": True, "noLocalReplace": True, + "durationSeconds": round(time.monotonic() - started, 2), + "packages": packages, "tests": tests, + "checks": ["remoteCLIInstall", "installedVersionDefault", "fixedRemoteModuleChecksum", + "downloadedWebsocketRaceTests", "remoteEchoAndGroupExample", "remoteBidirectionalRPCExample", + "remoteTwoConnectionSubscriptionRestore", "generatedCustomRoutesConsumer", + "oldServiceWithoutHook", "handwrittenHookPreserved", "customRouteFailureCleanup"], + "examples": {"websocket": echo.strip(), "websocketRPC": rpc.strip()}, + "sourceHashes": {name: hashlib.sha256((downloaded / name).read_bytes()).hexdigest() for name in + ("scripts/verifyWebsocketRelease.py", "cmd/go-omni/apigen/customRoutes_test.go")}, + "limitations": ["Functional release acceptance; not a vulnerability scan or browser UI test.", + "Known framework HTTP/2 dependency vulnerabilities remain documented in the validation report."], + } + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n") + print(f"PASS remote WebSocket release {args.version}; temporary files removed", flush=True) + + +if __name__ == "__main__": + main() diff --git a/ws/README.md b/ws/README.md new file mode 100644 index 0000000..d60ec46 --- /dev/null +++ b/ws/README.md @@ -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 由业务协议显式传播。 diff --git a/ws/client.go b/ws/client.go new file mode 100644 index 0000000..0f6c66c --- /dev/null +++ b/ws/client.go @@ -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} +} diff --git a/ws/client_test.go b/ws/client_test.go new file mode 100644 index 0000000..117bb89 --- /dev/null +++ b/ws/client_test.go @@ -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, "") +} diff --git a/ws/config.go b/ws/config.go new file mode 100644 index 0000000..18ada1c --- /dev/null +++ b/ws/config.go @@ -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 +} diff --git a/ws/errors.go b/ws/errors.go new file mode 100644 index 0000000..2a4c624 --- /dev/null +++ b/ws/errors.go @@ -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 } diff --git a/ws/hub/hub.go b/ws/hub/hub.go new file mode 100644 index 0000000..355ba31 --- /dev/null +++ b/ws/hub/hub.go @@ -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 +} diff --git a/ws/hub/hub_test.go b/ws/hub/hub_test.go new file mode 100644 index 0000000..65113ff --- /dev/null +++ b/ws/hub/hub_test.go @@ -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") + } + } + }) + } +} diff --git a/ws/limits_test.go b/ws/limits_test.go new file mode 100644 index 0000000..a5b0b22 --- /dev/null +++ b/ws/limits_test.go @@ -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) + } +} diff --git a/ws/observe/README.md b/ws/observe/README.md new file mode 100644 index 0000000..8686124 --- /dev/null +++ b/ws/observe/README.md @@ -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 +``` diff --git a/ws/observe/observe.go b/ws/observe/observe.go new file mode 100644 index 0000000..a7ed94d --- /dev/null +++ b/ws/observe/observe.go @@ -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 +} diff --git a/ws/observe/observe_test.go b/ws/observe/observe_test.go new file mode 100644 index 0000000..fda5665 --- /dev/null +++ b/ws/observe/observe_test.go @@ -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) + } +} diff --git a/ws/omni/handler.go b/ws/omni/handler.go new file mode 100644 index 0000000..550d440 --- /dev/null +++ b/ws/omni/handler.go @@ -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 +} diff --git a/ws/omni/handler_test.go b/ws/omni/handler_test.go new file mode 100644 index 0000000..00bdda8 --- /dev/null +++ b/ws/omni/handler_test.go @@ -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) + } +} diff --git a/ws/rpc/README.md b/ws/rpc/README.md new file mode 100644 index 0000000..89b1d77 --- /dev/null +++ b/ws/rpc/README.md @@ -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` 测协议解析与严格输入校验。 diff --git a/ws/rpc/config.go b/ws/rpc/config.go new file mode 100644 index 0000000..cf3d2db --- /dev/null +++ b/ws/rpc/config.go @@ -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 } diff --git a/ws/rpc/peer.go b/ws/rpc/peer.go new file mode 100644 index 0000000..6032a27 --- /dev/null +++ b/ws/rpc/peer.go @@ -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() + } +} diff --git a/ws/rpc/peer_test.go b/ws/rpc/peer_test.go new file mode 100644 index 0000000..9a86cdd --- /dev/null +++ b/ws/rpc/peer_test.go @@ -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() +} diff --git a/ws/rpc/protocol.go b/ws/rpc/protocol.go new file mode 100644 index 0000000..9790d12 --- /dev/null +++ b/ws/rpc/protocol.go @@ -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 +} diff --git a/ws/rpc/subscriptions.go b/ws/rpc/subscriptions.go new file mode 100644 index 0000000..eb14447 --- /dev/null +++ b/ws/rpc/subscriptions.go @@ -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 +} diff --git a/ws/server.go b/ws/server.go new file mode 100644 index 0000000..02efcd4 --- /dev/null +++ b/ws/server.go @@ -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} +} diff --git a/ws/server_test.go b/ws/server_test.go new file mode 100644 index 0000000..4e46e33 --- /dev/null +++ b/ws/server_test.go @@ -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) + } + } + }) + } +} diff --git a/ws/session.go b/ws/session.go new file mode 100644 index 0000000..b2f2ce0 --- /dev/null +++ b/ws/session.go @@ -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) +} diff --git a/ws/session_test.go b/ws/session_test.go new file mode 100644 index 0000000..01dc711 --- /dev/null +++ b/ws/session_test.go @@ -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 +} diff --git a/ws/testdata/fuzz/FuzzOrigin/1701c6390d7ef919 b/ws/testdata/fuzz/FuzzOrigin/1701c6390d7ef919 new file mode 100644 index 0000000..7c9b702 --- /dev/null +++ b/ws/testdata/fuzz/FuzzOrigin/1701c6390d7ef919 @@ -0,0 +1,2 @@ +go test fuzz v1 +string("https://::443")