128 lines
2.8 KiB
Go
128 lines
2.8 KiB
Go
package middleware
|
|
|
|
import (
|
|
"compress/gzip"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
)
|
|
|
|
// CompressionMiddleware 压缩中间件
|
|
type CompressionMiddleware struct {
|
|
// 压缩级别 (0-9)
|
|
level int
|
|
// 最小压缩大小
|
|
minSize int64
|
|
// 支持的内容类型
|
|
contentTypes []string
|
|
}
|
|
|
|
// NewCompressionMiddleware 创建压缩中间件
|
|
func NewCompressionMiddleware(level int, minSize int64) *CompressionMiddleware {
|
|
return &CompressionMiddleware{
|
|
level: level,
|
|
minSize: minSize,
|
|
contentTypes: []string{
|
|
"text/plain",
|
|
"text/html",
|
|
"text/css",
|
|
"text/javascript",
|
|
"application/javascript",
|
|
"application/json",
|
|
"application/xml",
|
|
"application/xml+rss",
|
|
"text/xml",
|
|
"application/x-yaml",
|
|
"text/yaml",
|
|
"application/x-www-form-urlencoded",
|
|
"application/x-protobuf",
|
|
"application/grpc",
|
|
"application/grpc+proto",
|
|
},
|
|
}
|
|
}
|
|
|
|
// Middleware 中间件处理函数
|
|
func (m *CompressionMiddleware) Middleware(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
// 处理请求压缩
|
|
if m.shouldCompress(r.Header.Get("Content-Encoding")) {
|
|
reader, err := gzip.NewReader(r.Body)
|
|
if err != nil {
|
|
http.Error(w, "Invalid gzip content", http.StatusBadRequest)
|
|
return
|
|
}
|
|
defer reader.Close()
|
|
r.Body = io.NopCloser(reader)
|
|
}
|
|
|
|
// 处理响应压缩
|
|
if m.shouldCompressResponse(r) {
|
|
gw := gzip.NewWriter(w)
|
|
defer gw.Close()
|
|
|
|
// 包装响应写入器
|
|
writer := &gzipResponseWriter{
|
|
ResponseWriter: w,
|
|
writer: gw,
|
|
}
|
|
|
|
// 设置响应头
|
|
w.Header().Set("Content-Encoding", "gzip")
|
|
w.Header().Set("Vary", "Accept-Encoding")
|
|
|
|
next.ServeHTTP(writer, r)
|
|
} else {
|
|
next.ServeHTTP(w, r)
|
|
}
|
|
})
|
|
}
|
|
|
|
// shouldCompress 检查是否应该压缩请求
|
|
func (m *CompressionMiddleware) shouldCompress(encoding string) bool {
|
|
return strings.Contains(encoding, "gzip")
|
|
}
|
|
|
|
// shouldCompressResponse 检查是否应该压缩响应
|
|
func (m *CompressionMiddleware) shouldCompressResponse(r *http.Request) bool {
|
|
// 检查客户端是否支持gzip
|
|
acceptEncoding := r.Header.Get("Accept-Encoding")
|
|
if !strings.Contains(acceptEncoding, "gzip") {
|
|
return false
|
|
}
|
|
|
|
// 检查内容类型
|
|
contentType := r.Header.Get("Content-Type")
|
|
for _, t := range m.contentTypes {
|
|
if strings.Contains(contentType, t) {
|
|
return true
|
|
}
|
|
}
|
|
|
|
return false
|
|
}
|
|
|
|
// gzipResponseWriter 包装的响应写入器
|
|
type gzipResponseWriter struct {
|
|
http.ResponseWriter
|
|
writer *gzip.Writer
|
|
}
|
|
|
|
// Write 写入数据
|
|
func (w *gzipResponseWriter) Write(b []byte) (int, error) {
|
|
return w.writer.Write(b)
|
|
}
|
|
|
|
// WriteHeader 写入状态码
|
|
func (w *gzipResponseWriter) WriteHeader(statusCode int) {
|
|
w.ResponseWriter.WriteHeader(statusCode)
|
|
}
|
|
|
|
// Flush 刷新数据
|
|
func (w *gzipResponseWriter) Flush() {
|
|
w.writer.Flush()
|
|
if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
|
|
flusher.Flush()
|
|
}
|
|
}
|