Allow chaff to be used as middleware on a single request (#1)

* Add concept of "detectors"

This will enable smarter chaff middleware

* Handle race in buffer

* Allow chaff to be used as middleware on a single request

* Fix vet and staticcheck errors

README.md

@@ -4,36 +4,50 @@
 [![Go](https://github.com/mikehelmick/go-chaff/workflows/Go/badge.svg?event=push)](https://github.com/mikehelmick/go-chaff/actions?query=workflow%3AGo)
 
 This package provides the necessary tools to allow for your server to handle
-chaff requests from clients. This technique can be used when you want to guard
-against the fact that clients are connecting to your server is meaningful.
+chaff (fake) requests from clients. This technique can be used when you want to
+guard against the fact that clients are connecting to your server is meaningful.
 
-Use of this method allows your clients to periodically connect to the server
-sending chaff instead of real requests.
+The tracker automatically captures metadata like average request time and
+response size, with the aim of making a chaff request indistinguishable from a
+real request. This is useful in situations where someone (e.g. server operator,
+network peer) should not be able to glean information about the system from
+requests, their size, or their frequency.
 
-There are two components, the middleware function that implements the tracking
-and an http.Handler that serves the chaff requests.
+Clients periodically send "chaff" requests. They denote the request is chaff via
+a header or similar identifier. If one of your goals is to obfuscate server
+logs, a dedicated URL is not recommended as this will be easily distinguisable
+in logs.
+
+There are two components:
+
+- a middleware function that implements tracking
+- an `http.Handler` that serves the chaff requests
 
 ## Usage
 
-1. Create the tracker and install the middleware on routes you want to
-   simulate the request latency and response size of.
-
-```golang
-r := mux.NewRouter()
-track := chaff.New()
-defer track.Close()
-{
-  sub := r.PathPrefix("").Subrouter()
-  sub.Use(track.Track)
-  sub.Handle(...) // your actual methods to simulate
-}
-```
-
-2. Install a handler to handle the chaff requests.
-
-```golang
-{
-  sub := r.PathPrefix("/chaff").Subrouter()
-  sub.Handle("", track).Methods("GET")
-}
-```
+1.  Option 1 - use a single handler, detect chaff based on a request property
+    like a header. This is most useful when you don't trust the server operator
+    and can have the performance hit of the branching logic in a single handler:
+
+    ```go
+    mux := http.NewServeMux()
+    mux.Handle("/", tracker.HandleTrack(chaff.HeaderDetector("X-Chaff"), myHandler))
+    ```
+
+    In this example, requests to `/` are served normally and the tracker
+    generates heuristics automatically. When a request includes an `X-Chaff`
+    header, the handler sends a chaff response.
+
+1.  Option 2 - create the tracker on specific routes and provide a dedicated
+    chaff endpoint. This is useful when you trust the server operator, but not
+    the network observer:
+
+    ```go
+    r := mux.NewRouter()
+    tracker := chaff.New()
+    defer tracker.Close()
+
+    mux := http.NewServeMux()
+    mux.Handle("/", tracker.Track())
+    mux.Handle("/chaff", tracker.HandleChaff())
+    ```

detector.go

@@ -0,0 +1,38 @@
+// Copyright 2020 Mike Helmick
+// Copyright 2020 Seth Vargo
+//
+// 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 Detector interface {
+	IsChaff(r *http.Request) bool
+}
+
+var _ Detector = (DetectorFunc)(nil)
+
+type DetectorFunc func(r *http.Request) bool
+
+func (d DetectorFunc) IsChaff(r *http.Request) bool {
+	return d(r)
+}
+
+// HeaderDetector is a detector that searches for the header's presence to mark
+// a request as chaff.
+func HeaderDetector(h string) Detector {
+	return DetectorFunc(func(r *http.Request) bool {
+		return r.Header.Get(h) != ""
+	})
+}

json.go

@@ -22,7 +22,7 @@ import (
 	"time"
 )
 
-// ProduceJSONFn
+// 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
@@ -56,7 +56,7 @@ func (j *JSONResponse) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			log.Printf("error: unable to marshal chaff JSON response: %v", err)
 			w.WriteHeader(http.StatusInternalServerError)
 			w.Header().Set("Content-Type", "application/json")
-			w.Write([]byte(fmt.Sprintf("{\"error\": \"%v\"}", err.Error())))
+			fmt.Fprintf(w, "{\"error\": \"%v\"}", err.Error())
 			return
 		}
 	}
