Determine buffering if req or beresp json_body response occurs in hcl…

….bodies #44

eval/buffer.go

@@ -0,0 +1,52 @@
+package eval
+
+import (
+	"reflect"
+
+	"github.com/hashicorp/hcl/v2"
+)
+
+type BufferOption uint8
+
+const (
+	BufferNone BufferOption = iota
+	BufferRequest
+	BufferResponse
+)
+
+// MustBuffer determines if any of the hcl.bodies makes use of 'post' or 'json_body'.
+func MustBuffer(ctxBodies []hcl.Body) BufferOption {
+	result := BufferNone
+	for _, body := range ctxBodies {
+		attrs, err := body.JustAttributes()
+		if err != nil {
+			return result
+		}
+		for _, attr := range attrs {
+			for _, traversal := range attr.Expr.Variables() {
+				rootName := traversal.RootName()
+				if rootName != "req" && rootName != "beresp" {
+					continue
+				}
+				for _, step := range traversal[1:] {
+					nameField := reflect.ValueOf(step).FieldByName("Name")
+					name := nameField.String()
+					switch name {
+					case "json_body":
+						switch rootName {
+						case "req":
+							result |= BufferRequest
+						case "beresp":
+							result |= BufferResponse
+						}
+					case "post":
+						if rootName == "req" {
+							result |= BufferRequest
+						}
+					}
+				}
+			}
+		}
+	}
+	return result
+}

eval/context.go

@@ -49,7 +49,7 @@ func NewENVContext(src []byte) *hcl.EvalContext {
 	}
 }
 
