Add MaxResponseSize to guard against OOMs

This commit adds an optional `PerformRequestOptions.MaxResponseSize` to guard against OOMs. It is added to the `SearchService` as well.

Close #929

client.go

@@ -1230,14 +1230,15 @@ func (c *Client) mustActiveConn() error {
 
 // PerformRequestOptions must be passed into PerformRequest.
 type PerformRequestOptions struct {
-	Method       string
-	Path         string
-	Params       url.Values
-	Body         interface{}
-	ContentType  string
-	IgnoreErrors []int
-	Retrier      Retrier
-	Headers      http.Header
+	Method          string
+	Path            string
+	Params          url.Values
+	Body            interface{}
+	ContentType     string
+	IgnoreErrors    []int
+	Retrier         Retrier
+	Headers         http.Header
+	MaxResponseSize int64
 }
 
 // PerformRequest does a HTTP request to Elasticsearch.
@@ -1376,14 +1377,14 @@ func (c *Client) PerformRequest(ctx context.Context, opt PerformRequestOptions)
 		if err := checkResponse((*http.Request)(req), res, opt.IgnoreErrors...); err != nil {
 			// No retry if request succeeded
 			// We still try to return a response.
-			resp, _ = c.newResponse(res)
+			resp, _ = c.newResponse(res, opt.MaxResponseSize)
 			return resp, err
 		}
 
 		// We successfully made a request with this connection
 		conn.MarkAsHealthy()
 
-		resp, err = c.newResponse(res)
+		resp, err = c.newResponse(res, opt.MaxResponseSize)
 		if err != nil {
 			return nil, err
 		}

client_test.go

@@ -1173,6 +1173,33 @@ func TestPerformRequestWithCustomLogger(t *testing.T) {
 	}
 }
 
+func TestPerformRequestWithMaxResponseSize(t *testing.T) {
+	client, err := NewClient()
+	if err != nil {
+		t.Fatal(err)
+	}
+	res, err := client.PerformRequest(context.TODO(), PerformRequestOptions{
+		Method:          "GET",
+		Path:            "/",
+		MaxResponseSize: 1000,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if res == nil {
+		t.Fatal("expected response to be != nil")
+	}
+
+	res, err = client.PerformRequest(context.TODO(), PerformRequestOptions{
+		Method:          "GET",
+		Path:            "/",
+		MaxResponseSize: 100,
+	})
+	if err != ErrResponseSize {
+		t.Fatal("expected response size error")
+	}
+}
+
 // failingTransport will run a fail callback if it sees a given URL path prefix.
 type failingTransport struct {
 	path string                                      // path prefix to look for

response.go

@@ -6,10 +6,17 @@ package elastic
 
 import (
 	"encoding/json"
+	"errors"
+	"io"
 	"io/ioutil"
 	"net/http"
 )
 
+var (
+	// ErrResponseSize is raised if a response body exceeds the given max body size.
+	ErrResponseSize = errors.New("response size too large")
+)
+
 // Response represents a response from Elasticsearch.
 type Response struct {
 	// StatusCode is the HTTP status code, e.g. 200.
@@ -22,16 +29,26 @@ type Response struct {
 }
 
 // newResponse creates a new response from the HTTP response.
-func (c *Client) newResponse(res *http.Response) (*Response, error) {
+func (c *Client) newResponse(res *http.Response, maxBodySize int64) (*Response, error) {
 	r := &Response{
 		StatusCode: res.StatusCode,
 		Header:     res.Header,
 	}
 	if res.Body != nil {
-		slurp, err := ioutil.ReadAll(res.Body)
+		body := io.Reader(res.Body)
+		if maxBodySize > 0 {
+			if res.ContentLength > maxBodySize {
+				return nil, ErrResponseSize
+			}
+			body = io.LimitReader(body, maxBodySize+1)
+		}
+		slurp, err := ioutil.ReadAll(body)
 		if err != nil {
 			return nil, err
 		}
+		if maxBodySize > 0 && int64(len(slurp)) > maxBodySize {
+			return nil, ErrResponseSize
+		}
 		// HEAD requests return a body but no content
 		if len(slurp) > 0 {
 			r.Body = json.RawMessage(slurp)

response_test.go

@@ -29,7 +29,7 @@ func BenchmarkResponse(b *testing.B) {
 			StatusCode: http.StatusOK,
 		}
 		var err error
-		resp, err = c.newResponse(res)
+		resp, err = c.newResponse(res, 0)
 		if err != nil {
 			b.Fatal(err)
 		}

search.go

@@ -31,6 +31,7 @@ type SearchService struct {
 	ignoreUnavailable *bool
 	allowNoIndices    *bool
 	expandWildcards   string
+	maxResponseSize   int64
 }
 
 // NewSearchService creates a new service for searching in Elasticsearch.
@@ -312,6 +313,13 @@ func (s *SearchService) ExpandWildcards(expandWildcards string) *SearchService {
 	return s
 }
 
+// MaxResponseSize sets an upper limit on the response body size that we accept,
+// to guard against OOM situations.
+func (s *SearchService) MaxResponseSize(maxResponseSize int64) *SearchService {
+	s.maxResponseSize = maxResponseSize
+	return s
+}
+
 // buildURL builds the URL for the operation.
 func (s *SearchService) buildURL() (string, url.Values, error) {
 	var err error
@@ -399,10 +407,11 @@ func (s *SearchService) Do(ctx context.Context) (*SearchResult, error) {
 		body = src
 	}
 	res, err := s.client.PerformRequest(ctx, PerformRequestOptions{
-		Method: "POST",
-		Path:   path,
-		Params: params,
-		Body:   body,
+		Method:          "POST",
+		Path:            path,
+		Params:          params,
+		Body:            body,
+		MaxResponseSize: s.maxResponseSize,
 	})
 	if err != nil {
 		return nil, err