From a4ad41c09fe81750800e2e10df7c9d0e212b4a7a Mon Sep 17 00:00:00 2001
From: Andrew Cholakian <andrew@andrewvc.com>
Date: Wed, 22 Jan 2020 15:45:54 -0600
Subject: [PATCH] [Heartbeat] Fix excessive memory usage when parsing bodies
 (#15639)

* [Heartbeat] Fix excessive memory usage when parsing bodies

In 7.5 we introduced a regression where we would allocate a 100MiB
buffer per HTTP request. This change fixes that to stream data instead.

(cherry picked from commit 080dedb6de75d5776525b78dcf65febeb5761199)
---
 CHANGELOG.next.asciidoc                    |  1 +
 heartbeat/monitors/active/http/respbody.go | 51 +++++++---------------
 2 files changed, 16 insertions(+), 36 deletions(-)

diff --git a/CHANGELOG.next.asciidoc b/CHANGELOG.next.asciidoc
index c6d8c6cfa02..7eebca4818c 100644
--- a/CHANGELOG.next.asciidoc
+++ b/CHANGELOG.next.asciidoc
@@ -83,6 +83,7 @@ https://github.com/elastic/beats/compare/v7.0.0-alpha2...master[Check the HEAD d
 *Heartbeat*
 
 - Fix recording of SSL cert metadata for Expired/Unvalidated x509 certs. {pull}13687[13687]
+- Fixed excessive memory usage introduced in 7.5 due to over-allocating memory for HTTP checks. {pull}15639[15639]
 
 *Journalbeat*
 
diff --git a/heartbeat/monitors/active/http/respbody.go b/heartbeat/monitors/active/http/respbody.go
index 7990b6324e7..b8b95e5b023 100644
--- a/heartbeat/monitors/active/http/respbody.go
+++ b/heartbeat/monitors/active/http/respbody.go
@@ -22,7 +22,9 @@ import (
 	"encoding/hex"
 	"io"
 	"net/http"
-	"unicode/utf8"
+	"strings"
+
+	"github.com/docker/go-units"
 
 	"github.com/elastic/beats/heartbeat/reason"
 	"github.com/elastic/beats/libbeat/common"
@@ -31,7 +33,7 @@ import (
 // maxBufferBodyBytes sets a hard limit on how much we're willing to buffer for any reason internally.
 // since we must buffer the whole body for body validators this is effectively a cap on that.
 // 100MiB out to be enough for everybody.
-const maxBufferBodyBytes = 100 * 1024 * 1024
+const maxBufferBodyBytes = 100 * units.MiB
 
 func processBody(resp *http.Response, config responseConfig, validator multiValidator) (common.MapStr, reason.Reason) {
 	// Determine how much of the body to actually buffer in memory
@@ -94,43 +96,20 @@ func readBody(resp *http.Response, maxSampleBytes int) (bodySample string, bodyS
 
 func readPrefixAndHash(body io.ReadCloser, maxPrefixSize int) (respSize int, prefix string, hashStr string, err error) {
 	hash := sha256.New()
-	// Function to lazily get the body of the response
-	rawBuf := make([]byte, 1024)
-
-	// Buffer to hold the prefix output along with tracking info
-	prefixBuf := make([]byte, maxPrefixSize)
-	prefixRemainingBytes := maxPrefixSize
-	prefixWriteOffset := 0
-	for {
-		readSize, readErr := body.Read(rawBuf)
-
-		respSize += readSize
-		hash.Write(rawBuf[:readSize])
-
-		if prefixRemainingBytes > 0 {
-			if readSize >= prefixRemainingBytes {
-				copy(prefixBuf[prefixWriteOffset:maxPrefixSize], rawBuf[:prefixRemainingBytes])
-				prefixWriteOffset += prefixRemainingBytes
-				prefixRemainingBytes = 0
-			} else {
-				copy(prefixBuf[prefixWriteOffset:prefixWriteOffset+readSize], rawBuf[:readSize])
-				prefixWriteOffset += readSize
-				prefixRemainingBytes -= readSize
-			}
-		}
 
-		if readErr == io.EOF {
-			break
-		}
+	var prefixBuf strings.Builder
 
-		if readErr != nil {
-			return 0, "", "", readErr
-		}
+	n, err := io.Copy(&prefixBuf, io.TeeReader(io.LimitReader(body, int64(maxPrefixSize)), hash))
+	if err == nil {
+		// finish streaming into hash if the body has not been fully consumed yet
+		var m int64
+		m, err = io.Copy(hash, body)
+		n += m
 	}
 
-	// We discard the body if it is not valid UTF-8
-	if utf8.Valid(prefixBuf[:prefixWriteOffset]) {
-		prefix = string(prefixBuf[:prefixWriteOffset])
+	if err != nil && err != io.EOF {
+		return 0, "", "", err
 	}
-	return respSize, prefix, hex.EncodeToString(hash.Sum(nil)), nil
+
+	return int(n), prefixBuf.String(), hex.EncodeToString(hash.Sum(nil)), nil
 }