diff --git a/go.mod b/go.mod index 6c71120..dbeb743 100644 --- a/go.mod +++ b/go.mod @@ -4,9 +4,10 @@ go 1.24.4 require ( github.com/WJQSERVER-STUDIO/go-utils/copyb v0.0.4 - github.com/WJQSERVER-STUDIO/httpc v0.6.0 + github.com/WJQSERVER-STUDIO/httpc v0.5.1 github.com/fenthope/reco v0.0.1 github.com/go-json-experiment/json v0.0.0-20250517221953-25912455fbc8 + github.com/gorilla/websocket v1.5.3 ) require github.com/valyala/bytebufferpool v1.0.0 // indirect diff --git a/go.sum b/go.sum index e5d2a67..69661e4 100644 --- a/go.sum +++ b/go.sum @@ -1,10 +1,12 @@ github.com/WJQSERVER-STUDIO/go-utils/copyb v0.0.4 h1:JLtFd00AdFg/TP+dtvIzLkdHwKUGPOAijN1sMtEYoFg= github.com/WJQSERVER-STUDIO/go-utils/copyb v0.0.4/go.mod h1:FZ6XE+4TKy4MOfX1xWKe6Rwsg0ucYFCdNh1KLvyKTfc= -github.com/WJQSERVER-STUDIO/httpc v0.6.0 h1:GXKc6BNGn5ALdDLkCsM+mP24jyi1S9QBLL2XQ1a2sPM= -github.com/WJQSERVER-STUDIO/httpc v0.6.0/go.mod h1:M7KNUZjjhCkzzcg9lBPs9YfkImI+7vqjAyjdA19+joE= +github.com/WJQSERVER-STUDIO/httpc v0.5.1 h1:+TKCPYBuj7PAHuiduGCGAqsHAa4QtsUfoVwRN777q64= +github.com/WJQSERVER-STUDIO/httpc v0.5.1/go.mod h1:M7KNUZjjhCkzzcg9lBPs9YfkImI+7vqjAyjdA19+joE= github.com/fenthope/reco v0.0.1 h1:GYcuXCEKYoctD0dFkiBC+t0RMTOyOiujBCin8bbLR3Y= github.com/fenthope/reco v0.0.1/go.mod h1:mDkGLHte5udWTIcjQTxrABRcf56SSdxBOCLgrRDwI/Y= github.com/go-json-experiment/json v0.0.0-20250517221953-25912455fbc8 h1:o8UqXPI6SVwQt04RGsqKp3qqmbOfTNMqDrWsc4O47kk= github.com/go-json-experiment/json v0.0.0-20250517221953-25912455fbc8/go.mod h1:TiCD2a1pcmjd7YnhGH0f/zKNcCD06B029pHhzV23c2M= +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/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw= github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= diff --git a/ws.go b/ws.go new file mode 100644 index 0000000..0ee0ee3 --- /dev/null +++ b/ws.go @@ -0,0 +1,130 @@ +package touka + +import ( + "errors" + "log" + "net/http" + + "github.com/gorilla/websocket" +) + +// WebSocketHandler 是用户提供的用于处理 WebSocket 连接的函数类型。 +// conn 是一个已经完成握手的 WebSocket 连接。 +type WebSocketHandler func(c *Context, conn *websocket.Conn) + +// WebSocketUpgradeOptions 用于配置 WebSocket 升级中间件。 +type WebSocketUpgradeOptions struct { + // Upgrader 是 gorilla/websocket.Upgrader 的实例。 + // 用户可以配置 ReadBufferSize, WriteBufferSize, CheckOrigin, Subprotocols 等。 + // 如果为 nil,将使用一个带有合理默认值的 Upgrader。 + Upgrader *websocket.Upgrader + + // Handler 是在 WebSocket 成功升级后调用的处理函数。 + // 这个字段是必需的。 + Handler WebSocketHandler + + // OnError 是一个可选的回调函数,用于处理升级过程中发生的错误。 + // 如果未提供,错误将导致一个标准的 HTTP 错误响应(例如 400 Bad Request)。 + OnError func(c *Context, status int, err error) +} + +// defaultWebSocketUpgrader 返回一个具有合理默认值的 websocket.Upgrader。 +func defaultWebSocketUpgrader() *websocket.Upgrader { + return &websocket.Upgrader{ + ReadBufferSize: 1024, + WriteBufferSize: 1024, + // CheckOrigin 应该由用户根据其安全需求来配置。 + // 默认情况下,如果 Origin 头部存在且与 Host 头部不匹配,会拒绝连接。 + // 对于开发,可以暂时设置为 func(r *http.Request) bool { return true } + // 但在生产环境中必须小心配置。 + CheckOrigin: func(r *http.Request) bool { + // 简单的同源检查或允许所有 (根据需要调整) + // return r.Header.Get("Origin") == "" || strings.HasPrefix(r.Header.Get("Origin"), "http://"+r.Host) || strings.HasPrefix(r.Header.Get("Origin"), "https://"+r.Host) + return true // 示例:允许所有,生产环境请谨慎 + }, + } +} + +// defaultWebSocketOnError 是默认的错误处理函数。 +func defaultWebSocketOnError(c *Context, status int, err error) { + // 使用框架的错误处理机制或简单的字符串响应 + // 确保不要写入一个已经开始的响应 + if !c.Writer.Written() { + // 返回英文错误信息 + errMsg := http.StatusText(status) + if err != nil { + errMsg = err.Error() // 可以考虑是否暴露详细错误 + } + c.String(status, "%s", errMsg) // 或者 c.engine.errorHandle.handler(c, status) + } + c.Abort() // 总是中止 +} + +// WebSocketUpgrade 返回一个 WebSocket 升级中间件。 +// 它能自动感知 HTTP/1.1 的 Upgrade 请求和 HTTP/2 的扩展 CONNECT 请求 (RFC 8441)。 +func WebSocketUpgrade(opts WebSocketUpgradeOptions) HandlerFunc { + if opts.Handler == nil { + panic("touka: WebSocketUpgradeOptions.Handler cannot be nil") + } + + upgrader := opts.Upgrader + if upgrader == nil { + upgrader = defaultWebSocketUpgrader() + } + + onError := opts.OnError + if onError == nil { + onError = defaultWebSocketOnError + } + + return func(c *Context) { + // 调试日志,查看请求详情 + // reqBytes, _ := httputil.DumpRequest(c.Request, true) + // log.Printf("WebSocketUpgrade: Incoming request for path %s:\n%s", c.Request.URL.Path, string(reqBytes)) + // log.Printf("Request Proto: %s, Method: %s", c.Request.Proto, c.Request.Method) + + // 对于我们的目的,让 gorilla/websocket 的 Upgrade 方法去判断更佳, + // 它已经实现了 RFC 8441 的支持。 + + // 我们不再需要手动区分 HTTP/1.1 和 HTTP/2 的逻辑, + // gorilla/websocket.Upgrader.Upgrade 会自动处理。 + // 它会检查请求是 HTTP/1.1 Upgrade 还是 HTTP/2 CONNECT with :protocol=websocket。 + + // 对于 HTTP/2,Upgrade() 方法不会发送 101,而是处理 CONNECT 的 200 OK。 + // 它也不会调用 Hijack,因为连接已经在 HTTP/2 流上。 + conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + // 升级失败。gorilla/websocket.Upgrade 会处理错误响应的发送。 + // (对于 HTTP/1.1 会是 400/403 等;对于 HTTP/2 也是类似的非 2xx 响应) + var httpErr websocket.HandshakeError + statusCode := http.StatusBadRequest // 默认 + if errors.As(err, &httpErr) { + // 尝试获取更具体的错误信息,但状态码可能不直接暴露 + } + + // 使用英文记录日志 + log.Printf("WebSocket upgrade/handshake failed for %s (Proto: %s): %v", c.Request.RemoteAddr, c.Request.Proto, err) + onError(c, statusCode, err) + if !c.IsAborted() { + c.Abort() + } + return + } + + // 升级/握手成功 + // 使用英文记录日志 + log.Printf("WebSocket connection established for %s (Proto: %s)", c.Request.RemoteAddr, c.Request.Proto) + + if !c.IsAborted() { + c.Abort() // 确保 HTTP 处理链中止 + } + + defer func() { + // 使用英文记录日志 + log.Printf("Closing WebSocket connection for %s", conn.RemoteAddr()) + _ = conn.Close() + }() + + opts.Handler(c, conn) // 执行用户定义的 WebSocket 逻辑 + } +}