diff --git a/proxy/addon.go b/proxy/addon.go index d10e055..7ed029d 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(*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 { diff --git a/proxy/flow.go b/proxy/flow.go index 2bec3b6..2c09c87 100644 --- a/proxy/flow.go +++ b/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. diff --git a/proxy/proxy.go b/proxy/proxy.go index ab0f641..3a68425 100644 --- a/proxy/proxy.go +++ b/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) }