From a94d64bed1deb470b79d51861ee078a07573a2b6 Mon Sep 17 00:00:00 2001
From: "yu.deng" <dengyu326@gmail.com>
Date: Thu, 30 Jun 2022 14:58:26 +0800
Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E6=8C=81=20stream=20=E7=9A=84?=
 =?UTF-8?q?=E4=BF=AE=E6=94=B9?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 proxy/addon.go | 13 +++++++++++++
 proxy/proxy.go |  6 ++++++
 2 files changed, 19 insertions(+)

diff --git a/proxy/addon.go b/proxy/addon.go
index d10e055..4ca013a 100644
--- a/proxy/addon.go
+++ b/proxy/addon.go
@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"io"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -33,6 +34,12 @@ type Addon interface {
 
 	// The full HTTP response has been read.
 	Response(*Flow)
+
+	// Stream request body modifier
+	StreamRequestModifier(io.Reader) io.Reader
+
+	// Stream response body modifier
+	StreamResponseModifier(io.Reader) io.Reader
 }
 
 // BaseAddon do nothing
@@ -49,6 +56,12 @@ func (addon *BaseAddon) Requestheaders(*Flow)  {}
 func (addon *BaseAddon) Request(*Flow)         {}
 func (addon *BaseAddon) Responseheaders(*Flow) {}
 func (addon *BaseAddon) Response(*Flow)        {}
+func (addon *BaseAddon) StreamRequestModifier(in io.Reader) io.Reader {
+	return in
+}
+func (addon *BaseAddon) StreamResponseModifier(in io.Reader) io.Reader {
+	return in
+}
 
 // LogAddon log connection and flow
 type LogAddon struct {
diff --git a/proxy/proxy.go b/proxy/proxy.go
index ab0f641..88d87a4 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -189,6 +189,9 @@ func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 		}
 	}
 
+	for _, addon := range proxy.Addons {
+		reqBody = addon.StreamRequestModifier(reqBody)
+	}
 	proxyReq, err := http.NewRequest(f.Request.Method, f.Request.URL.String(), reqBody)
 	if err != nil {
 		log.Error(err)
@@ -247,6 +250,9 @@ func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 			}
 		}
 	}
+	for _, addon := range proxy.Addons {
+		resBody = addon.StreamResponseModifier(resBody)
+	}
 
 	reply(f.Response, resBody)
 }