From a565edd9268633147ebd2e85d92ba2d9198c93cc Mon Sep 17 00:00:00 2001
From: nyanpassu <nyanpassu@outlook.com>
Date: Wed, 1 Sep 2021 23:43:24 +0800
Subject: [PATCH] fix bugs when transfer chunked result

---
 proxy/proxy.go |  5 +++++
 utils/utils.go | 60 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/proxy/proxy.go b/proxy/proxy.go
index 7e62c0b..afe455d 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -66,6 +66,11 @@ func (ph HTTPProxyHandler) proxy(response http.ResponseWriter, request *http.Req
 		return
 	}
 	if resp.StatusCode != http.StatusSwitchingProtocols {
+		if request.Method == http.MethodGet && utils.IsChunkedEncoding(resp) {
+			log.Debug("[dispatch] Forward chunked response")
+			utils.ForwardChunked(response, resp)
+			return
+		}
 		log.Debug("[dispatch] Forward http response")
 		if err := utils.Forward(resp, response); err != nil {
 			log.Errorf("[dispatch] forward docker socket response failed %v", err)
diff --git a/utils/utils.go b/utils/utils.go
index 4e7f6f6..9c1764c 100644
--- a/utils/utils.go
+++ b/utils/utils.go
@@ -24,6 +24,66 @@ func Initialize(bufSize int) {
 	debug = log.GetLevel() == log.DebugLevel
 }
 
+// IsChunkedEncoding .
+func IsChunkedEncoding(src *http.Response) bool {
+	for _, value := range src.TransferEncoding {
+		if lower := strings.ToLower(strings.Trim(value, " ")); lower == "chunked" {
+			return true
+		}
+	}
+	return false
+}
+
+// ForwardChunked .
+func ForwardChunked(response http.ResponseWriter, resp *http.Response) {
+	log.Info("[ForwardChunked] Will forward chunked response")
+	initResponseHeader(response, resp.StatusCode, resp.Header)
+
+	buffer := make([]byte, 256)
+	for {
+		cnt, err := resp.Body.Read(buffer)
+		readed := buffer[:cnt]
+
+		log.Infof("%d bytes readed", cnt)
+
+		if _, err := response.Write(readed); err != nil {
+			log.WithError(err).Error("[ForwardChunked] server response write error")
+			return
+		}
+		if flusher, ok := response.(http.Flusher); ok {
+			flusher.Flush()
+		} else {
+			log.Warn("[ForwardChunked] server response is not http flusher")
+		}
+
+		if err == io.EOF {
+			log.Info("[ForwardChunked] client response end")
+			return
+		}
+
+		if err != nil {
+			log.WithError(err).Error("[ForwardChunked] copy io error")
+			return
+		}
+	}
+}
+
+func initResponseHeader(response http.ResponseWriter, statusCode int, header http.Header) {
+	PrintHeaders("ServerResponse", header)
+	responseHeader := response.Header()
+	for key, values := range header {
+		for _, value := range values {
+			responseHeader.Add(key, value)
+		}
+	}
+	response.WriteHeader(statusCode)
+	if flusher, ok := response.(http.Flusher); ok {
+		flusher.Flush()
+	} else {
+		log.Error("[WriteToServerResponse] Can't make flush to http.flusher")
+	}
+}
+
 // Forward .
 func Forward(src *http.Response, dst http.ResponseWriter) error {
 	copyHeader(src, dst)