@@ -67,7 +67,7 @@ func (j *JSONResponse) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		w.Header().Add(Header, randomData(details.headerSize))
 	}
 	w.Header().Set("Content-Type", "application/json")
-	w.Write(bodyData)
+	fmt.Fprintf(w, "%s", bodyData)
 
 	j.t.normalizeLatnecy(start, details.latencyMs)
 }

json_test.go

@@ -58,10 +58,10 @@ func TestJSONChaff(t *testing.T) {
 		t.Errorf("wrong code, want: %v, got: %v", http.StatusOK, w.Code)
 	}
 
-	if header, ok := w.HeaderMap[Header]; !ok {
+	if header := w.Header().Get(Header); header == "" {
 		t.Errorf("expected header '%v' missing", Header)
 	} else {
-		checkLength(t, 100, len(header[0]))
+		checkLength(t, 100, len(header))
 	}
 
 	var response Example

tracker.go

@@ -21,6 +21,7 @@ import (
 	"log"
 	"net/http"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -50,14 +51,14 @@ type Tracker struct {
 }
 
 type request struct {
-	latencyMs  int64
-	bodySize   int
-	headerSize int
+	latencyMs  uint64
+	bodySize   uint64
+	headerSize uint64
 }
 
-func newRequest(start, end time.Time, headerSize, bodySize int) *request {
+func newRequest(start, end time.Time, headerSize, bodySize uint64) *request {
 	return &request{
-		latencyMs:  end.Sub(start).Milliseconds(),
+		latencyMs:  uint64(end.Sub(start).Milliseconds()),
 		headerSize: headerSize,
 		bodySize:   bodySize,
 	}
@@ -133,21 +134,22 @@ func (t *Tracker) CalculateProfile() *request {
 		return &request{}
 	}
 
-	var latency, hSize, bSize int64
+	var latency, hSize, bSize uint64
 	for _, r := range t.buffer {
 		latency += r.latencyMs
-		hSize += int64(r.headerSize)
-		bSize += int64(r.bodySize)
+		hSize += uint64(r.headerSize)
+		bSize += uint64(r.bodySize)
 	}
-	divisor := int64(t.size)
+	divisor := uint64(t.size)
+
 	return &request{
 		latencyMs:  latency / divisor,
-		headerSize: int(hSize / divisor),
-		bodySize:   int(bSize / divisor),
+		headerSize: uint64(hSize / divisor),
+		bodySize:   uint64(bSize / divisor),
 	}
 }
 
-func randomData(size int) string {
+func randomData(size uint64) string {
 	// Account for base64 overhead
 	size = 3 * size / 4
 	buffer := make([]byte, size)
@@ -158,38 +160,86 @@ func randomData(size int) string {
 	return base64.StdEncoding.EncodeToString(buffer)
 }
 
-// ServerHTTP 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.
+// ServeHTTP implements http.Handler. See HandleChaff for more details.
 func (t *Tracker) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	start := time.Now()
-	details := t.CalculateProfile()
+	t.HandleChaff().ServeHTTP(w, r)
+}
 
-	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)
+// 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 {
+	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)
+			}
+		}
+
+		t.normalizeLatnecy(start, details.latencyMs)
+	})
+}
+
+// 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.
+func (t *Tracker) Track(next http.Handler) http.Handler {
+	return t.HandleTrack(nil, next)
+}
+
+// HandleTrack wraps the given http handler and detector. If the request is
+// deemed to be chaff (as determined by the Detector), the system sends a chaff
+// response. Otherwise it returns the real response and adds it to the tracker.
+func (t *Tracker) HandleTrack(d Detector, next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if d != nil && d.IsChaff(r) {
+			// Send chaff response
+			t.HandleChaff().ServeHTTP(w, r)
+			return
+		}
+
+		// Handle the real request, gathering metadata
+		start := time.Now()
+		proxyWriter := &writeThrough{w: w}
+		next.ServeHTTP(proxyWriter, r)
+		end := time.Now()
 
-	t.normalizeLatnecy(start, details.latencyMs)
+		// Grab the size of the headers that are present.
+		var headerSize uint64
+		for k, vals := range w.Header() {
+			headerSize += uint64(len(k))
+			for _, v := range vals {
+				headerSize += uint64(len(v))
+			}
+		}
+
+		// Save metadata
+		select {
+		case t.ch <- newRequest(start, end, headerSize, proxyWriter.Size()):
+		default: // channel full, drop request.
+		}
+	})
 }
 
