From 4569185e3212e47ffeecbe8ebfb08829bdbd8dde Mon Sep 17 00:00:00 2001
From: Mike Helmick <helmick@google.com>
Date: Tue, 16 Feb 2021 15:17:32 -0800
Subject: [PATCH] max size random data

---
 json.go         |  6 ++++--
 tracker.go      |  9 +++++++++
 tracker_test.go | 17 +++++++++++++++++
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/json.go b/json.go
index 3375cf5..ed780e7 100644
--- a/json.go
+++ b/json.go
@@ -79,10 +79,12 @@ func (j *JSONResponder) Write(headerSize, bodySize uint64, w http.ResponseWriter
 		}
 	}
 
+	const headerDiff = contentHeaderSize + uint64(len(Header))
+
 	w.WriteHeader(http.StatusOK)
 	// Generate the response details.
-	if headerSize > contentHeaderSize {
-		w.Header().Add(Header, RandomData(headerSize-contentHeaderSize-uint64(len(Header))))
+	if headerSize > headerDiff {
+		w.Header().Add(Header, RandomData(headerSize-headerDiff))
 	}
 	w.Header().Set("Content-Type", "application/json")
 	fmt.Fprintf(w, "%s", bodyData)
diff --git a/tracker.go b/tracker.go
index cc1fde7..cf978b4 100644
--- a/tracker.go
+++ b/tracker.go
@@ -28,6 +28,7 @@ import (
 const (
 	Header          = "X-Chaff"
 	DefaultCapacity = 100
+	MaxRandomBytes  = 1000000
 )
 
 // Tracker represents the status of a latency and request size tracker.
@@ -158,9 +159,17 @@ func (t *Tracker) CalculateProfile() *request {
 	}
 }
 
+// RandomData generates size bytes of random base64 data.
 func RandomData(size uint64) string {
 	// Account for base64 overhead
 	size = 3 * size / 4
+	if size <= 0 {
+		return ""
+	}
+	if size > MaxRandomBytes {
+		size = MaxRandomBytes
+	}
+
 	buffer := make([]byte, size)
 	_, err := rand.Read(buffer)
 	if err != nil {
diff --git a/tracker_test.go b/tracker_test.go
index f0c716f..0bdc4e4 100644
--- a/tracker_test.go
+++ b/tracker_test.go
@@ -15,6 +15,7 @@
 package chaff
 
 import (
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -27,6 +28,22 @@ import (
 	"github.com/google/go-cmp/cmp"
 )
 
+func TestRandomData(t *testing.T) {
+	d := RandomData(0)
+	if d != "" {
+		t.Fatalf("expected empty string, got: %q", d)
+	}
+
+	d = RandomData(MaxRandomBytes * 2)
+	b, err := base64.StdEncoding.DecodeString(d)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if l := len(b); l < int(float32(MaxRandomBytes)*0.99) || l > int(float32(MaxRandomBytes)*1.01) {
+		t.Fatalf("length is outside of 1pct of expected, want: %d got: %d", MaxRandomBytes, l)
+	}
+}
+
 func checkLength(t *testing.T, expected int, length int) {
 	t.Helper()
 	lower := float64(expected) * 0.99