/
utils.go
49 lines (42 loc) · 1.02 KB
/
utils.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
/*************************************************************************
* Copyright 2023 Gravwell, Inc. All rights reserved.
* Contact: <legal@gravwell.io>
*
* This software may be modified and distributed under the terms of the
* BSD 2-clause license. See the LICENSE file for details.
**************************************************************************/
package client
import (
"errors"
"io"
"net/http"
)
const (
maxDataDrain = 1024 * 1024 * 4
)
var (
ErrMaxBodyDrained = errors.New("too many response bytes in body, closing")
)
// drainResponse will drain up to 4MB of data then close the response Body.
// We do this so that http requests can re-use connects as per docs.
func drainResponse(resp *http.Response) {
if resp.Body == nil {
return
}
var nw nilWriter
io.Copy(&nw, resp.Body)
resp.Body.Close()
}
type nilWriter struct {
n int
}
func (nw *nilWriter) Write(b []byte) (r int, err error) {
if nw.n > maxDataDrain {
r = -1
err = ErrMaxBodyDrained
return
}
nw.n += len(b)
r = len(b)
return
}