From d8993aca7c4d3d864f5d84ee82b19353016d3c35 Mon Sep 17 00:00:00 2001
From: Mariano Cano <mariano.cano@gmail.com>
Date: Tue, 25 Feb 2025 11:24:14 -0800
Subject: [PATCH] Add option to specify a client timeout

This commit adds to the ca Client a new option to specify the client
timeout.

Fixes #2176
---
 ca/adminClient.go |  2 +-
 ca/client.go      | 21 ++++++++++++++++++---
 ca/client_test.go | 28 ++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/ca/adminClient.go b/ca/adminClient.go
index 464df92f..f96bc6ac 100644
--- a/ca/adminClient.go
+++ b/ca/adminClient.go
@@ -77,7 +77,7 @@ func NewAdminClient(endpoint string, opts ...ClientOption) (*AdminClient, error)
 	}
 
 	return &AdminClient{
-		client:      newClient(tr),
+		client:      newClient(tr, o.timeout),
 		endpoint:    u,
 		retryFunc:   o.retryFunc,
 		opts:        opts,
diff --git a/ca/client.go b/ca/client.go
index 165eac2c..a44445bc 100644
--- a/ca/client.go
+++ b/ca/client.go
@@ -22,6 +22,7 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/pkg/errors"
 	"golang.org/x/net/http2"
@@ -53,10 +54,11 @@ type uaClient struct {
 	Client *http.Client
 }
 
-func newClient(transport http.RoundTripper) *uaClient {
+func newClient(transport http.RoundTripper, timeout time.Duration) *uaClient {
 	return &uaClient{
 		Client: &http.Client{
 			Transport: transport,
+			Timeout:   timeout,
 		},
 	}
 }
@@ -149,6 +151,7 @@ type ClientOption func(o *clientOptions) error
 
 type clientOptions struct {
 	transport            http.RoundTripper
+	timeout              time.Duration
 	rootSHA256           string
 	rootFilename         string
 	rootBundle           []byte
@@ -388,6 +391,16 @@ func WithRetryFunc(fn RetryFunc) ClientOption {
 	}
 }
 
+// WithTimeout defines the time limit for requests made by this client. The
+// timeout includes connection time, any redirects, and reading the response
+// body.
+func WithTimeout(d time.Duration) ClientOption {
+	return func(o *clientOptions) error {
+		o.timeout = d
+		return nil
+	}
+}
+
 func getTransportFromFile(filename string) (http.RoundTripper, error) {
 	data, err := os.ReadFile(filename)
 	if err != nil {
@@ -548,6 +561,7 @@ type Client struct {
 	client    *uaClient
 	endpoint  *url.URL
 	retryFunc RetryFunc
+	timeout   time.Duration
 	opts      []ClientOption
 }
 
@@ -568,9 +582,10 @@ func NewClient(endpoint string, opts ...ClientOption) (*Client, error) {
 	}
 
 	return &Client{
-		client:    newClient(tr),
+		client:    newClient(tr, o.timeout),
 		endpoint:  u,
 		retryFunc: o.retryFunc,
+		timeout:   o.timeout,
 		opts:      opts,
 	}, nil
 }
@@ -890,7 +905,7 @@ func (c *Client) RevokeWithContext(ctx context.Context, req *api.RevokeRequest,
 	var uaClient *uaClient
 retry:
 	if tr != nil {
-		uaClient = newClient(tr)
+		uaClient = newClient(tr, c.timeout)
 	} else {
 		uaClient = c.client
 	}
diff --git a/ca/client_test.go b/ca/client_test.go
index bd05614b..e3b7283d 100644
--- a/ca/client_test.go
+++ b/ca/client_test.go
@@ -1017,6 +1017,34 @@ func TestClient_GetCaURL(t *testing.T) {
 	}
 }
 
+func TestClient_WithTimeout(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(100 * time.Millisecond)
+		render.JSONStatus(w, r, api.HealthResponse{Status: "ok"}, 200)
+	}))
+	defer srv.Close()
+
+	tests := []struct {
+		name      string
+		options   []ClientOption
+		assertion assert.ErrorAssertionFunc
+	}{
+		{"ok", []ClientOption{WithTransport(http.DefaultTransport)}, assert.NoError},
+		{"ok with timeout", []ClientOption{WithTransport(http.DefaultTransport), WithTimeout(time.Second)}, assert.NoError},
+		{"fail with timeout", []ClientOption{WithTransport(http.DefaultTransport), WithTimeout(100 * time.Millisecond)}, assert.Error},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c, err := NewClient(srv.URL, tt.options...)
+			require.NoError(t, err)
+			_, err = c.Health()
+			tt.assertion(t, err)
+		})
+	}
+
+}
+
 func Test_enforceRequestID(t *testing.T) {
 	set := httptest.NewRequest(http.MethodGet, "https://example.com", http.NoBody)
 	set.Header.Set("X-Request-Id", "already-set")