create responder interface so that it's easier to extend a tracker wi…

…th custom responders

json.go

@@ -17,57 +17,75 @@ package chaff
 import (
 	"encoding/json"
 	"fmt"
-	"log"
 	"net/http"
-	"time"
+)
+
+const (
+	// Number of bytes added by the content type header for application/json
+	contentHeaderSize = uint64(31)
 )
 
 // ProduceJSONFn is a function for producing JSON responses.
 type ProduceJSONFn func(string) interface{}
 
-// JSONResponse is an HTTP handler that can wrap a tracker and response with
-// a custom JSON object.
+// The default JSON object that is returned.
+type BasicPadding struct {
+	Padding string `json:"padding"`
+}
+
+// The default ProduceJSONFn
+func PaddingWriterFn(randomData string) interface{} {
+	return &BasicPadding{
+		Padding: randomData,
+	}
+}
+
+// JSONResponder implements the Responder interface and
+// allows you to reply to chaff reqiests with a custom JSON object.
 //
-// To use, you must provide an instanciated Tracker as well as a function
+// To use a function
 // that will be given the heuristically sized payload so you can transform
 // it into the struct that you want to serialize.
-type JSONResponse struct {
-	t  *Tracker
+type JSONResponder struct {
 	fn ProduceJSONFn
 }
 
 // NewJSONResponse creates a new JSON responder
 // Requres a ProduceJSONFn that will be given the random data payload
 // and is responsible for putting it into a struct that can be marshalled
 // as the JSON response.
-func NewJSONResponse(t *Tracker, fn ProduceJSONFn) *JSONResponse {
-	return &JSONResponse{t, fn}
+func NewJSONResponder(fn ProduceJSONFn) Responder {
+	return &JSONResponder{
+		fn: fn,
+	}
 }
 
-func (j *JSONResponse) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	start := time.Now()
-	details := j.t.CalculateProfile()
+func DefaultJSONResponder() Responder {
+	return &JSONResponder{
+		fn: PaddingWriterFn,
+	}
+}
 
+func (j *JSONResponder) Write(headerSize, bodySize uint64, w http.ResponseWriter, r *http.Request) error {
 	var bodyData []byte
 	var err error
-	if details.bodySize > 0 {
-		bodyData, err = json.Marshal(j.fn(randomData(details.bodySize)))
+	if bodySize > 0 {
+		bodyData, err = json.Marshal(j.fn(RandomData(bodySize)))
 		if err != nil {
-			log.Printf("error: unable to marshal chaff JSON response: %v", err)
 			w.WriteHeader(http.StatusInternalServerError)
 			w.Header().Set("Content-Type", "application/json")
 			fmt.Fprintf(w, "{\"error\": \"%v\"}", err.Error())
-			return
+			return err
 		}
 	}
 
 	w.WriteHeader(http.StatusOK)
 	// Generate the response details.
-	if details.headerSize > 0 {
-		w.Header().Add(Header, randomData(details.headerSize))
+	if headerSize > contentHeaderSize {
+		w.Header().Add(Header, RandomData(headerSize-contentHeaderSize-uint64(len(Header))))
 	}
 	w.Header().Set("Content-Type", "application/json")
 	fmt.Fprintf(w, "%s", bodyData)
 
-	j.t.normalizeLatnecy(start, details.latencyMs)
+	return nil
 }

json_test.go

@@ -44,10 +44,10 @@ func TestJSONChaff(t *testing.T) {
 		t.Fatalf("http.NewRequest: %v", err)
 	}
 
-	handler := NewJSONResponse(track, produceExample)
+	responder := NewJSONResponder(produceExample)
 
 	before := time.Now()
-	handler.ServeHTTP(w, r)
+	track.ChaffHandler(responder).ServeHTTP(w, r)
 	after := time.Now()
 
 	if d := after.Sub(before); d < 25*time.Millisecond {
@@ -58,11 +58,14 @@ func TestJSONChaff(t *testing.T) {
 		t.Errorf("wrong code, want: %v, got: %v", http.StatusOK, w.Code)
 	}
 
-	if header := w.Header().Get(Header); header == "" {
-		t.Errorf("expected header '%v' missing", Header)
-	} else {
-		checkLength(t, 100, len(header))
+	headerSize := 0
+	for k, v := range w.Header() {
+		headerSize += len(k)
+		if len(v) > 0 {
+			headerSize += len(v[0])
+		}
 	}
+	checkLength(t, 100, headerSize)
 
 	var response Example
 	if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {

plain.go

@@ -0,0 +1,34 @@
+// Copyright 2020 Mike Helmick
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package chaff
+
+import "net/http"
+
+type PlainResponder struct {
+}
+
+func (pr *PlainResponder) Write(headerSize, bodySize uint64, w http.ResponseWriter, r *http.Request) error {
+	w.WriteHeader(http.StatusOK)
+	// Generate the response details.
+	if headerSize > 0 {
+		w.Header().Add(Header, RandomData(headerSize))
+	}
+	if bodySize > 0 {
+		if _, err := w.Write([]byte(RandomData(bodySize))); err != nil {
+			return err
+		}
+	}
+	return nil
+}

responder.go

@@ -0,0 +1,24 @@
+// Copyright 2020 Mike Helmick
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package chaff
+
+import "net/http"
+
+// Responder allows you to extend the chaff library with custom
+// responders.
+type Responder interface {
+	// Writes the appropriately sized header and body in the desired format.
+	Write(headerSize, bodySize uint64, w http.ResponseWriter, r *http.Request) error
+}

tracker.go

@@ -149,7 +149,7 @@ func (t *Tracker) CalculateProfile() *request {
 	}
 }
 
-func randomData(size uint64) string {
+func RandomData(size uint64) string {
 	// Account for base64 overhead
 	size = 3 * size / 4
 	buffer := make([]byte, size)
@@ -165,29 +165,27 @@ func (t *Tracker) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	t.HandleChaff().ServeHTTP(w, r)
 }
 
-// HandleChaff is the chaff request handler. Based on the current request
-// profile the requst will be held for a certian period of time and then return
-// approximate size random data.
-func (t *Tracker) HandleChaff() http.Handler {
+func (t *Tracker) ChaffHandler(responder Responder) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		start := time.Now()
 		details := t.CalculateProfile()
 
-		w.WriteHeader(http.StatusOK)
-		// Generate the response details.
-		if details.headerSize > 0 {
-			w.Header().Add(Header, randomData(details.headerSize))
-		}
-		if details.bodySize > 0 {
-			if _, err := w.Write([]byte(randomData(details.bodySize))); err != nil {
-				log.Printf("chaff request failed to write: %v", err)
-			}
+		if err := responder.Write(details.headerSize, details.bodySize, w, r); err != nil {
+			log.Printf("error writing chaff response: %v", err)
 		}
 
 		t.normalizeLatnecy(start, details.latencyMs)
 	})
 }
 
+// HandleChaff is the chaff request handler. Based on the current request
+// profile the requst will be held for a certian period of time and then return
+// approximate size random data.
+func (t *Tracker) HandleChaff() http.Handler {
+	responder := &PlainResponder{}
+	return t.ChaffHandler(responder)
+}
+
 // Track wraps a http handler and collects metrics about the request for
 // replaying later during a chaff response. It's suitable for use as a
 // middleware function in common Go web frameworks.