add auto websocket

This commit is contained in:
wjqserver 2025-05-30 21:32:22 +08:00
parent d6aa97bbe5
commit dd4afc276b
3 changed files with 133 additions and 0 deletions

1
go.mod
View file

@ -6,6 +6,7 @@ require (
github.com/WJQSERVER-STUDIO/go-utils/copyb v0.0.4
github.com/WJQSERVER-STUDIO/httpc v0.5.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

2
go.sum
View file

@ -4,5 +4,7 @@ github.com/WJQSERVER-STUDIO/httpc v0.5.1 h1:+TKCPYBuj7PAHuiduGCGAqsHAa4QtsUfoVwR
github.com/WJQSERVER-STUDIO/httpc v0.5.1/go.mod h1:M7KNUZjjhCkzzcg9lBPs9YfkImI+7vqjAyjdA19+joE=
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=

130
ws.go Normal file
View file

@ -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/2Upgrade() 方法不会发送 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 逻辑
}
}