diff --git a/client/client.go b/client/client.go index 93d9f4f49a..d8e69953d5 100644 --- a/client/client.go +++ b/client/client.go @@ -59,6 +59,7 @@ import ( "net/http" "net/url" "path" + "slices" "strings" "sync" "sync/atomic" @@ -241,6 +242,13 @@ func New(ops ...Opt) (*Client, error) { c.client.Transport = otelhttp.NewTransport(c.client.Transport, c.traceOpts...) + if len(cfg.responseHooks) > 0 { + c.client.Transport = &responseHookTransport{ + base: c.client.Transport, + hooks: slices.Clone(cfg.responseHooks), + } + } + return c, nil } diff --git a/client/client_options.go b/client/client_options.go index ae80c4e5cf..c37fc9bceb 100644 --- a/client/client_options.go +++ b/client/client_options.go @@ -2,6 +2,7 @@ package client import ( "context" + "errors" "fmt" "net" "net/http" @@ -55,10 +56,19 @@ type clientConfig struct { // takes precedence. Either field disables API-version negotiation. envAPIVersion string + // responseHooks is a list of custom response hooks to call on responses. + responseHooks []ResponseHook + // traceOpts is a list of options to configure the tracing span. traceOpts []otelhttp.Option } +// ResponseHook is called for each HTTP response returned by the daemon. +// Hooks are invoked in the order they were added. +// +// Hooks must not read or close resp.Body. +type ResponseHook func(*http.Response) error + // Opt is a configuration option to initialize a [Client]. type Opt func(*clientConfig) error @@ -348,3 +358,18 @@ func WithTraceOptions(opts ...otelhttp.Option) Opt { return nil } } + +// WithResponseHook adds a ResponseHook to the client. ResponseHooks are called +// for each HTTP response returned by the daemon. Hooks are invoked in the order +// they were added. +// +// Hooks must not read or close resp.Body. +func WithResponseHook(h ResponseHook) Opt { + return func(c *clientConfig) error { + if h == nil { + return errors.New("invalid response hook: hook is nil") + } + c.responseHooks = append(c.responseHooks, h) + return nil + } +} diff --git a/client/client_options_test.go b/client/client_options_test.go index a0af40f349..63ed47127a 100644 --- a/client/client_options_test.go +++ b/client/client_options_test.go @@ -2,6 +2,8 @@ package client import ( "crypto/tls" + "errors" + "io" "net/http" "net/http/cookiejar" "runtime" @@ -390,3 +392,106 @@ func TestWithHTTPClient(t *testing.T) { cmpopts.IgnoreUnexported(http.Transport{}, tls.Config{}), cmpopts.EquateComparable(&cookiejar.Jar{})) } + +func TestWithResponseHook(t *testing.T) { + const hdrKey = "X-Test-Header" + const hdrVal = "hello-world" + + t.Run("single hook", func(t *testing.T) { + var got string + c, err := New( + WithResponseHook(func(resp *http.Response) error { + got = resp.Header.Get(hdrKey) + return nil + }), + WithBaseMockClient(func(req *http.Request) (*http.Response, error) { + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + } + resp.Header.Set(hdrKey, hdrVal) + return resp, nil + }), + ) + assert.NilError(t, err) + + _, err = c.Ping(t.Context(), PingOptions{}) + assert.NilError(t, err) + assert.Check(t, is.Equal(got, hdrVal)) + + assert.NilError(t, c.Close()) + }) + + t.Run("invalid hook", func(t *testing.T) { + _, err := New(WithResponseHook(nil)) + assert.Error(t, err, "invalid response hook: hook is nil") + }) + + t.Run("multiple hooks", func(t *testing.T) { + var triggered []string + + c, err := New( + WithResponseHook(func(*http.Response) error { + triggered = append(triggered, "hook 1: "+hdrVal) + return nil + }), + WithResponseHook(func(*http.Response) error { + triggered = append(triggered, "hook 2: "+hdrVal) + return nil + }), + WithBaseMockClient(func(req *http.Request) (*http.Response, error) { + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + } + resp.Header.Set(hdrKey, hdrVal) + return resp, nil + }), + ) + assert.NilError(t, err) + + _, err = c.Ping(t.Context(), PingOptions{}) + assert.NilError(t, err) + assert.Check(t, is.DeepEqual(triggered, []string{"hook 1: " + hdrVal, "hook 2: " + hdrVal})) + + assert.NilError(t, c.Close()) + }) + + t.Run("hook error", func(t *testing.T) { + closed := false + expError := errors.New("hook failed") + + c, err := New( + WithResponseHook(func(*http.Response) error { + return expError + }), + WithBaseMockClient(func(req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: &closeTracker{onClose: func() { closed = true }}, + }, nil + }), + ) + assert.NilError(t, err) + + _, err = c.Ping(t.Context(), PingOptions{}) + assert.Check(t, is.ErrorIs(err, expError)) + assert.Check(t, closed) + + assert.NilError(t, c.Close()) + }) +} + +type closeTracker struct { + onClose func() +} + +func (c *closeTracker) Read(p []byte) (int, error) { return 0, io.EOF } + +func (c *closeTracker) Close() error { + if c.onClose != nil { + c.onClose() + } + return nil +} diff --git a/client/client_responsehook.go b/client/client_responsehook.go new file mode 100644 index 0000000000..45f96605b8 --- /dev/null +++ b/client/client_responsehook.go @@ -0,0 +1,26 @@ +package client + +import ( + "net/http" +) + +type responseHookTransport struct { + base http.RoundTripper + hooks []ResponseHook +} + +func (t *responseHookTransport) RoundTrip(req *http.Request) (*http.Response, error) { + resp, err := t.base.RoundTrip(req) + if err != nil { + return resp, err + } + + for _, h := range t.hooks { + if err := h(resp); err != nil { + _ = resp.Body.Close() + return nil, err + } + } + + return resp, nil +} diff --git a/vendor/github.com/moby/moby/client/client.go b/vendor/github.com/moby/moby/client/client.go index 93d9f4f49a..d8e69953d5 100644 --- a/vendor/github.com/moby/moby/client/client.go +++ b/vendor/github.com/moby/moby/client/client.go @@ -59,6 +59,7 @@ import ( "net/http" "net/url" "path" + "slices" "strings" "sync" "sync/atomic" @@ -241,6 +242,13 @@ func New(ops ...Opt) (*Client, error) { c.client.Transport = otelhttp.NewTransport(c.client.Transport, c.traceOpts...) + if len(cfg.responseHooks) > 0 { + c.client.Transport = &responseHookTransport{ + base: c.client.Transport, + hooks: slices.Clone(cfg.responseHooks), + } + } + return c, nil } diff --git a/vendor/github.com/moby/moby/client/client_options.go b/vendor/github.com/moby/moby/client/client_options.go index ae80c4e5cf..c37fc9bceb 100644 --- a/vendor/github.com/moby/moby/client/client_options.go +++ b/vendor/github.com/moby/moby/client/client_options.go @@ -2,6 +2,7 @@ package client import ( "context" + "errors" "fmt" "net" "net/http" @@ -55,10 +56,19 @@ type clientConfig struct { // takes precedence. Either field disables API-version negotiation. envAPIVersion string + // responseHooks is a list of custom response hooks to call on responses. + responseHooks []ResponseHook + // traceOpts is a list of options to configure the tracing span. traceOpts []otelhttp.Option } +// ResponseHook is called for each HTTP response returned by the daemon. +// Hooks are invoked in the order they were added. +// +// Hooks must not read or close resp.Body. +type ResponseHook func(*http.Response) error + // Opt is a configuration option to initialize a [Client]. type Opt func(*clientConfig) error @@ -348,3 +358,18 @@ func WithTraceOptions(opts ...otelhttp.Option) Opt { return nil } } + +// WithResponseHook adds a ResponseHook to the client. ResponseHooks are called +// for each HTTP response returned by the daemon. Hooks are invoked in the order +// they were added. +// +// Hooks must not read or close resp.Body. +func WithResponseHook(h ResponseHook) Opt { + return func(c *clientConfig) error { + if h == nil { + return errors.New("invalid response hook: hook is nil") + } + c.responseHooks = append(c.responseHooks, h) + return nil + } +} diff --git a/vendor/github.com/moby/moby/client/client_responsehook.go b/vendor/github.com/moby/moby/client/client_responsehook.go new file mode 100644 index 0000000000..45f96605b8 --- /dev/null +++ b/vendor/github.com/moby/moby/client/client_responsehook.go @@ -0,0 +1,26 @@ +package client + +import ( + "net/http" +) + +type responseHookTransport struct { + base http.RoundTripper + hooks []ResponseHook +} + +func (t *responseHookTransport) RoundTrip(req *http.Request) (*http.Response, error) { + resp, err := t.base.RoundTrip(req) + if err != nil { + return resp, err + } + + for _, h := range t.hooks { + if err := h(resp); err != nil { + _ = resp.Body.Close() + return nil, err + } + } + + return resp, nil +}