diff --git a/ecw.go b/ecw.go index 8f1417a..e43db3a 100644 --- a/ecw.go +++ b/ecw.go @@ -1,6 +1,9 @@ package touka import ( + "bufio" + "errors" + "net" "net/http" "sync" ) @@ -44,9 +47,9 @@ func (ecw *errorCapturingResponseWriter) reset(w http.ResponseWriter, r *http.Re // AcquireErrorCapturingResponseWriter 从对象池获取一个 errorCapturingResponseWriter 实例 // 必须在处理完成后调用 ReleaseErrorCapturingResponseWriter -func AcquireErrorCapturingResponseWriter(c *Context, eh ErrorHandler) *errorCapturingResponseWriter { +func AcquireErrorCapturingResponseWriter(c *Context) *errorCapturingResponseWriter { ecw := errorResponseWriterPool.Get().(*errorCapturingResponseWriter) - ecw.reset(c.Writer, c.Request, c, eh) // 传入 Touka Context 的 Writer + ecw.reset(c.Writer, c.Request, c, c.engine.errorHandle.handler) // 传入 Touka Context 的 Writer return ecw } @@ -74,9 +77,9 @@ func (ecw *errorCapturingResponseWriter) WriteHeader(statusCode int) { if ecw.responseStarted { return // 响应已开始, 忽略后续的 WriteHeader 调用 } - ecw.statusCode = statusCode // 总是记录 FileServer 意图的状态码 + ecw.statusCode = statusCode - if statusCode >= http.StatusBadRequest { + if ecw.Status() >= 400 { ecw.capturedErrorSignal = true // 是一个错误状态码 (>=400), 激活错误信号 // 不会将这个 WriteHeader 传递给原始的 w, 等待 processAfterFileServer 处理 @@ -108,7 +111,7 @@ func (ecw *errorCapturingResponseWriter) Write(data []byte) (int, error) { for k, v := range ecw.headerSnapshot { ecw.w.Header()[k] = v // 直接赋值 []string, 保留所有值 } - ecw.w.WriteHeader(ecw.statusCode) // 发送实际的状态码 (可能是 200 或之前设置的 2xx) + ecw.w.WriteHeader(ecw.Status()) // 发送实际的状态码 (可能是 200 或之前设置的 2xx) ecw.responseStarted = true } return ecw.w.Write(data) // 写入数据到原始 ResponseWriter @@ -133,7 +136,7 @@ func (ecw *errorCapturingResponseWriter) processAfterFileServer() { ecw.ctx.Next() } else { // 调用用户自定义的 ErrorHandlerFunc, 由它负责完整的错误响应 - ecw.errorHandlerFunc(ecw.ctx, ecw.statusCode) + ecw.errorHandlerFunc(ecw.ctx, ecw.Status()) ecw.ctx.Abort() } } @@ -141,3 +144,57 @@ func (ecw *errorCapturingResponseWriter) processAfterFileServer() { // 如果 ecw.capturedErrorSignal && ecw.responseStarted, 表示在捕获错误信号之前, // 成功路径的响应已经开始, 此时无法再进行错误处理覆盖 } + +// Status 返回当前记录的状态码 +func (ecw *errorCapturingResponseWriter) Status() int { + if ecw.statusCode == 0 && !ecw.responseStarted { + // 如果还没有显式设置状态码, 并且响应尚未开始, + // 则尝试从底层 ResponseWriter 获取状态码 (如果它实现了 Statuser) + if tw, ok := ecw.w.(ResponseWriter); ok { + return tw.Status() + } + // 否则, 默认返回 200 OK (Go HTTP server 的默认行为) + return http.StatusOK + } + return ecw.statusCode +} + +// Size 返回已写入响应体的字节数 +func (ecw *errorCapturingResponseWriter) Size() int { + // ecw 在捕获错误信号时会丢弃 FileServer 写入的数据, 所以 Size 应返回 0 + if ecw.capturedErrorSignal { + return 0 + } + // 否则, 尝试从底层 ResponseWriter 获取已写入的字节数 + if tw, ok := ecw.w.(ResponseWriter); ok { + return tw.Size() + } + // 对于其他类型的 ResponseWriter, 无法可靠获取, 只能返回 0 + return 0 +} + +// Written方式 +func (ecw *errorCapturingResponseWriter) Written() bool { + // 如果响应已经通过这个包装器开始写入 (WriteHeader 或 Write 成功调用) + // 或者如果原始 ResponseWriter 已经标记为 Written (例如, 如果它是 touka.ResponseWriterImpl) + // 则认为响应已开始 + if ecw.responseStarted { + return true + } + // 检查原始 ResponseWriter 是否已经写入 + if tw, ok := ecw.w.(ResponseWriter); ok { + return tw.Written() + } + // 对于其他类型的 ResponseWriter, 无法可靠判断是否已写入, 只能依赖 responseStarted 标记 + return false +} + +// Hijack 实现 http.Hijacker 接口 +// 它将 Hijack 调用委托给底层的 ResponseWriter +func (ecw *errorCapturingResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + hijacker, ok := ecw.w.(http.Hijacker) + if !ok { + return nil, nil, errors.New("the underlying ResponseWriter does not support the Hijacker interface") + } + return hijacker.Hijack() +} diff --git a/engine.go b/engine.go index c3e4743..b97e90d 100644 --- a/engine.go +++ b/engine.go @@ -2,6 +2,7 @@ package touka import ( "context" + "log" "reflect" "runtime" "strings" @@ -52,7 +53,8 @@ type Engine struct { noRoute HandlerFunc - unMatchFS UnMatchFS // 未匹配下的处理 + unMatchFS UnMatchFS // 未匹配下的处理 + unMatchFileServer http.Handler // 处理handle serverProtocols *http.Protocols //服务协议 Protocols ProtocolsConfig //协议版本配置 @@ -73,6 +75,9 @@ func defaultErrorHandle(c *Context, code int) { // 检查客户端是否已断 return default: + if c.Writer.Written() { + return + } // 输出json 状态码与状态码对应描述 c.JSON(code, H{ "code": code, @@ -84,6 +89,22 @@ func defaultErrorHandle(c *Context, code int) { // 检查客户端是否已断 } } +// 默认errorhandle包装 避免竞争意外问题, 保证稳定性 +func defaultErrorWarp(handler ErrorHandler) ErrorHandler { + return func(c *Context, code int) { + select { + case <-c.Request.Context().Done(): + return + default: + if c.Writer.Written() { + log.Printf("errpage: response already started for status %d, skipping error page rendering", code) + return + } + } + handler(c, code) + } +} + type UnMatchFS struct { FSForUnmatched http.FileSystem ServeUnmatchedAsFS bool @@ -146,7 +167,7 @@ func Default() *Engine { // 设置自定义错误处理 func (engine *Engine) SetErrorHandler(handler ErrorHandler) { engine.errorHandle.useDefault = false - engine.errorHandle.handler = handler + engine.errorHandle.handler = defaultErrorWarp(handler) } // 获取一个默认错误处理handle @@ -159,8 +180,10 @@ func (engine *Engine) SetUnMatchFS(fs http.FileSystem) { if fs != nil { engine.unMatchFS.FSForUnmatched = fs engine.unMatchFS.ServeUnmatchedAsFS = true + engine.unMatchFileServer = http.FileServer(fs) } else { engine.unMatchFS.ServeUnmatchedAsFS = false + engine.unMatchFileServer = nil } } @@ -442,14 +465,20 @@ func (engine *Engine) handleRequest(c *Context) { func unMatchFSHandle() HandlerFunc { return func(c *Context) { engine := c.engine + // 确保 engine.unMatchFileServer 存在 + if !engine.unMatchFS.ServeUnmatchedAsFS || engine.unMatchFileServer == nil { + c.Next() // 如果未配置或 FileSystem 为 nil，则继续处理链 + return + } if c.Request.Method == http.MethodGet || c.Request.Method == http.MethodHead { // 使用 http.FileServer 处理未匹配的请求 - fileServer := http.FileServer(engine.unMatchFS.FSForUnmatched) + //fileServer := http.FileServer(engine.unMatchFS.FSForUnmatched) //ecw := newErrorCapturingResponseWriter(c, c.engine.errorHandle.handler) - ecw := AcquireErrorCapturingResponseWriter(c, c.engine.errorHandle.handler) + ecw := AcquireErrorCapturingResponseWriter(c) defer ReleaseErrorCapturingResponseWriter(ecw) - fileServer.ServeHTTP(ecw, c.Request) + c.engine.unMatchFileServer.ServeHTTP(ecw, c.Request) ecw.processAfterFileServer() + c.Abort() return } else { if engine.noRoute == nil { @@ -726,7 +755,7 @@ func (engine *Engine) Static(relativePath, rootPath string) { // 使用自定义的 ResponseWriter 包装器来捕获 FileServer 可能返回的错误状态码 // 这样我们可以在 FileServer 返回 404 或 403 时，使用 Engine 的 ErrorHandler 进行统一处理 - ecw := AcquireErrorCapturingResponseWriter(c, c.engine.errorHandle.handler) + ecw := AcquireErrorCapturingResponseWriter(c) defer ReleaseErrorCapturingResponseWriter(ecw) // @@ -790,7 +819,7 @@ func (group *RouterGroup) Static(relativePath, rootPath string) { // 使用自定义的 ResponseWriter 包装器来捕获 FileServer 可能返回的错误状态码 // 这样我们可以在 FileServer 返回 404 或 403 时，使用 Engine 的 ErrorHandler 进行统一处理 - ecw := AcquireErrorCapturingResponseWriter(c, group.engine.errorHandle.handler) + ecw := AcquireErrorCapturingResponseWriter(c) defer ReleaseErrorCapturingResponseWriter(ecw) // diff --git a/go.mod b/go.mod index a798659..239260c 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index 6dfccbb..7d9a3dd 100644 --- a/go.sum +++ b/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/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/tgzip.go b/tgzip.go new file mode 100644 index 0000000..419cc6f --- /dev/null +++ b/tgzip.go @@ -0,0 +1,357 @@ +package touka + +import ( + "bufio" + "compress/gzip" + "errors" + "io" + "net" + "net/http" + "strconv" + "strings" + "sync" +) + +const ( + headerAcceptEncoding = "Accept-Encoding" // 请求头部，客户端声明接受的编码 + headerContentEncoding = "Content-Encoding" // 响应头部，服务器声明使用的编码 + headerContentLength = "Content-Length" // 响应头部，内容长度 + headerContentType = "Content-Type" // 响应头部，内容类型 + headerVary = "Vary" // 响应头部，指示缓存行为 + encodingGzip = "gzip" // Gzip 编码名称 +) + +var ( + // 默认可压缩的 MIME 类型 + defaultCompressibleTypes = []string{ + "text/html", "text/css", "text/plain", "text/javascript", + "application/javascript", "application/x-javascript", "application/json", + "application/xml", "image/svg+xml", + } +) + +// GzipOptions 用于配置 Gzip 中间件。 +type GzipOptions struct { + // Level 设置 Gzip 压缩级别。 + // 例如: gzip.DefaultCompression, gzip.BestSpeed, gzip.BestCompression。 + Level int + // MinContentLength 是应用 Gzip 的最小内容长度。 + // 如果响应的 Content-Length 小于此值，则不应用 Gzip。 + // 默认为 0 (无最小长度限制)。 + MinContentLength int64 + // CompressibleTypes 是要压缩的 MIME 类型列表。 + // 如果为空，将使用 defaultCompressibleTypes。 + CompressibleTypes []string + // DecompressFn 是一个可选函数，用于解压缩请求体 (如果请求体是 gzipped)。 + // 如果为 nil，则禁用请求体解压缩。 + // 注意: 本次实现主要关注响应压缩，请求解压可以作为扩展。 + // DecompressFn func(c *Context) +} + +// gzipResponseWriter 包装了 touka.ResponseWriter 以提供 Gzip 压缩功能。 +type gzipResponseWriter struct { + ResponseWriter // 底层的 ResponseWriter (可能是 ecw 或 responseWriterImpl) + gzWriter *gzip.Writer // compress/gzip 的 writer + options *GzipOptions // Gzip 配置 + wroteHeader bool // 标记 Header 是否已写入 + doCompression bool // 标记是否执行压缩 + statusCode int // 存储状态码，在实际写入底层 Writer 前使用 +} + +// --- 对象池 --- +var gzipResponseWriterPool = sync.Pool{ + New: func() interface{} { + return &gzipResponseWriter{} + }, +} + +// gzip.Writer 实例的对象池。 +// 注意: gzip.Writer.Reset() 不会改变压缩级别，所以对象池需要提供已正确初始化级别的 writer。 +// 我们为每个可能的级别创建一个池。 +var gzipWriterPools [gzip.BestCompression - gzip.BestSpeed + 2]*sync.Pool // 覆盖 -1 (Default) 到 9 (BestCompression) + +func init() { + for i := gzip.BestSpeed; i <= gzip.BestCompression; i++ { + level := i // 捕获循环变量用于闭包 + gzipWriterPools[level-gzip.BestSpeed] = &sync.Pool{ + New: func() interface{} { + // 初始化时 writer 为 nil，在 Reset 时设置 + w, _ := gzip.NewWriterLevel(nil, level) + return w + }, + } + } + // 为 gzip.DefaultCompression (-1) 映射一个索引 + defaultLevelIndex := gzip.BestCompression - gzip.BestSpeed + 1 + gzipWriterPools[defaultLevelIndex] = &sync.Pool{ + New: func() interface{} { + w, _ := gzip.NewWriterLevel(nil, gzip.DefaultCompression) + return w + }, + } +} + +// 从对象池获取一个 gzip.Writer +func getGzipWriterFromPool(level int, underlyingWriter io.Writer) *gzip.Writer { + var poolIndex int + if level == gzip.DefaultCompression { + poolIndex = gzip.BestCompression - gzip.BestSpeed + 1 + } else if level >= gzip.BestSpeed && level <= gzip.BestCompression { + poolIndex = level - gzip.BestSpeed + } else { // 无效级别，使用默认级别 + poolIndex = gzip.BestCompression - gzip.BestSpeed + 1 + level = gzip.DefaultCompression // 保证一致性 + } + + gz := gzipWriterPools[poolIndex].Get().(*gzip.Writer) + gz.Reset(underlyingWriter) // 重置并关联到底层的 io.Writer + return gz +} + +// 将 gzip.Writer 返还给对象池 +func putGzipWriterToPool(gz *gzip.Writer, level int) { + var poolIndex int + if level == gzip.DefaultCompression { + poolIndex = gzip.BestCompression - gzip.BestSpeed + 1 + } else if level >= gzip.BestSpeed && level <= gzip.BestCompression { + poolIndex = level - gzip.BestSpeed + } else { // 不应该发生，如果 getGzipWriterFromPool 进行了标准化 + poolIndex = gzip.BestCompression - gzip.BestSpeed + 1 + } + gzipWriterPools[poolIndex].Put(gz) +} + +// 从对象池获取一个 gzipResponseWriter +func acquireGzipResponseWriter(underlying ResponseWriter, opts *GzipOptions) *gzipResponseWriter { + gzw := gzipResponseWriterPool.Get().(*gzipResponseWriter) + gzw.ResponseWriter = underlying + gzw.options = opts + gzw.wroteHeader = false + gzw.doCompression = false + gzw.statusCode = 0 // 重置状态码 + // gzWriter 将在 WriteHeader 中如果需要时获取 + return gzw +} + +// 将 gzipResponseWriter 返还给对象池 +func releaseGzipResponseWriter(gzw *gzipResponseWriter) { + if gzw.gzWriter != nil { + // 确保它被关闭并返回到池中 + _ = gzw.gzWriter.Close() // 关闭会 flush + putGzipWriterToPool(gzw.gzWriter, gzw.options.Level) + gzw.gzWriter = nil + } + gzw.ResponseWriter = nil // 断开引用 + gzw.options = nil + gzipResponseWriterPool.Put(gzw) +} + +// --- gzipResponseWriter 方法实现 --- + +// Header 返回底层 ResponseWriter 的头部 map。 +func (gzw *gzipResponseWriter) Header() http.Header { + return gzw.ResponseWriter.Header() +} + +// WriteHeader 发送 HTTP 响应头部和指定的状态码。 +// 在这里决定是否进行压缩。 +func (gzw *gzipResponseWriter) WriteHeader(statusCode int) { + if gzw.wroteHeader { + return + } + gzw.wroteHeader = true + gzw.statusCode = statusCode // 存储状态码 + + // 在修改头部以进行压缩之前进行条件检查 + // 1. 如果状态码是信息性(1xx)、重定向(3xx)、无内容(204)、重置内容(205)或未修改(304)，则不压缩 + if statusCode < http.StatusOK || statusCode == http.StatusNoContent || statusCode == http.StatusResetContent || statusCode == http.StatusNotModified { + gzw.ResponseWriter.WriteHeader(statusCode) + return + } + // 2. 如果响应已经被编码，则不压缩 + if gzw.Header().Get(headerContentEncoding) != "" { + gzw.ResponseWriter.WriteHeader(statusCode) + return + } + // 3. 检查 Content-Type + contentType := strings.ToLower(strings.TrimSpace(strings.Split(gzw.Header().Get(headerContentType), ";")[0])) + compressibleTypes := gzw.options.CompressibleTypes + if len(compressibleTypes) == 0 { + compressibleTypes = defaultCompressibleTypes + } + isCompressible := false + for _, t := range compressibleTypes { + if strings.HasPrefix(contentType, t) { // 使用 HasPrefix 以匹配如 "text/html; charset=utf-8" + isCompressible = true + break + } + } + if !isCompressible { + gzw.ResponseWriter.WriteHeader(statusCode) + return + } + // 4. 检查 MinContentLength + if gzw.options.MinContentLength > 0 { + if clStr := gzw.Header().Get(headerContentLength); clStr != "" { + if cl, err := strconv.ParseInt(clStr, 10, 64); err == nil && cl < gzw.options.MinContentLength { + gzw.ResponseWriter.WriteHeader(statusCode) + return + } + } + // 如果未设置 Content-Length，但设置了 MinContentLength，我们可能仍会压缩。 + // 这是一个权衡：可能会压缩小的动态内容。 + } + + // 所有检查通过，进行压缩 + gzw.doCompression = true + gzw.Header().Set(headerContentEncoding, encodingGzip) + gzw.Header().Add(headerVary, headerAcceptEncoding) // 使用 Add 以避免覆盖其他 Vary 值 + gzw.Header().Del(headerContentLength) // Gzip 会改变内容长度，所以删除它 + + // 从池中获取 gzWriter，并将其 Reset 指向实际的底层 ResponseWriter + // 注意：gzw.ResponseWriter 是被 Gzip 包装的 writer (例如，原始的 responseWriterImpl 或 ecw) + gzw.gzWriter = getGzipWriterFromPool(gzw.options.Level, gzw.ResponseWriter) + gzw.ResponseWriter.WriteHeader(statusCode) // 调用原始的 WriteHeader +} + +// Write 将数据写入连接作为 HTTP 回复的一部分。 +func (gzw *gzipResponseWriter) Write(data []byte) (int, error) { + if !gzw.wroteHeader { + // 如果在 WriteHeader 之前调用 Write，根据 http.ResponseWriter 规范， + // 应写入 200 OK 头部。 + gzw.WriteHeader(http.StatusOK) + } + + if gzw.doCompression { + return gzw.gzWriter.Write(data) + } + return gzw.ResponseWriter.Write(data) +} + +// Close 确保 gzip writer 被关闭并释放资源。 +// 中间件应该在 c.Next() 之后调用它（通常在 defer 中）。 +func (gzw *gzipResponseWriter) Close() error { + if gzw.gzWriter != nil { + err := gzw.gzWriter.Close() // Close 会 Flush + putGzipWriterToPool(gzw.gzWriter, gzw.options.Level) + gzw.gzWriter = nil // 标记为已返回 + return err + } + return nil +} + +// Flush 将所有缓冲数据发送到客户端。 +// 实现 http.Flusher。 +func (gzw *gzipResponseWriter) Flush() { + if gzw.doCompression && gzw.gzWriter != nil { + _ = gzw.gzWriter.Flush() // 确保 gzip writer 的缓冲被刷新 + } + // 然后刷新底层的 writer (如果它支持) + if fl, ok := gzw.ResponseWriter.(http.Flusher); ok { + fl.Flush() + } +} + +// Hijack 允许调用者接管连接。 +// 实现 http.Hijacker。 +func (gzw *gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + // 如果正在压缩，hijack 的意义不大或不安全。 + // 然而，WriteHeader 应该会阻止对 101 状态码的压缩。 + // 此调用必须转到实际的底层 ResponseWriter。 + if hj, ok := gzw.ResponseWriter.(http.Hijacker); ok { + return hj.Hijack() + } + // 返回英文错误 + return nil, nil, errors.New("touka.gzipResponseWriter: underlying ResponseWriter does not implement http.Hijacker") +} + +// Status 返回已写入的 HTTP 状态码。委托给底层的 ResponseWriter。 +// 这确保了与 ecw 或其他可能跟踪状态的包装器的兼容性。 +func (gzw *gzipResponseWriter) Status() int { + if gzw.statusCode != 0 { // 如果我们在 WriteHeader 期间存储了它 + return gzw.statusCode + } + return gzw.ResponseWriter.Status() // 委托 +} + +// Size 返回已写入的字节数。委托给底层的 ResponseWriter。 +// 如果已压缩，这将是压缩后的大小 (由底层 writer 记录)。 +func (gzw *gzipResponseWriter) Size() int { + return gzw.ResponseWriter.Size() // GzipResponseWriter 本身不直接跟踪大小，依赖底层 +} + +// Written 返回 WriteHeader 是否已被调用。委托给底层的 ResponseWriter。 +func (gzw *gzipResponseWriter) Written() bool { + // 如果 gzw.wroteHeader 为 true，说明 WriteHeader 至少被 gzw 处理过。 + // 但最终是否写入底层取决于 gzw 的逻辑。 + // 更可靠的是询问底层 writer。 + return gzw.ResponseWriter.Written() // 委托 +} + +// --- Gzip 中间件 --- + +// Gzip 返回一个使用 Gzip 压缩 HTTP 响应的中间件。 +// 它会检查客户端的 "Accept-Encoding" 头部和响应的 "Content-Type" +// 来决定是否应用压缩。 +// level 参数指定压缩级别 (例如 gzip.DefaultCompression)。 +// opts 参数是可选的 GzipOptions。 +func Gzip(level int, opts ...GzipOptions) HandlerFunc { + config := GzipOptions{ // 初始化默认配置 + Level: level, + MinContentLength: 0, // 默认：无最小长度 + CompressibleTypes: defaultCompressibleTypes, + } + if len(opts) > 0 { // 如果传入了 GzipOptions，则覆盖默认值 + opt := opts[0] + config.Level = opt.Level // 允许通过结构体覆盖级别 + if opt.MinContentLength > 0 { + config.MinContentLength = opt.MinContentLength + } + if len(opt.CompressibleTypes) > 0 { + config.CompressibleTypes = opt.CompressibleTypes + } + } + // 验证级别 + if config.Level < gzip.DefaultCompression || config.Level > gzip.BestCompression { + config.Level = gzip.DefaultCompression + } + + return func(c *Context) { + // 1. 检查客户端是否接受 gzip + if !strings.Contains(c.Request.Header.Get(headerAcceptEncoding), encodingGzip) { + c.Next() + return + } + + // 2. 包装 ResponseWriter + originalWriter := c.Writer + gzw := acquireGzipResponseWriter(originalWriter, &config) + c.Writer = gzw // 替换上下文的 writer + + // defer 确保即使后续处理函数发生 panic，也能进行清理， + // 尽管恢复中间件应该自己处理 panic 响应。 + defer func() { + // 必须关闭 gzip writer 以刷新其缓冲区。 + // 这也会将 gzip.Writer 返回到其对象池。 + if err := gzw.Close(); err != nil { + // 记录关闭 gzip writer 时的错误，但不应覆盖已发送的响应 + // 通常这个错误不严重，因为数据可能已经大部分发送 + // 使用英文记录日志 + // log.Printf("Error closing gzip writer: %v", err) + c.AddError(err) // 可以选择将错误添加到 Context 中 + } + + // 恢复原始 writer 并将 gzipResponseWriter 返回到其对象池 + c.Writer = originalWriter + releaseGzipResponseWriter(gzw) + }() + + // 3. 调用链中的下一个处理函数 + c.Next() + + // c.Next() 执行完毕后，响应头部应该已经设置。 + // gzw.WriteHeader 会被显式调用或通过第一次 Write 隐式调用。 + // 如果 gzw.doCompression 为 true，响应体已写入 gzw.gzWriter。 + // defer 中的 gzw.Close() 会刷新最终的压缩字节。 + } +} 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 逻辑 + } +}