Merge pull request #56 from SpectoLabs/develop

Develop

Author: Karolis Rusenas

admin.go

@@ -152,12 +152,11 @@ func (d *DBClient) ImportRecordsHandler(w http.ResponseWriter, req *http.Request
 	payloads := requests.Data
 	if len(payloads) > 0 {
 		for _, pl := range payloads {
-			bts, err := json.Marshal(pl)
-
+			bts, err := pl.encode()
 			if err != nil {
 				log.WithFields(log.Fields{
 					"error": err.Error(),
-				}).Error("Failed to marshal json")
+				}).Error("Failed to encode payload")
 			} else {
 				// recalculating request hash and storing it in database
 				r := request{details: pl.Request}

cache.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"encoding/json"
 	"fmt"
 
 	log "github.com/Sirupsen/logrus"
@@ -66,16 +65,14 @@ func (c *Cache) GetAllRequests() (payloads []Payload, err error) {
 		c := b.Cursor()
 
 		for k, v := c.First(); k != nil; k, v = c.Next() {
-			var pl Payload
-			err = json.Unmarshal(v, &pl)
-
+			pl, err := decodePayload(v)
 			if err != nil {
 				log.WithFields(log.Fields{
 					"error": err.Error(),
 					"json":  v,
-				}).Warning("Failed to deserialize json")
+				}).Warning("Failed to deserialize bytes to payload.")
 			} else {
-				payloads = append(payloads, pl)
+				payloads = append(payloads, *pl)
 			}
 		}
 		return nil

models.go

@@ -3,7 +3,7 @@ package main
 import (
 	"bytes"
 	"crypto/md5"
-	"encoding/json"
+	"encoding/gob"
 	"fmt"
 	"io"
 	"net/http"
@@ -62,6 +62,29 @@ type Payload struct {
 	ID       string         `json:"id"`
 }
 
+// encode method encodes all exported Payload fields to bytes
+func (p *Payload) encode() ([]byte, error) {
+	buf := new(bytes.Buffer)
+	enc := gob.NewEncoder(buf)
+	err := enc.Encode(p)
+	if err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+// decodePayload decodes supplied bytes into Payload structure
+func decodePayload(data []byte) (*Payload, error) {
+	var p *Payload
+	buf := bytes.NewBuffer(data)
+	dec := gob.NewDecoder(buf)
+	err := dec.Decode(&p)
+	if err != nil {
+		return nil, err
+	}
+	return p, nil
+}
+
 // recordRequest saves request for later playback
 func (d *DBClient) captureRequest(req *http.Request) (*http.Response, error) {
 
@@ -212,18 +235,16 @@ func (d *DBClient) save(req *http.Request, resp *http.Response, respBody []byte,
 			Request:  requestObj,
 			ID:       key,
 		}
-		// converting it to json bytes
-		bts, err := json.Marshal(payload)
 
+		bts, err := payload.encode()
 		if err != nil {
 			log.WithFields(log.Fields{
 				"error": err.Error(),
-			}).Error("Failed to marshal json")
+			}).Error("Failed to serialize payload")
 		} else {
 			d.cache.Set([]byte(key), bts)
 		}
 	}
-
 }
 
 // getRequestFingerprint returns request hash
@@ -237,19 +258,20 @@ func getRequestFingerprint(req *http.Request) string {
 func (d *DBClient) getResponse(req *http.Request) *http.Response {
 
 	key := getRequestFingerprint(req)
-	var payload Payload
+	//	var payload Payload
 
 	payloadBts, err := d.cache.Get([]byte(key))
 
 	if err == nil {
 		// getting cache response
-		err = json.Unmarshal(payloadBts, &payload)
+		payload, err := decodePayload(payloadBts)
 		if err != nil {
-			log.Error(err)
-			// what now?
+			log.WithFields(log.Fields{
+				"error": err.Error(),
+			}).Error("Failed to decode payload")
 		}
 
-		c := NewConstructor(req, payload)
+		c := NewConstructor(req, *payload)
 
 		if d.cfg.middleware != "" {
 			_ = c.ApplyMiddleware(d.cfg.middleware)

models_test.go

@@ -2,7 +2,6 @@ package main
 
 import (
 	"bytes"
-	"encoding/json"
 	"fmt"
 	"io/ioutil"
 	"net/http"
@@ -58,16 +57,10 @@ func TestRequestBodyCaptured(t *testing.T) {
 	fp := getRequestFingerprint(req)
 
 	payloadBts, err := dbClient.cache.Get([]byte(fp))
-
-	var payload Payload
-
 	expect(t, err, nil)
 
-	// getting cache response
-	err = json.Unmarshal(payloadBts, &payload)
-
+	payload, err := decodePayload(payloadBts)
 	expect(t, err, nil)
-
 	expect(t, payload.Request.Body, "fizz=buzz")
 
 }
@@ -98,3 +91,37 @@ func TestDeleteAllRecords(t *testing.T) {
 	err := dbClient.cache.DeleteBucket(dbClient.cache.requestsBucket)
 	expect(t, err, nil)
 }
+
+func TestPayloadEncodeDecode(t *testing.T) {
+	resp := response{
+		Status: 200,
+		Body:   "body here",
+	}
+
+	payload := Payload{Response: resp}
+
+	bts, err := payload.encode()
+	expect(t, err, nil)
+
+	pl, err := decodePayload(bts)
+	expect(t, err, nil)
+	expect(t, pl.Response.Body, resp.Body)
+	expect(t, pl.Response.Status, resp.Status)
+
+}
+
+func TestPayloadEncodeEmpty(t *testing.T) {
+	payload := Payload{}
+
+	bts, err := payload.encode()
+	expect(t, err, nil)
+
+	_, err = decodePayload(bts)
+	expect(t, err, nil)
+}
+
+func TestDecodeRandomBytes(t *testing.T) {
+	bts := []byte("some random stuff here")
+	_, err := decodePayload(bts)
+	refute(t, err, nil)
+}