From 84e61f776507cd2f385d1cb1bbdeddb609286055 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E7=A5=A5=E8=B6=85?= <iwind.liu@gmail.com>
Date: Mon, 15 Jan 2024 10:31:42 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/compressions/utils.go | 30 ++++++++++++++++++++++++++++++
 internal/nodes/http_writer.go  | 25 ++++++++++++-------------
 2 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/internal/compressions/utils.go b/internal/compressions/utils.go
index c466c91..74719c2 100644
--- a/internal/compressions/utils.go
+++ b/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
+}
diff --git a/internal/nodes/http_writer.go b/internal/nodes/http_writer.go
index 756214a..ac14905 100644
--- a/internal/nodes/http_writer.go
+++ b/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,18 +563,18 @@ func (this *HTTPWriter) PrepareWebP(resp *http.Response, size int64) {
 		}
 
 		var contentEncoding = this.GetHeader("Content-Encoding")
-		switch contentEncoding {
-		case "gzip", "deflate", "br", "zstd":
-			reader, err := compressions.NewReader(resp.Body, contentEncoding)
-			if err != nil {
+		if len(contentEncoding) > 0 {
+			if compressions.SupportEncoding(contentEncoding) {
+				reader, err := compressions.NewReader(resp.Body, contentEncoding)
+				if err != nil {
+					return
+				}
+				this.Header().Del("Content-Encoding")
+				this.Header().Del("Content-Length")
+				this.rawReader = reader
+			} else {
 				return
 			}
-			this.Header().Del("Content-Encoding")
-			this.Header().Del("Content-Length")
-			this.rawReader = reader
-		case "": // 空
-		default:
-			return
 		}
 
 		this.webpOriginContentType = contentType
@@ -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
 	}