0904

internal/route/gzip.go

@@ -0,0 +1,150 @@
+// Copyright 2026 tamakyi. All rights reserved.
+// Use of this source code is governed by a MIT-style
+// license that can be found in the LICENSE file.
+
+package route
+
+import (
+	"bufio"
+	"compress/gzip"
+	"errors"
+	"net"
+	"net/http"
+	"strings"
+)
+
+// GzipHandler 返回一个压缩中间件，用于包装 http.Handler。
+// 它对可压缩的响应（HTML / CSS / JS / JSON 等文本类型）做 gzip 压缩，
+// 自动跳过图片、视频、字体等已经压缩过的格式。
+//
+// 由于 flamego 的 Context 无法替换底层 ResponseWriter，这里采用在
+// http.Server 最外层包装 Handler 的方式，保证对所有响应（含静态文件）生效。
+func GzipHandler(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// 客户端不支持 gzip 时直接透传
+		if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		gw := &gzipResponseWriter{
+			ResponseWriter: w,
+			gz:             gzip.NewWriter(w),
+		}
+		defer func() {
+			// 只有真正写入过 body 才关闭 gzip 流，避免 302/204 等
+			// 无 body 响应被追加空的 gzip 尾块
+			if gw.compress && gw.wrote {
+				_ = gw.gz.Close()
+			}
+		}()
+		next.ServeHTTP(gw, r)
+	})
+}
+
+// gzipResponseWriter 包装 http.ResponseWriter，在 WriteHeader 时根据
+// Content-Type 决定是否启用 gzip，保持接口兼容（Flusher / Hijacker / Pusher）。
+type gzipResponseWriter struct {
+	http.ResponseWriter
+	gz          *gzip.Writer
+	compress    bool
+	wrote       bool
+	wroteHeader bool
+}
+
+// WriteHeader 记录状态码，并根据 Content-Type 决定是否压缩。
+// 由于 Content-Encoding 改变后 Content-Length 不再准确，需要删除该头。
+func (w *gzipResponseWriter) WriteHeader(status int) {
+	if w.wroteHeader {
+		return
+	}
+	w.wroteHeader = true
+
+	// 无实体响应（204 No Content / 304 Not Modified）不应附加 Content-Encoding。
+	if status == http.StatusNoContent || status == http.StatusNotModified {
+		w.ResponseWriter.WriteHeader(status)
+		return
+	}
+
+	if shouldCompressContentType(w.Header().Get("Content-Type")) {
+		w.compress = true
+		w.Header().Del("Content-Length")
+		w.Header().Set("Content-Encoding", "gzip")
+		w.Header().Add("Vary", "Accept-Encoding")
+	}
+	w.ResponseWriter.WriteHeader(status)
+}
+
+// Write 写入响应体。若尚未 WriteHeader（handler 依赖 net/http 自动嗅探
+// Content-Type 的情况），先用首个字节块探测类型再决策。
+func (w *gzipResponseWriter) Write(b []byte) (int, error) {
+	if !w.wroteHeader {
+		ct := w.Header().Get("Content-Type")
+		if ct == "" {
+			ct = http.DetectContentType(b)
+			// 将嗅探到的类型写回响应头，使 WriteHeader 的压缩判定能正确生效
+			// （否则空 Content-Type 会被判定为"不压缩"，导致本应压缩的纯文本响应未压缩）。
+			w.Header().Set("Content-Type", ct)
+		}
+		w.WriteHeader(http.StatusOK)
+	}
+	if w.compress {
+		w.wrote = true
+		return w.gz.Write(b)
+	}
+	return w.ResponseWriter.Write(b)
+}
+
+// Flush 转发底层 Flusher，并在压缩模式下先刷新 gzip 缓冲。
+func (w *gzipResponseWriter) Flush() {
+	if w.compress && w.gz != nil {
+		_ = w.gz.Flush()
+	}
+	if f, ok := w.ResponseWriter.(http.Flusher); ok {
+		f.Flush()
+	}
+}
+
+// Hijack 转发底层 Hijacker，保持对 WebSocket 等升级协议兼容。
+func (w *gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if h, ok := w.ResponseWriter.(http.Hijacker); ok {
+		return h.Hijack()
+	}
+	return nil, nil, errors.New("gzipResponseWriter: hijacking not supported")
+}
+
+// Push 转发 HTTP/2 Server Push（如果底层支持）。
+func (w *gzipResponseWriter) Push(target string, opts *http.PushOptions) error {
+	if p, ok := w.ResponseWriter.(http.Pusher); ok {
+		return p.Push(target, opts)
+	}
+	return http.ErrNotSupported
+}
+
+// shouldCompressContentType 判断响应 Content-Type 是否值得压缩。
+// 空类型、二进制图片/视频/音频/字体等已压缩格式直接跳过。
+func shouldCompressContentType(ct string) bool {
+	ct = strings.ToLower(strings.TrimSpace(ct))
+	if ct == "" {
+		return false
+	}
+
+	// SVG 是文本格式，值得压缩；其余图片格式已是压缩格式
+	if strings.HasPrefix(ct, "image/") {
+		return strings.Contains(ct, "svg")
+	}
+	// 视频 / 音频 / 字体已是压缩格式
+	if strings.HasPrefix(ct, "video/") || strings.HasPrefix(ct, "audio/") ||
+		strings.HasPrefix(ct, "font/") {
+		return false
+	}
+	// 已知压缩或二进制格式
+	switch ct {
+	case "application/zip", "application/gzip", "application/x-7z-compressed",
+		"application/x-rar-compressed", "application/pdf",
+		"application/vnd.ms-fontobject", "application/wasm",
+		"application/octet-stream", "application/x-protobuf", "application/grpc":
+		return false
+	}
+	return true
+}