From 9080c0e1160f08a5636e1d9bd3bb4ca6d0e05a18 Mon Sep 17 00:00:00 2001
From: "yu.deng" <dengyu326@gmail.com>
Date: Thu, 30 Jun 2022 15:54:55 +0800
Subject: [PATCH] need flow context in stream modifier provide extra info

---
 proxy/addon.go |  8 ++++----
 proxy/proxy.go | 24 +++++++++++-------------
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/proxy/addon.go b/proxy/addon.go
index 4ca013a..7ed029d 100644
--- a/proxy/addon.go
+++ b/proxy/addon.go
@@ -36,10 +36,10 @@ type Addon interface {
 	Response(*Flow)
 
 	// Stream request body modifier
-	StreamRequestModifier(io.Reader) io.Reader
+	StreamRequestModifier(*Flow, io.Reader) io.Reader
 
 	// Stream response body modifier
-	StreamResponseModifier(io.Reader) io.Reader
+	StreamResponseModifier(*Flow, io.Reader) io.Reader
 }
 
 // BaseAddon do nothing
@@ -56,10 +56,10 @@ 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 {
+func (addon *BaseAddon) StreamRequestModifier(f *Flow, in io.Reader) io.Reader {
 	return in
 }
-func (addon *BaseAddon) StreamResponseModifier(in io.Reader) io.Reader {
+func (addon *BaseAddon) StreamResponseModifier(f *Flow, in io.Reader) io.Reader {
 	return in
 }
 
diff --git a/proxy/proxy.go b/proxy/proxy.go
index 8776329..3a68425 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -132,18 +132,16 @@ func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 				logErr(log, err)
 			}
 		}
-		if response.Body != nil {
-			if response.BodyReader != nil {
-				_, err := io.Copy(res, response.BodyReader)
-				if err != nil {
-					logErr(log, err)
-				}
+		if response.BodyReader != nil {
+			_, err := io.Copy(res, response.BodyReader)
+			if err != nil {
+				logErr(log, err)
 			}
-			if len(response.Body) > 0 {
-				_, err := res.Write(response.Body)
-				if err != nil {
-					logErr(log, err)
-				}
+		}
+		if response.Body != nil && len(response.Body) > 0 {
+			_, err := res.Write(response.Body)
+			if err != nil {
+				logErr(log, err)
 			}
 		}
 	}
@@ -199,7 +197,7 @@ func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 	}
 
 	for _, addon := range proxy.Addons {
-		reqBody = addon.StreamRequestModifier(reqBody)
+		reqBody = addon.StreamRequestModifier(f, reqBody)
 	}
 	proxyReq, err := http.NewRequest(f.Request.Method, f.Request.URL.String(), reqBody)
 	if err != nil {
@@ -260,7 +258,7 @@ func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 		}
 	}
 	for _, addon := range proxy.Addons {
-		resBody = addon.StreamResponseModifier(resBody)
+		resBody = addon.StreamResponseModifier(f, resBody)
 	}
 
 	reply(f.Response, resBody)