client: fix `Agent` use after free (#2037)

* fix

* fix

* fix
pull/2040/head
Trim21 2022-08-22 14:01:40 +08:00 committed by GitHub
parent 32d311cd59
commit 80a6fdc215
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 21 additions and 11 deletions

View File

@ -741,9 +741,14 @@ func (a *Agent) RetryIf(retryIf RetryIfFunc) *Agent {
/************************** End Agent Setting **************************/ /************************** End Agent Setting **************************/
// Bytes returns the status code, bytes body and errors of url. // Bytes returns the status code, bytes body and errors of url.
//
// it's not safe to use Agent after calling [Agent.Bytes]
func (a *Agent) Bytes() (code int, body []byte, errs []error) { func (a *Agent) Bytes() (code int, body []byte, errs []error) {
defer a.release() defer a.release()
return a.bytes()
}
func (a *Agent) bytes() (code int, body []byte, errs []error) {
if errs = append(errs, a.errs...); len(errs) > 0 { if errs = append(errs, a.errs...); len(errs) > 0 {
return return
} }
@ -802,16 +807,22 @@ func printDebugInfo(req *Request, resp *Response, w io.Writer) {
} }
// String returns the status code, string body and errors of url. // String returns the status code, string body and errors of url.
//
// it's not safe to use Agent after calling [Agent.String]
func (a *Agent) String() (int, string, []error) { func (a *Agent) String() (int, string, []error) {
code, body, errs := a.Bytes() defer a.release()
code, body, errs := a.bytes()
return code, utils.UnsafeString(body), errs return code, utils.UnsafeString(body), errs
} }
// Struct returns the status code, bytes body and errors of url. // Struct returns the status code, bytes body and errors of url.
// And bytes body will be unmarshalled to given v. // And bytes body will be unmarshalled to given v.
//
// it's not safe to use Agent after calling [Agent.Struct]
func (a *Agent) Struct(v interface{}) (code int, body []byte, errs []error) { func (a *Agent) Struct(v interface{}) (code int, body []byte, errs []error) {
if code, body, errs = a.Bytes(); len(errs) > 0 { defer a.release()
if code, body, errs = a.bytes(); len(errs) > 0 {
return return
} }
@ -856,8 +867,12 @@ func (a *Agent) reset() {
} }
var ( var (
clientPool sync.Pool clientPool sync.Pool
agentPool sync.Pool agentPool = sync.Pool{
New: func() interface{} {
return &Agent{req: &Request{}}
},
}
responsePool sync.Pool responsePool sync.Pool
argsPool sync.Pool argsPool sync.Pool
formFilePool sync.Pool formFilePool sync.Pool
@ -895,11 +910,7 @@ func ReleaseClient(c *Client) {
// no longer needed. This allows Agent recycling, reduces GC pressure // no longer needed. This allows Agent recycling, reduces GC pressure
// and usually improves performance. // and usually improves performance.
func AcquireAgent() *Agent { func AcquireAgent() *Agent {
v := agentPool.Get() return agentPool.Get().(*Agent)
if v == nil {
return &Agent{req: &Request{}}
}
return v.(*Agent)
} }
// ReleaseAgent returns a acquired via AcquireAgent to Agent pool. // ReleaseAgent returns a acquired via AcquireAgent to Agent pool.

View File

@ -4,6 +4,7 @@ import (
"bytes" "bytes"
"crypto/tls" "crypto/tls"
"encoding/base64" "encoding/base64"
"encoding/json"
"errors" "errors"
"fmt" "fmt"
"io" "io"
@ -16,8 +17,6 @@ import (
"testing" "testing"
"time" "time"
"encoding/json"
"github.com/gofiber/fiber/v2/internal/tlstest" "github.com/gofiber/fiber/v2/internal/tlstest"
"github.com/gofiber/fiber/v2/utils" "github.com/gofiber/fiber/v2/utils"
"github.com/valyala/fasthttp/fasthttputil" "github.com/valyala/fasthttp/fasthttputil"