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() 会刷新最终的压缩字节。 + } +}