diff --git a/pkg/client/client.go b/pkg/client/client.go index a43cd521..15f1e6e9 100644 --- a/pkg/client/client.go +++ b/pkg/client/client.go @@ -99,22 +99,23 @@ func New(opts ...Option) (*Client, error) { } } - // Instantiate the HTTP client - transport := &http.Transport{ - TLSClientConfig: client.tlsConfig, - } + if client.httpClient == nil { + transport := &http.Transport{ + TLSClientConfig: client.tlsConfig, + } - if client.dialer != nil { - transport.DialContext = client.dialer.DialContext - } + if client.dialer != nil { + transport.DialContext = client.dialer.DialContext + } - client.httpClient = &http.Client{ - // The default is zero, which means no timeout, which means that - // the requests may hang indefinitely. See [1] for more details. - // - // [1]: https://github.com/cirruslabs/orchard/issues/152#issuecomment-1927091747 - Timeout: defaultHTTPTimeout, - Transport: transport, + client.httpClient = &http.Client{ + // The default is zero, which means no timeout, which means that + // the requests may hang indefinitely. See [1] for more details. + // + // [1]: https://github.com/cirruslabs/orchard/issues/152#issuecomment-1927091747 + Timeout: defaultHTTPTimeout, + Transport: transport, + } } url, err := url.Parse(client.address) diff --git a/pkg/client/option.go b/pkg/client/option.go index 537999c2..6915b101 100644 --- a/pkg/client/option.go +++ b/pkg/client/option.go @@ -3,6 +3,7 @@ package client import ( "crypto/x509" "fmt" + "net/http" "strings" "github.com/cirruslabs/orchard/internal/dialer" @@ -11,6 +12,15 @@ import ( type Option func(*Client) +// WithHTTPClient uses httpClient for HTTP and WebSocket requests. +// The caller configures its timeout and transport, including TLS and dialing, +// and owns connection cleanup. Passing nil uses the default HTTP client. +func WithHTTPClient(httpClient *http.Client) Option { + return func(client *Client) { + client.httpClient = httpClient + } +} + func WithAddress(address string) Option { return func(client *Client) { client.address = address diff --git a/pkg/client/option_test.go b/pkg/client/option_test.go new file mode 100644 index 00000000..c65b0be9 --- /dev/null +++ b/pkg/client/option_test.go @@ -0,0 +1,122 @@ +package client_test + +import ( + "context" + "io" + "net" + "net/http" + "net/http/httptest" + "net/http/httptrace" + "sync/atomic" + "testing" + "time" + + "github.com/cirruslabs/orchard/pkg/client" + "github.com/coder/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const requestTimeout = 5 * time.Second + +func TestWithHTTPClientControlsRedirectsAndConnections(t *testing.T) { + var calls atomic.Int32 + server := httptest.NewTLSServer( + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + http.Redirect(w, r, "/redirected", http.StatusTemporaryRedirect) + }), + ) + t.Cleanup(server.Close) + + httpClient := server.Client() + httpClient.Timeout = requestTimeout + httpClient.CheckRedirect = func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse + } + t.Cleanup(httpClient.CloseIdleConnections) + + controller, err := client.New(client.WithAddress(server.URL), client.WithHTTPClient(httpClient)) + require.NoError(t, err) + + var connection net.Conn + ctx := httptrace.WithClientTrace( + t.Context(), + &httptrace.ClientTrace{ + GotConn: func(info httptrace.GotConnInfo) { + connection = info.Conn + }, + }, + ) + _, err = controller.Workers().Get(ctx, "worker") + var apiError *client.APIError + require.ErrorAs(t, err, &apiError) + require.Equal(t, http.StatusTemporaryRedirect, apiError.StatusCode) + require.EqualValues(t, 1, calls.Load()) + + require.NotNil(t, connection) + httpClient.CloseIdleConnections() + _, err = connection.Write([]byte("closed")) + require.ErrorIs(t, err, net.ErrClosed) +} + +func TestWithHTTPClientSupportsWebSockets(t *testing.T) { + server := httptest.NewTLSServer( + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + connection, err := websocket.Accept(w, r, nil) + if !assert.NoError(t, err) { + return + } + defer connection.CloseNow() + + messageType, payload, err := connection.Read(r.Context()) + if !assert.NoError(t, err) { + return + } + + assert.NoError(t, connection.Write(r.Context(), messageType, payload)) + }), + ) + t.Cleanup(server.Close) + + httpClient := server.Client() + t.Cleanup(httpClient.CloseIdleConnections) + + controller, err := client.New(client.WithAddress(server.URL), client.WithHTTPClient(httpClient)) + require.NoError(t, err) + + ctx, cancel := context.WithTimeout(t.Context(), requestTimeout) + defer cancel() + + connection, err := controller.Workers().PortForward(ctx, "worker", 22) + require.NoError(t, err) + defer connection.Close() + + _, err = connection.Write([]byte("hello")) + require.NoError(t, err) + + payload := make([]byte, len("hello")) + _, err = io.ReadFull(connection, payload) + require.NoError(t, err) + require.Equal(t, "hello", string(payload)) +} + +func TestWithHTTPClientNilKeepsDefaults(t *testing.T) { + server := httptest.NewTLSServer( + http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, `{"name":"worker"}`) + }), + ) + t.Cleanup(server.Close) + + controller, err := client.New( + client.WithAddress(server.URL), + client.WithHTTPClient(nil), + client.WithTrustedCertificate(server.Certificate()), + ) + require.NoError(t, err) + + worker, err := controller.Workers().Get(t.Context(), "worker") + require.NoError(t, err) + require.Equal(t, "worker", worker.Name) +}