From d254512f25ae89e18d76a711f622c89752dc4164 Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Fri, 18 Dec 2020 15:51:01 +0800
Subject: [PATCH] content-encoding

---
 README.md        |  1 +
 addon/dumper.go  |  9 ++++--
 flow/encoding.go | 73 ++++++++++++++++++++++++++++++++++++++++++++++++
 flow/flow.go     |  8 ++++++
 proxy/proxy.go   |  5 ++--
 5 files changed, 90 insertions(+), 6 deletions(-)
 create mode 100644 flow/encoding.go

diff --git a/README.md b/README.md
index 2c9f38a..aa542ba 100644
--- a/README.md
+++ b/README.md
@@ -57,6 +57,7 @@ Usage of mitmproxy:
 - [ ] http2
 - [ ] websocket 解析
 - [ ] web 界面
+- [ ] Content-Encoding 相关
 
 ## License
 
diff --git a/addon/dumper.go b/addon/dumper.go
index 636ae2f..3141a3f 100644
--- a/addon/dumper.go
+++ b/addon/dumper.go
@@ -68,9 +68,12 @@ func (d *Dumper) Requestheaders(f *flow.Flow) {
 			}
 			buf.WriteString("\r\n")
 
-			if d.level == 1 && f.Response.Body != nil && len(f.Response.Body) > 0 && CanPrint(f.Response.Body) {
-				buf.Write(f.Response.Body)
-				buf.WriteString("\r\n\r\n")
+			if d.level == 1 && f.Response.Body != nil && len(f.Response.Body) > 0 {
+				body, _ := f.Response.DecodedBody()
+				if len(body) > 0 && CanPrint(body) {
+					buf.Write(body)
+					buf.WriteString("\r\n\r\n")
+				}
 			}
 		}
 
diff --git a/flow/encoding.go b/flow/encoding.go
new file mode 100644
index 0000000..58c99f6
--- /dev/null
+++ b/flow/encoding.go
@@ -0,0 +1,73 @@
+package flow
+
+import (
+	"bytes"
+	"compress/gzip"
+	"errors"
+	"io"
+)
+
+// handle http header: content-encoding
+
+var EncodingNotSupport = errors.New("content-encoding not support")
+
+func (r *Response) DecodedBody() ([]byte, bool) {
+	if r.decodedBody != nil {
+		return r.decodedBody, r.decoded
+	}
+
+	if r.decodedErr != nil {
+		return nil, r.decoded
+	}
+
+	if r.Body == nil {
+		return nil, r.decoded
+	}
+
+	if len(r.Body) == 0 {
+		r.decodedBody = r.Body
+		return r.decodedBody, r.decoded
+	}
+
+	enc := r.Header.Get("Content-Encoding")
+	if enc == "" {
+		r.decodedBody = r.Body
+		return r.decodedBody, r.decoded
+	}
+
+	r.decodedBody, r.decodedErr = Decode(enc, r.Body)
+	if r.decodedErr != nil {
+		log.Error(r.decodedErr)
+	} else {
+		r.decoded = true
+	}
+
+	return r.decodedBody, r.decoded
+}
+
+// 当 Response.Body 替换为解压的内容时调用
+func (r *Response) RemoveEncodingHeader() {
+	r.Header.Del("Content-Encoding")
+	r.Header.Del("Content-Length")
+}
+
+func Decode(enc string, body []byte) ([]byte, error) {
+	if enc == "gzip" {
+		zr, err := gzip.NewReader(bytes.NewReader(body))
+		if err != nil {
+			return nil, err
+		}
+		buf := bytes.NewBuffer(make([]byte, 0))
+		_, err = io.Copy(buf, zr)
+		if err != nil {
+			return nil, err
+		}
+		err = zr.Close()
+		if err != nil {
+			return nil, err
+		}
+		return buf.Bytes(), nil
+	}
+
+	return nil, EncodingNotSupport
+}
diff --git a/flow/flow.go b/flow/flow.go
index a63521a..2c4cb7d 100644
--- a/flow/flow.go
+++ b/flow/flow.go
@@ -3,8 +3,12 @@ package flow
 import (
 	"net/http"
 	"net/url"
+
+	_log "github.com/sirupsen/logrus"
 )
 
+var log = _log.WithField("at", "flow")
+
 type Request struct {
 	Method string
 	URL    *url.URL
@@ -33,6 +37,10 @@ type Response struct {
 	StatusCode int
 	Header     http.Header
 	Body       []byte
+
+	decodedBody []byte
+	decoded     bool // decoded reports whether the response was sent compressed but was decoded to decodedBody.
+	decodedErr  error
 }
 
 type Flow struct {
diff --git a/proxy/proxy.go b/proxy/proxy.go
index 336968b..bd8c297 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -48,9 +48,8 @@ func NewProxy(opts *Options) (*Proxy, error) {
 			IdleConnTimeout:       90 * time.Second,
 			TLSHandshakeTimeout:   10 * time.Second,
 			ExpectContinueTimeout: 1 * time.Second,
-
-			ForceAttemptHTTP2:  false, // disable http2
-			DisableCompression: true,
+			ForceAttemptHTTP2:     false, // disable http2
+			DisableCompression:    true,  // To get the original response from the server, set Transport.DisableCompression to true.
 			TLSClientConfig: &tls.Config{
 				KeyLogWriter: GetTlsKeyLogWriter(),
 			},