diff --git a/internal/utils/io.go b/internal/utils/io.go index 70a5bb21..cb5db7e4 100644 --- a/internal/utils/io.go +++ b/internal/utils/io.go @@ -36,6 +36,11 @@ type ( pSrcDst *Pipe pDstSrc *Pipe } + + HookCloser struct { + c io.ReadCloser + hook func() + } ) func NewContextReader(ctx context.Context, r io.Reader) *ContextReader { @@ -143,6 +148,41 @@ const copyBufSize = synk.SizedPoolThreshold var bytesPool = synk.GetBytesPool() +// ReadAllBody reads the body of the response into a buffer and returns it and a function to release the buffer. +func ReadAllBody(resp *http.Response) (buf []byte, release func(), err error) { + if contentLength := resp.ContentLength; contentLength > 0 { + buf = bytesPool.GetSized(int(contentLength)) + _, err = io.ReadFull(resp.Body, buf) + if err != nil { + bytesPool.Put(buf) + return nil, nil, err + } + return buf, func() { bytesPool.Put(buf) }, nil + } + buf, err = io.ReadAll(resp.Body) + if err != nil { + bytesPool.Put(buf) + return nil, nil, err + } + return buf, func() { bytesPool.Put(buf) }, nil +} + +// NewHookCloser wraps a io.ReadCloser and calls the hook function when the closer is closed. +func NewHookCloser(c io.ReadCloser, hook func()) *HookCloser { + return &HookCloser{hook: hook, c: c} +} + +// Close calls the hook function and closes the underlying reader +func (r *HookCloser) Close() error { + r.hook() + return r.c.Close() +} + +// Read reads from the underlying reader. +func (r *HookCloser) Read(p []byte) (int, error) { + return r.c.Read(p) +} + // Copyright 2009 The Go Authors. All rights reserved. // Use of this source code is governed by a BSD-style // This is a copy of io.Copy with context and HTTP flusher handling