-func (t *Tracker) normalizeLatnecy(start time.Time, targetMs int64) {
-	elapsed := time.Now().Sub(start)
-	if rem := targetMs - elapsed.Milliseconds(); rem > 0 {
+func (t *Tracker) normalizeLatnecy(start time.Time, targetMs uint64) {
+	elapsed := time.Since(start)
+	if rem := targetMs - uint64(elapsed.Milliseconds()); rem > 0 {
 		time.Sleep(time.Duration(rem) * time.Millisecond)
 	}
 }
 
-// write through wraps an http.ResponseWriter so that we can count the
-// number of bytes that are written by the delegate handler.
+// write through wraps an http.ResponseWriter so that we can count the number of
+// bytes that are written by the delegate handler.
 type writeThrough struct {
-	size int
+	size uint64
 	w    http.ResponseWriter
 }
 
@@ -198,33 +248,14 @@ func (wt *writeThrough) Header() http.Header {
 }
 
 func (wt *writeThrough) Write(b []byte) (int, error) {
-	wt.size += len(b)
+	atomic.AddUint64(&wt.size, uint64(len(b)))
 	return wt.w.Write(b)
 }
 
 func (wt *writeThrough) WriteHeader(statusCode int) {
 	wt.w.WriteHeader(statusCode)
 }
 
-// Track provides the necessary http middleware function.
-func (t *Tracker) Track(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		start := time.Now()
-		proxyWriter := &writeThrough{w: w}
-		next.ServeHTTP(proxyWriter, r)
-		end := time.Now()
-
-		// grab the size of the headers that are present.
-		headerSize := 0
-		for k, vals := range w.Header() {
-			headerSize += len(k)
-			for _, v := range vals {
-				headerSize += len(v)
-			}
-		}
-		select {
-		case t.ch <- newRequest(start, end, headerSize, proxyWriter.size):
-		default: // channel full, drop request.
-		}
-	})
+func (wt *writeThrough) Size() uint64 {
+	return atomic.LoadUint64(&wt.size)
 }

tracker_test.go

@@ -15,6 +15,7 @@
 package chaff
 
 import (
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -59,10 +60,10 @@ func TestChaff(t *testing.T) {
 		t.Errorf("wrong code, want: %v, got: %v", http.StatusOK, w.Code)
 	}
 
-	if header, ok := w.HeaderMap[Header]; !ok {
+	if header := w.Header().Get(Header); header == "" {
 		t.Errorf("expected header '%v' missing", Header)
 	} else {
-		checkLength(t, 100, len(header[0]))
+		checkLength(t, 100, len(header))
 	}
 	checkLength(t, 250, len(w.Body.Bytes()))
 }
@@ -85,7 +86,7 @@ func TestTracking(t *testing.T) {
 				time.Sleep(1 * time.Millisecond)
 				w.WriteHeader(http.StatusAccepted)
 				w.Header().Add("padding", strings.Repeat("a", i+1))
-				w.Write([]byte(strings.Repeat("b", i+1)))
+				fmt.Fprintf(w, "%s", strings.Repeat("b", i+1))
 			}))
 
 		recorder := httptest.NewRecorder()