-func NewHTTPContext(baseCtx *hcl.EvalContext, buffer bool, req, bereq *http.Request, beresp *http.Response) *hcl.EvalContext {
+func NewHTTPContext(baseCtx *hcl.EvalContext, bufOpt BufferOption, req, bereq *http.Request, beresp *http.Response) *hcl.EvalContext {
 	if req == nil {
 		return baseCtx
 	}
@@ -66,7 +66,7 @@ func NewHTTPContext(baseCtx *hcl.EvalContext, buffer bool, req, bereq *http.Requ
 		id = uid
 	}
 
-	if buffer {
+	if (bufOpt & BufferRequest) == BufferRequest {
 		switch req.Method {
 		case http.MethodPost, http.MethodPut, http.MethodPatch:
 			setGetBody(req)
@@ -80,7 +80,7 @@ func NewHTTPContext(baseCtx *hcl.EvalContext, buffer bool, req, bereq *http.Requ
 		"url":       cty.StringVal(newRawURL(req.URL).String()),
 		"query":     seetie.ValuesMapToValue(req.URL.Query()),
 		"post":      seetie.ValuesMapToValue(parseForm(req).PostForm),
-		"json_body": seetie.MapToValue(parseJSON(req)),
+		"json_body": seetie.MapToValue(parseReqJSON(req)),
 	}.Merge(newVariable(httpCtx, req.Cookies(), req.Header))))
 
 	if beresp != nil {
@@ -91,8 +91,14 @@ func NewHTTPContext(baseCtx *hcl.EvalContext, buffer bool, req, bereq *http.Requ
 			"query":  seetie.ValuesMapToValue(bereq.URL.Query()),
 			"post":   seetie.ValuesMapToValue(parseForm(bereq).PostForm),
 		}.Merge(newVariable(httpCtx, bereq.Cookies(), bereq.Header)))
+
+		var jsonBody map[string]interface{}
+		if (bufOpt & BufferResponse) == BufferResponse {
+			jsonBody = parseRespJSON(beresp)
+		}
 		evalCtx.Variables["beresp"] = cty.ObjectVal(ContextMap{
-			"status": cty.StringVal(strconv.Itoa(beresp.StatusCode)),
+			"status":    cty.StringVal(strconv.Itoa(beresp.StatusCode)),
+			"json_body": seetie.MapToValue(jsonBody),
 		}.Merge(newVariable(httpCtx, beresp.Cookies(), beresp.Header)))
 	}
 
@@ -144,26 +150,50 @@ func parseForm(r *http.Request) *http.Request {
 	return r
 }
 
-func parseJSON(r *http.Request) map[string]interface{} {
-	if r.GetBody == nil {
-		return nil
-	}
+func isJSONMediaType(contentType string) bool {
+	m, _, _ := mime.ParseMediaType(contentType)
+	return m == "application/json"
+}
 
-	m, _, _ := mime.ParseMediaType(r.Header.Get("Content-Type"))
-	if m != "application/json" {
+func parseJSON(r io.Reader) map[string]interface{} {
+	if r == nil {
 		return nil
 	}
-
 	var result map[string]interface{}
-	b, err := ioutil.ReadAll(r.Body)
+	b, err := ioutil.ReadAll(r)
 	if err != nil {
 		return nil
 	}
 	_ = json.Unmarshal(b, &result)
-	r.Body, _ = r.GetBody() // reset
 	return result
 }
 
+func parseReqJSON(req *http.Request) map[string]interface{} {
+	if req.GetBody == nil {
+		return nil
+	}
+
+	if !isJSONMediaType(req.Header.Get("Content-Type")) {
+		return nil
+	}
+
+	result := parseJSON(req.Body)
+	req.Body, _ = req.GetBody() // reset
+	return result
+}
+
+func parseRespJSON(beresp *http.Response) map[string]interface{} {
+	if !isJSONMediaType(beresp.Header.Get("Content-Type")) {
+		return nil
+	}
+
+	buf := &bytes.Buffer{}
+	io.Copy(buf, beresp.Body) // TODO: err handling
+	// reset
+	beresp.Body = newReadCloser(bytes.NewBuffer(buf.Bytes()), beresp.Body)
+	return parseJSON(buf)
+}
+
 func newRawURL(u *url.URL) *url.URL {
 	rawURL := *u
 	rawURL.RawQuery = ""

handler/proxy.go

@@ -10,7 +10,6 @@ import (
 	"net"
 	"net/http"
 	"net/url"
-	"reflect"
 	"strconv"
 	"strings"
 	"time"
@@ -41,13 +40,13 @@ var (
 )
 
 type Proxy struct {
-	evalContext *hcl.EvalContext
-	log         *logrus.Entry
-	mustBuffer  bool
-	options     *ProxyOptions
-	originURL   *url.URL
-	transport   *http.Transport
-	upstreamLog *logging.AccessLog
+	evalContext  *hcl.EvalContext
+	log          *logrus.Entry
+	bufferOption eval.BufferOption
+	options      *ProxyOptions
+	originURL    *url.URL
+	transport    *http.Transport
+	upstreamLog  *logging.AccessLog
 }
 
 type ProxyOptions struct {
@@ -117,12 +116,12 @@ func NewProxy(options *ProxyOptions, log *logrus.Entry, evalCtx *hcl.EvalContext
 	env.DecodeWithPrefix(&logConf, "BACKEND_")
 
 	proxy := &Proxy{
-		evalContext: evalCtx,
-		log:         log,
-		mustBuffer:  mustBuffer(options),
-		options:     options,
-		originURL:   originURL,
-		upstreamLog: logging.NewAccessLog(&logConf, log.Logger),
+		bufferOption: eval.MustBuffer(options.Context),
+		evalContext:  evalCtx,
+		log:          log,
+		options:      options,
+		originURL:    originURL,
+		upstreamLog:  logging.NewAccessLog(&logConf, log.Logger),
 	}
 
 	var tlsConf *tls.Config
@@ -315,7 +314,7 @@ func (p *Proxy) setRoundtripContext(req *http.Request, beresp *http.Response) {
 		headerCtx = req.Header
 	}
 
-	evalCtx := eval.NewHTTPContext(p.evalContext, p.mustBuffer, req, bereq, beresp)
+	evalCtx := eval.NewHTTPContext(p.evalContext, p.bufferOption, req, bereq, beresp)
 
 	// Remove blacklisted headers after evaluation to be accessible within our context configuration.
 	if attrCtx == attrReqHeaders {
@@ -337,32 +336,6 @@ func (p *Proxy) setRoundtripContext(req *http.Request, beresp *http.Response) {
 	}
 }
 
-// mustBuffer determines if any of the hcl.bodies makes use of 'post' or 'json_body'.
-func mustBuffer(opts *ProxyOptions) bool {
-	for _, body := range opts.Context {
-		attrs, err := body.JustAttributes()
-		if err != nil {
-			return false
-		}
-		for _, attr := range attrs {
-			for _, traversal := range attr.Expr.Variables() {
-				if traversal.RootName() != "req" {
-					continue
-				}
-				for _, step := range traversal[1:] {
-					nameField := reflect.ValueOf(step).FieldByName("Name")
-					name := nameField.String()
-					switch name {
-					case "json_body", "post":
-						return true
-					}
-				}
-			}
-		}
-	}
-	return false
-}
-
 func isCorsRequest(req *http.Request) bool {
 	return req.Header.Get("Origin") != ""
 }