mirror of
https://github.com/infinite-iroha/touka.git
synced 2026-02-03 00:41:10 +08:00
add auto websocket
This commit is contained in:
parent
d6aa97bbe5
commit
dd4afc276b
3 changed files with 133 additions and 0 deletions
1
go.mod
1
go.mod
|
|
@ -6,6 +6,7 @@ require (
|
||||||
github.com/WJQSERVER-STUDIO/go-utils/copyb v0.0.4
|
github.com/WJQSERVER-STUDIO/go-utils/copyb v0.0.4
|
||||||
github.com/WJQSERVER-STUDIO/httpc v0.5.1
|
github.com/WJQSERVER-STUDIO/httpc v0.5.1
|
||||||
github.com/go-json-experiment/json v0.0.0-20250517221953-25912455fbc8
|
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
|
require github.com/valyala/bytebufferpool v1.0.0 // indirect
|
||||||
|
|
|
||||||
2
go.sum
2
go.sum
|
|
@ -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/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 h1:o8UqXPI6SVwQt04RGsqKp3qqmbOfTNMqDrWsc4O47kk=
|
||||||
github.com/go-json-experiment/json v0.0.0-20250517221953-25912455fbc8/go.mod h1:TiCD2a1pcmjd7YnhGH0f/zKNcCD06B029pHhzV23c2M=
|
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 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
|
||||||
github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
|
github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
|
||||||
|
|
|
||||||
130
ws.go
Normal file
130
ws.go
Normal 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/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 逻辑
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Add table
Add a link
Reference in a new issue