Merge pull request #69 from SpectoLabs/develop

Develop

cache.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"fmt"
 
 	log "github.com/Sirupsen/logrus"
@@ -48,7 +49,12 @@ func (c *Cache) Get(key []byte) (value []byte, err error) {
 		if bucket == nil {
 			return fmt.Errorf("Bucket %q not found!", c.requestsBucket)
 		}
-		value = bucket.Get(key)
+
+		// "Byte slices returned from Bolt are only valid during a transaction."
+		var buffer bytes.Buffer
+		buffer.Write(bucket.Get(key))
+
+		value = buffer.Bytes()
 		return nil
 	})
 

examples/middleware/reflect_body/reflect_body.py

@@ -1,7 +1,9 @@
 #!/usr/bin/env python
 import sys
 import json
+import logging
 
+logging.basicConfig(filename='middleware_reflect.log', level=logging.DEBUG)
 
 def main():
     """
@@ -16,6 +18,8 @@ def main():
     payload_dict['response']['body'] = payload_dict['request']['body']
     payload_dict['response']['status'] = 200
 
+
+    logging.debug(payload_dict)
     # returning new payload
     print(json.dumps(payload_dict))
 

middleware.go

@@ -98,6 +98,7 @@ func ExecuteMiddleware(command string, payload Payload) (Payload, error) {
 		if err != nil {
 			log.WithFields(log.Fields{
 				"mwOutput": string(mwOutput),
+				"error":    err.Error(),
 			}).Error("Failed to unmarshal JSON from middleware")
 		} else {
 			// payload unmarshalled into Payload struct, returning it

middleware_test.go

@@ -47,3 +47,18 @@ func TestMakeCustom404(t *testing.T) {
 	expect(t, newPayload.Response.Status, 404)
 	expect(t, newPayload.Response.Headers["middleware"][0], "changed response")
 }
+
+func TestReflectBody(t *testing.T) {
+	command := "./examples/middleware/reflect_body/reflect_body.py"
+
+	req := requestDetails{Path: "/", Method: "GET", Destination: "hostname-x", Query: "", Body: "request_body_here"}
+
+	payload := Payload{Request: req}
+
+	newPayload, err := ExecuteMiddleware(command, payload)
+
+	expect(t, err, nil)
+	expect(t, newPayload.Response.Body, req.Body)
+	expect(t, newPayload.Request.Method, req.Method)
+	expect(t, newPayload.Request.Destination, req.Destination)
+}

models.go

@@ -39,10 +39,22 @@ type requestDetails struct {
 	Headers     map[string][]string `json:"headers"`
 }
 
+func (r *request) concatenate() string {
+	var buffer bytes.Buffer
+
+	buffer.WriteString(r.details.Destination)
+	buffer.WriteString(r.details.Path)
+	buffer.WriteString(r.details.Method)
+	buffer.WriteString(r.details.Query)
+	buffer.WriteString(r.details.Body)
+
+	return buffer.String()
+}
+
 // hash returns unique hash key for request
 func (r *request) hash() string {
 	h := md5.New()
-	io.WriteString(h, fmt.Sprintf("%s%s%s%s", r.details.Destination, r.details.Path, r.details.Method, r.details.Query))
+	io.WriteString(h, r.concatenate())
 	return fmt.Sprintf("%x", h.Sum(nil))
 }
 
@@ -120,7 +132,7 @@ func (d *DBClient) captureRequest(req *http.Request) (*http.Response, error) {
 		}
 
 		// saving response body with request/response meta to cache
-		go d.save(req, resp, respBody, reqBody)
+		d.save(req, reqBody, resp, respBody)
 	}
 
 	// return new response or error here
@@ -197,9 +209,9 @@ func (d *DBClient) doRequest(request *http.Request) (*http.Response, error) {
 }
 
 // save gets request fingerprint, extracts request body, status code and headers, then saves it to cache
-func (d *DBClient) save(req *http.Request, resp *http.Response, respBody []byte, reqBody []byte) {
+func (d *DBClient) save(req *http.Request, reqBody []byte, resp *http.Response, respBody []byte) {
 	// record request here
-	key := getRequestFingerprint(req)
+	key := getRequestFingerprint(req, reqBody)
 
 	if resp == nil {
 		resp = emptyResp
@@ -248,17 +260,31 @@ func (d *DBClient) save(req *http.Request, resp *http.Response, respBody []byte,
 }
 
 // getRequestFingerprint returns request hash
-func getRequestFingerprint(req *http.Request) string {
-	details := requestDetails{Path: req.URL.Path, Method: req.Method, Destination: req.Host, Query: req.URL.RawQuery}
+func getRequestFingerprint(req *http.Request, requestBody []byte) string {
+	details := requestDetails{
+		Path:        req.URL.Path,
+		Method:      req.Method,
+		Destination: req.Host,
+		Query:       req.URL.RawQuery,
+		Body:        string(requestBody),
+	}
+
 	r := request{details: details}
 	return r.hash()
 }
 
 // getResponse returns stored response from cache
 func (d *DBClient) getResponse(req *http.Request) *http.Response {
 
-	key := getRequestFingerprint(req)
-	//	var payload Payload
+	reqBody, err := ioutil.ReadAll(req.Body)
+
+	if err != nil {
+		log.WithFields(log.Fields{
+			"error": err.Error(),
+		}).Error("Got error when reading request body")
+	}
+
+	key := getRequestFingerprint(req, reqBody)
 
 	payloadBts, err := d.cache.Get([]byte(key))
 

models_test.go

@@ -7,7 +7,6 @@ import (
 	"net/http"
 	"os"
 	"testing"
-	"time"
 )
 
 func TestMain(m *testing.M) {
@@ -43,25 +42,68 @@ func TestRequestBodyCaptured(t *testing.T) {
 	server, dbClient := testTools(200, `{'message': 'here'}`)
 	defer server.Close()
 
-	body := ioutil.NopCloser(bytes.NewBuffer([]byte("fizz=buzz")))
+	requestBody := []byte("fizz=buzz")
+
+	body := ioutil.NopCloser(bytes.NewBuffer(requestBody))
 
 	req, err := http.NewRequest("POST", "http://capture_body.com", body)
 	expect(t, err, nil)
 
 	_, err = dbClient.captureRequest(req)
 	expect(t, err, nil)
 
-	// since capture
-	time.Sleep(10 * time.Millisecond)
-
-	fp := getRequestFingerprint(req)
+	fp := getRequestFingerprint(req, requestBody)
 
 	payloadBts, err := dbClient.cache.Get([]byte(fp))
 	expect(t, err, nil)
 
 	payload, err := decodePayload(payloadBts)
 	expect(t, err, nil)
 	expect(t, payload.Request.Body, "fizz=buzz")
+}
+
+func TestMatchOnRequestBody(t *testing.T) {
+
+	server, dbClient := testTools(200, `{'message': 'here'}`)
+	defer server.Close()
+
+	// preparing and saving requests/responses with unique bodies
+	for i := 0; i < 5; i++ {
+		requestBody := []byte(fmt.Sprintf("fizz=buzz, number=%d", i))
+		body := ioutil.NopCloser(bytes.NewBuffer(requestBody))
+
+		request, err := http.NewRequest("POST", "http://capture_body.com", body)
+		expect(t, err, nil)
+
+		resp := response{
+			Status: 200,
+			Body:   fmt.Sprintf("body here, number=%d", i),
+		}
+		payload := Payload{Response: resp}
+
+		// creating response
+		c := NewConstructor(request, payload)
+		response := c.reconstructResponse()
+
+		dbClient.save(request, requestBody, response, []byte(resp.Body))
+	}
+
+	// now getting responses
+	for i := 0; i < 5; i++ {
+		requestBody := []byte(fmt.Sprintf("fizz=buzz, number=%d", i))
+		body := ioutil.NopCloser(bytes.NewBuffer(requestBody))
+
+		request, _ := http.NewRequest("POST", "http://capture_body.com", body)
+
+		response := dbClient.getResponse(request)
+
+		responseBody, err := ioutil.ReadAll(response.Body)
+		response.Body.Close()
+
+		expect(t, err, nil)
+		expect(t, string(responseBody), fmt.Sprintf("body here, number=%d", i))
+
+	}
 
 }
 
@@ -71,7 +113,7 @@ func TestRequestFingerprint(t *testing.T) {
 	req, err := http.NewRequest("GET", "http://example.com", nil)
 	expect(t, err, nil)
 
-	fp := getRequestFingerprint(req)
+	fp := getRequestFingerprint(req, []byte(""))
 
 	expect(t, fp, "92a65ed4ca2b7100037a4cba9afd15ea")
 

synthesize.go

@@ -1,17 +1,33 @@
 package main
 
 import (
-	"bufio"
+	"bytes"
 	"fmt"
+	"io/ioutil"
 	"net/http"
+
+	log "github.com/Sirupsen/logrus"
 )
 
 // synthesizeResponse calls middleware to populate response data, nothing gets pass proxy
 func synthesizeResponse(req *http.Request, middleware string) *http.Response {
 
-	b := bufio.NewScanner(req.Body)
+	// this is mainly for testing, since when you create
+	if req.Body == nil {
+		req.Body = ioutil.NopCloser(bytes.NewBuffer([]byte("")))
+	}
+	responseBody, err := ioutil.ReadAll(req.Body)
+	req.Body.Close()
 
-	bodyStr := b.Text()
+	var bodyStr string
+	if err != nil {
+		log.WithFields(log.Fields{
+			"middleware": middleware,
+			"error":      err.Error(),
+		}).Error("Failed to read request body when synthesizing response")
+	} else {
+		bodyStr = string(responseBody)
+	}
 
 	request := requestDetails{
 		Path:        req.URL.Path,
@@ -24,6 +40,12 @@ func synthesizeResponse(req *http.Request, middleware string) *http.Response {
 	}
 	payload := Payload{Request: request}
 
+	log.WithFields(log.Fields{
+		"middleware":  middleware,
+		"body":        bodyStr,
+		"destination": request.Destination,
+	}).Debug("Synthesizing new response")
+
 	c := NewConstructor(req, payload)
 
 	if middleware != "" {