143 lines
4.2 KiB
Go
143 lines
4.2 KiB
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 对可压缩的文本响应（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) {
|
||
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 才关闭，避免 302/204 等无 body 响应被追加空的 gzip 尾块
|
||
if gw.compress && gw.wrote {
|
||
_ = gw.gz.Close()
|
||
}
|
||
}()
|
||
next.ServeHTTP(gw, r)
|
||
})
|
||
}
|
||
|
||
// gzipResponseWriter 在 WriteHeader 时按 Content-Type 决定是否启用 gzip，
|
||
// 并保持 Flusher / Hijacker / Pusher 接口兼容。
|
||
type gzipResponseWriter struct {
|
||
http.ResponseWriter
|
||
gz *gzip.Writer
|
||
compress bool
|
||
wrote bool
|
||
wroteHeader bool
|
||
}
|
||
|
||
// WriteHeader 记录状态码并决定是否压缩；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，先用首个字节块探测类型再决策。
|
||
func (w *gzipResponseWriter) Write(b []byte) (int, error) {
|
||
if !w.wroteHeader {
|
||
ct := w.Header().Get("Content-Type")
|
||
if ct == "" {
|
||
ct = http.DetectContentType(b)
|
||
// 嗅探结果须写回头，否则空 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
|
||
}
|
||
|
||
if strings.HasPrefix(ct, "image/") {
|
||
return strings.Contains(ct, "svg") // 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
|
||
}
|