client: implement WithResponseHook option

This options allows a client to call hooks on API responses, for
example, to get response header, or other information from the
response.

Signed-off-by: Sebastiaan van Stijn <github@gone.nl>

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
 }
 

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
+	}
+}

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
+}

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
+}

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
 }
 

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
+	}
+}

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
+}