fix(plugin): CloseIdleConnections silent no-op when tokenTransport wraps HTTP transport
tokenTransport now delegates CloseIdleConnections to the underlying *http.Transport so http.Client.CloseIdleConnections() works through the wrapper. Three plugin Close() methods updated to unwrap via HTTPClientTransport before calling CloseIdleConnections. RoundTrip now short-circuits when caller already set Authorization, avoiding an unnecessary req.Clone allocation on the pass-through path.
This commit is contained in:
ginuerzh
+55
-11
@@ -104,29 +104,73 @@ func (c *rpcCredentials) RequireTransportSecurity() bool {
|
||||
}
|
||||
|
||||
// NewHTTPClient creates an http.Client from opts. If opts is nil, a default
|
||||
// client with no timeout is returned. Use HTTPClientTransport to access the
|
||||
// underlying transport and close idle connections when the client is no longer
|
||||
// needed.
|
||||
// client with no timeout is returned. When opts.Token is set, the returned
|
||||
// client's transport is wrapped so that each outgoing request carries an
|
||||
// "Authorization: Bearer <token>" header (unless the caller has already set
|
||||
// its own Authorization header). Use HTTPClientTransport to access the
|
||||
// underlying *http.Transport and close idle connections when the client is no
|
||||
// longer needed.
|
||||
func NewHTTPClient(opts *Options) *http.Client {
|
||||
if opts == nil {
|
||||
return &http.Client{}
|
||||
}
|
||||
var transport http.RoundTripper = &http.Transport{
|
||||
TLSClientConfig: opts.TLSConfig,
|
||||
}
|
||||
if opts.Token != "" {
|
||||
transport = &tokenTransport{base: transport, token: opts.Token}
|
||||
}
|
||||
return &http.Client{
|
||||
Timeout: opts.Timeout,
|
||||
Transport: &http.Transport{
|
||||
TLSClientConfig: opts.TLSConfig,
|
||||
},
|
||||
Timeout: opts.Timeout,
|
||||
Transport: transport,
|
||||
}
|
||||
}
|
||||
|
||||
// tokenTransport is an http.RoundTripper that injects a bearer token into the
|
||||
// "Authorization" header of every outgoing request. It does not overwrite an
|
||||
// Authorization header that the caller has already set.
|
||||
type tokenTransport struct {
|
||||
base http.RoundTripper
|
||||
token string
|
||||
}
|
||||
|
||||
// RoundTrip injects the bearer token when the request has no Authorization
|
||||
// header. It clones the request before mutating its headers so that the
|
||||
// caller's request value is not modified.
|
||||
func (t *tokenTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
if req.Header.Get("Authorization") != "" {
|
||||
return t.base.RoundTrip(req)
|
||||
}
|
||||
r := req.Clone(req.Context())
|
||||
r.Header.Set("Authorization", "Bearer "+t.token)
|
||||
return t.base.RoundTrip(r)
|
||||
}
|
||||
|
||||
// CloseIdleConnections delegates to the underlying transport so that
|
||||
// http.Client.CloseIdleConnections() works through the tokenTransport
|
||||
// wrapper instead of silently becoming a no-op.
|
||||
func (t *tokenTransport) CloseIdleConnections() {
|
||||
type closeIdler interface {
|
||||
CloseIdleConnections()
|
||||
}
|
||||
if tr, ok := t.base.(closeIdler); ok {
|
||||
tr.CloseIdleConnections()
|
||||
}
|
||||
}
|
||||
|
||||
// HTTPClientTransport returns the *http.Transport underlying c, or nil if c's
|
||||
// transport is not an *http.Transport. Callers that hold the client for a
|
||||
// long time should call tr.CloseIdleConnections() when the client is no longer
|
||||
// needed.
|
||||
// transport is not an *http.Transport. When the client was created with a
|
||||
// non-empty Token, the *http.Transport is unwrapped from the tokenTransport
|
||||
// wrapper. Callers that hold the client for a long time should call
|
||||
// tr.CloseIdleConnections() when the client is no longer needed.
|
||||
func HTTPClientTransport(c *http.Client) *http.Transport {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
tr, _ := c.Transport.(*http.Transport)
|
||||
rt := c.Transport
|
||||
if tt, ok := rt.(*tokenTransport); ok {
|
||||
rt = tt.base
|
||||
}
|
||||
tr, _ := rt.(*http.Transport)
|
||||
return tr
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ package plugin
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
@@ -110,3 +111,107 @@ func TestNewHTTPClient_TransportTLSConfig(t *testing.T) {
|
||||
t.Error("expected nil TLSClientConfig")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewHTTPClient_WithToken_SetsAuthorizationHeader(t *testing.T) {
|
||||
var gotAuth string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := NewHTTPClient(&Options{Token: "secret-token"})
|
||||
req, err := http.NewRequest(http.MethodGet, srv.URL, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp, err := c.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
|
||||
if want := "Bearer secret-token"; gotAuth != want {
|
||||
t.Errorf("Authorization header: got %q, want %q", gotAuth, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewHTTPClient_NoToken_NoAuthorizationHeader(t *testing.T) {
|
||||
var gotAuth string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := NewHTTPClient(&Options{})
|
||||
req, err := http.NewRequest(http.MethodGet, srv.URL, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp, err := c.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
|
||||
if gotAuth != "" {
|
||||
t.Errorf("Authorization header: got %q, want empty", gotAuth)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewHTTPClient_Token_DoesNotOverrideExistingAuthorization(t *testing.T) {
|
||||
var gotAuth string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := NewHTTPClient(&Options{Token: "from-options"})
|
||||
req, err := http.NewRequest(http.MethodGet, srv.URL, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer caller-set")
|
||||
resp, err := c.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
|
||||
if want := "Bearer caller-set"; gotAuth != want {
|
||||
t.Errorf("Authorization header: got %q, want %q (caller's value must be preserved)", gotAuth, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewHTTPClient_Token_DoesNotMutateCallerRequest(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
c := NewHTTPClient(&Options{Token: "secret-token"})
|
||||
req, err := http.NewRequest(http.MethodGet, srv.URL, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp, err := c.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
|
||||
if got := req.Header.Get("Authorization"); got != "" {
|
||||
t.Errorf("caller's request Authorization header was mutated to %q; expected empty", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPClientTransport_UnwrapsTokenTransport(t *testing.T) {
|
||||
c := NewHTTPClient(&Options{Token: "tok"})
|
||||
tr := HTTPClientTransport(c)
|
||||
if tr == nil {
|
||||
t.Fatal("expected non-nil *http.Transport unwrapped from tokenTransport")
|
||||
}
|
||||
tr.CloseIdleConnections()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user