优化代码

internal/compressions/utils.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
 	"io"
+	"net/http"
 )
 
 type ContentEncoding = string
@@ -58,3 +59,32 @@ func NewWriter(writer io.Writer, compressType serverconfigs.HTTPCompressionType,
 	}
 	return nil, errors.New("invalid compression type '" + compressType + "'")
 }
+
+// SupportEncoding 检查是否支持某个编码
+func SupportEncoding(encoding string) bool {
+	return encoding == ContentEncodingBr ||
+		encoding == ContentEncodingGzip ||
+		encoding == ContentEncodingDeflate ||
+		encoding == ContentEncodingZSTD
+}
+
+// WrapHTTPResponse 包装http.Response对象
+func WrapHTTPResponse(resp *http.Response) {
+	if resp == nil {
+		return
+	}
+
+	var contentEncoding = resp.Header.Get("Content-Encoding")
+	if len(contentEncoding) == 0 || !SupportEncoding(contentEncoding) {
+		return
+	}
+
+	reader, err := NewReader(resp.Body, contentEncoding)
+	if err != nil {
+		// unable to decode, we ignore the error
+		return
+	}
+	resp.Header.Del("Content-Encoding")
+	resp.Header.Del("Content-Length")
+	resp.Body = reader
+}

internal/nodes/http_writer.go

@@ -18,7 +18,6 @@ import (
 	setutils "github.com/TeaOSLab/EdgeNode/internal/utils/sets"
 	"github.com/TeaOSLab/EdgeNode/internal/utils/writers"
 	_ "github.com/biessek/golang-ico"
-	"github.com/iwind/TeaGo/lists"
 	"github.com/iwind/TeaGo/types"
 	"github.com/iwind/gowebp"
 	_ "golang.org/x/image/bmp"
@@ -564,8 +563,8 @@ func (this *HTTPWriter) PrepareWebP(resp *http.Response, size int64) {
 		}
 
 		var contentEncoding = this.GetHeader("Content-Encoding")
-		switch contentEncoding {
+		if len(contentEncoding) > 0 {
-		case "gzip", "deflate", "br", "zstd":
+			if compressions.SupportEncoding(contentEncoding) {
 				reader, err := compressions.NewReader(resp.Body, contentEncoding)
 				if err != nil {
 					return
@@ -573,10 +572,10 @@ func (this *HTTPWriter) PrepareWebP(resp *http.Response, size int64) {
 				this.Header().Del("Content-Encoding")
 				this.Header().Del("Content-Length")
 				this.rawReader = reader
-		case "": // 空
+			} else {
-		default:
 				return
 			}
+		}
 
 		this.webpOriginContentType = contentType
 		this.webpIsEncoding = true
@@ -603,7 +602,7 @@ func (this *HTTPWriter) PrepareCompression(resp *http.Response, size int64) {
 	var contentEncoding = this.GetHeader("Content-Encoding")
 
 	if this.compressionConfig == nil || !this.compressionConfig.IsOn {
-		if lists.ContainsString([]string{"gzip", "deflate", "br", "zstd"}, contentEncoding) && !httpAcceptEncoding(acceptEncodings, contentEncoding) {
+		if compressions.SupportEncoding(contentEncoding) && !httpAcceptEncoding(acceptEncodings, contentEncoding) {
 			reader, err := compressions.NewReader(resp.Body, contentEncoding)
 			if err != nil {
 				return
@@ -625,7 +624,7 @@ func (this *HTTPWriter) PrepareCompression(resp *http.Response, size int64) {
 	}
 
 	// 如果已经有编码则不处理
-	if len(contentEncoding) > 0 && (!this.compressionConfig.DecompressData || !lists.ContainsString([]string{"gzip", "deflate", "br", "zstd"}, contentEncoding)) {
+	if len(contentEncoding) > 0 && (!this.compressionConfig.DecompressData || !compressions.SupportEncoding(contentEncoding)) {
 		return
 	}