58 lines
1.3 KiB
Go
58 lines
1.3 KiB
Go
package middleware
|
|
|
|
import (
|
|
"compress/gzip"
|
|
"github.com/gin-gonic/gin"
|
|
"net/http"
|
|
"strings"
|
|
)
|
|
|
|
func GzipMiddleware() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
// 1. 处理请求解压
|
|
if c.Request.Header.Get("Content-Encoding") == "gzip" {
|
|
gzReader, err := gzip.NewReader(c.Request.Body)
|
|
if err != nil {
|
|
c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "Invalid gzip format"})
|
|
return
|
|
}
|
|
defer gzReader.Close()
|
|
c.Request.Body = gzReader
|
|
}
|
|
|
|
// 2. 设置响应压缩支持
|
|
if strings.Contains(c.Request.Header.Get("Accept-Encoding"), "gzip") {
|
|
// 创建gzip writer
|
|
gzWriter := gzip.NewWriter(c.Writer)
|
|
defer gzWriter.Close()
|
|
|
|
// 替换原始writer为压缩writer
|
|
originalWriter := c.Writer
|
|
c.Writer = &gzipResponseWriter{
|
|
ResponseWriter: originalWriter,
|
|
gzWriter: gzWriter,
|
|
}
|
|
|
|
// 设置响应头
|
|
c.Header("Content-Encoding", "gzip")
|
|
c.Header("Vary", "Accept-Encoding")
|
|
}
|
|
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
// 自定义ResponseWriter实现gzip压缩
|
|
type gzipResponseWriter struct {
|
|
gin.ResponseWriter
|
|
gzWriter *gzip.Writer
|
|
}
|
|
|
|
func (w *gzipResponseWriter) Write(data []byte) (int, error) {
|
|
return w.gzWriter.Write(data)
|
|
}
|
|
|
|
func (w *gzipResponseWriter) WriteString(s string) (int, error) {
|
|
return w.gzWriter.Write([]byte(s))
|
|
}
|