Merge pull request #23 from pepesi/feature-another-stream-modifier

支持 stream 的修改

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(*Flow, io.Reader) io.Reader
+
+	// Stream response body modifier
+	StreamResponseModifier(*Flow, 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(f *Flow, in io.Reader) io.Reader {
+	return in
+}
+func (addon *BaseAddon) StreamResponseModifier(f *Flow, in io.Reader) io.Reader {
+	return in
+}
 
 // LogAddon log connection and flow
 type LogAddon struct {

proxy/flow.go

@@ -3,6 +3,7 @@ package proxy
 import (
 	"encoding/json"
 	"errors"
+	"io"
 	"net/http"
 	"net/url"
 
@@ -96,6 +97,7 @@ type Response struct {
 	StatusCode int         `json:"statusCode"`
 	Header     http.Header `json:"header"`
 	Body       []byte      `json:"-"`
+	BodyReader io.Reader
 
 	decodedBody []byte
 	decoded     bool // decoded reports whether the response was sent compressed but was decoded to decodedBody.

proxy/proxy.go

@@ -131,7 +131,14 @@ func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 			if err != nil {
 				logErr(log, err)
 			}
-		} else if response.Body != nil && len(response.Body) > 0 {
+		}
+		if response.BodyReader != nil {
+			_, err := io.Copy(res, response.BodyReader)
+			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)
@@ -189,6 +196,9 @@ func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 		}
 	}
 
+	for _, addon := range proxy.Addons {
+		reqBody = addon.StreamRequestModifier(f, reqBody)
+	}
 	proxyReq, err := http.NewRequest(f.Request.Method, f.Request.URL.String(), reqBody)
 	if err != nil {
 		log.Error(err)
@@ -247,6 +257,9 @@ func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 			}
 		}
 	}
+	for _, addon := range proxy.Addons {
+		resBody = addon.StreamResponseModifier(f, resBody)
+	}
 
 	reply(f.Response, resBody)
 }