tamabox/internal/route/gzip.go

// 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
}