http2: limit canonical header cache by bytes, not entries

The canonical header cache is a per-connection cache mapping header
keys to their canonicalized form. (For example, "foo-bar" => "Foo-Bar").
We limit the number of entries in the cache to prevent an attacker
from consuming unbounded amounts of memory by sending many unique
keys, but a small number of very large keys can still consume an
unreasonable amount of memory.

Track the amount of memory consumed by the cache and limit it based
on memory rather than number of entries.

Thanks to Josselin Costanzi for reporting this issue.

For golang/go#56350

Change-Id: I41db4c9823ed5bf371a9881accddff1268489b16
Reviewed-on: https://go-review.googlesource.com/c/net/+/455635
Reviewed-by: Jenny Rakoczy <jenny@golang.org>
Run-TryBot: Damien Neil <dneil@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
diff --git a/http2/server.go b/http2/server.go
index e35a76c..4eb7617 100644
--- a/http2/server.go
+++ b/http2/server.go
@@ -588,6 +588,7 @@
 	maxFrameSize                int32
 	peerMaxHeaderListSize       uint32            // zero means unknown (default)
 	canonHeader                 map[string]string // http2-lower-case -> Go-Canonical-Case
+	canonHeaderKeysSize         int               // canonHeader keys size in bytes
 	writingFrame                bool              // started writing a frame (on serve goroutine or separate)
 	writingFrameAsync           bool              // started a frame on its own goroutine but haven't heard back on wroteFrameCh
 	needsFrameFlush             bool              // last frame write wasn't a flush
@@ -766,6 +767,13 @@
 	}
 }
 
+// maxCachedCanonicalHeadersKeysSize is an arbitrarily-chosen limit on the size
+// of the entries in the canonHeader cache.
+// This should be larger than the size of unique, uncommon header keys likely to
+// be sent by the peer, while not so high as to permit unreasonable memory usage
+// if the peer sends an unbounded number of unique header keys.
+const maxCachedCanonicalHeadersKeysSize = 2048
+
 func (sc *serverConn) canonicalHeader(v string) string {
 	sc.serveG.check()
 	buildCommonHeaderMapsOnce()
@@ -781,14 +789,10 @@
 		sc.canonHeader = make(map[string]string)
 	}
 	cv = http.CanonicalHeaderKey(v)
-	// maxCachedCanonicalHeaders is an arbitrarily-chosen limit on the number of
-	// entries in the canonHeader cache. This should be larger than the number
-	// of unique, uncommon header keys likely to be sent by the peer, while not
-	// so high as to permit unreasonable memory usage if the peer sends an unbounded
-	// number of unique header keys.
-	const maxCachedCanonicalHeaders = 32
-	if len(sc.canonHeader) < maxCachedCanonicalHeaders {
+	size := 100 + len(v)*2 // 100 bytes of map overhead + key + value
+	if sc.canonHeaderKeysSize+size <= maxCachedCanonicalHeadersKeysSize {
 		sc.canonHeader[v] = cv
+		sc.canonHeaderKeysSize += size
 	}
 	return cv
 }
diff --git a/http2/server_test.go b/http2/server_test.go
index 3760871..a1e086c 100644
--- a/http2/server_test.go
+++ b/http2/server_test.go
@@ -4543,3 +4543,28 @@
 		})
 	}
 }
+
+// TestCanonicalHeaderCacheGrowth verifies that the canonical header cache
+// size is capped to a reasonable level.
+func TestCanonicalHeaderCacheGrowth(t *testing.T) {
+	defer disableGoroutineTracking()()
+	for _, size := range []int{1, (1 << 20) - 10} {
+		base := strings.Repeat("X", size)
+		sc := &serverConn{}
+		const count = 1000
+		for i := 0; i < count; i++ {
+			h := fmt.Sprintf("%v-%v", base, i)
+			c := sc.canonicalHeader(h)
+			if len(h) != len(c) {
+				t.Errorf("sc.canonicalHeader(%q) = %q, want same length", h, c)
+			}
+		}
+		total := 0
+		for k, v := range sc.canonHeader {
+			total += len(k) + len(v) + 100
+		}
+		if total > maxCachedCanonicalHeadersKeysSize {
+			t.Errorf("after adding %v ~%v-byte headers, canonHeader cache is ~%v bytes, want <%v", count, size, total, maxCachedCanonicalHeadersKeysSize)
+		}
+	}
+}