feat: enable auth for local requests to Corrosion API (resolves #110)

internal/corrosion/client.go

@@ -33,39 +33,44 @@ type APIClient struct {
 	newResubBackoff func() backoff.BackOff
 }
 
-// NewAPIClient creates a new Corrosion API client. The client retries on network errors using an exponential backoff
-// policy with a maximum interval of 1 second and a maximum elapsed time of 10 seconds.
+// NewAPIClient creates a new Corrosion API client. The bearerToken is sent in the Authorization header of every
+// request to authenticate against Corrosion API.
+// The client retries on network errors using an exponential backoff policy with a maximum interval of 1 second and
+// a maximum elapsed time of 10 seconds.
 // It automatically resubscribes to active subscriptions if an error occurs using an exponential backoff policy with a
 // maximum interval of 1 second and a maximum elapsed time of 60 seconds.
 // Use the WithHTTP2Client option to provide a custom HTTP client and the WithResubscribeBackoff option to change the
 // backoff policy for resubscribing to a query.
-func NewAPIClient(addr netip.AddrPort, opts ...APIClientOption) (*APIClient, error) {
+func NewAPIClient(addr netip.AddrPort, bearerToken string, opts ...APIClientOption) (*APIClient, error) {
 	baseURL, err := url.Parse(fmt.Sprintf("http://%s", addr))
 	if err != nil {
 		return nil, fmt.Errorf("invalid URL: %w", err)
 	}
-	c := &APIClient{
-		baseURL: baseURL,
-		client: &http.Client{
-			Transport: &RetryRoundTripper{
-				Base: &http2.Transport{
-					AllowHTTP: true,
-					DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
-						dialer := &net.Dialer{
-							Timeout: http2ConnectTimeout,
-						}
-						return dialer.DialContext(ctx, network, addr)
-					},
-				},
-				NewBackoff: func() backoff.BackOff {
-					return backoff.NewExponentialBackOff(
-						backoff.WithInitialInterval(100*time.Millisecond),
-						backoff.WithMaxInterval(1*time.Second),
-						backoff.WithMaxElapsedTime(http2MaxRetryTime),
-					)
+
+	transport := &AuthRoundTripper{
+		Base: &RetryRoundTripper{
+			Base: &http2.Transport{
+				AllowHTTP: true,
+				DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
+					dialer := &net.Dialer{
+						Timeout: http2ConnectTimeout,
+					}
+					return dialer.DialContext(ctx, network, addr)
 				},
 			},
+			NewBackoff: func() backoff.BackOff {
+				return backoff.NewExponentialBackOff(
+					backoff.WithInitialInterval(100*time.Millisecond),
+					backoff.WithMaxInterval(1*time.Second),
+					backoff.WithMaxElapsedTime(http2MaxRetryTime),
+				)
+			},
 		},
+		Token: bearerToken,
+	}
+	c := &APIClient{
+		baseURL: baseURL,
+		client:  &http.Client{Transport: transport},
 		newResubBackoff: func() backoff.BackOff {
 			return backoff.NewExponentialBackOff(
 				backoff.WithInitialInterval(100*time.Millisecond),
@@ -83,6 +88,8 @@ func NewAPIClient(addr netip.AddrPort, opts ...APIClientOption) (*APIClient, err
 
 type APIClientOption func(*APIClient)
 
+// WithHTTP2Client replaces the client's HTTP transport. The provided client bypasses the built-in bearer-token
+// injection, so the caller is responsible for setting the Authorization header.
 func WithHTTP2Client(client *http.Client) APIClientOption {
 	return func(c *APIClient) {
 		c.client = client
@@ -97,6 +104,22 @@ func WithResubscribeBackoff(newBackoff func() backoff.BackOff) APIClientOption {
 	}
 }
 
+// AuthRoundTripper sets the Authorization header on every outgoing request. An empty Token leaves the header untouched.
+type AuthRoundTripper struct {
+	Base  http.RoundTripper
+	Token string
+}
+
+func (rt *AuthRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	if rt.Token == "" {
+		return rt.Base.RoundTrip(req)
+	}
+	// RoundTripper contract: must not mutate the caller's request.
+	req = req.Clone(req.Context())
+	req.Header.Set("Authorization", "Bearer "+rt.Token)
+	return rt.Base.RoundTrip(req)
+}
+
 type RetryRoundTripper struct {
 	Base http.RoundTripper
 	// NewBackoff creates a new backoff policy for each request.
@@ -107,8 +130,7 @@ func (rt *RetryRoundTripper) RoundTrip(req *http.Request) (*http.Response, error
 	roundTrip := func() (*http.Response, error) {
 		resp, err := rt.Base.RoundTrip(req)
 		if err != nil {
-			var opErr *net.OpError
-			if errors.As(err, &opErr) {
+			if _, ok := errors.AsType[*net.OpError](err); ok {
 				// Not certain, but I expect operational errors should generally be retryable.
 				slog.Debug("Retrying corrosion API request due to network error.", "error", err)
 				return nil, err

internal/corrosion/client_test.go

@@ -0,0 +1,60 @@
+package corrosion
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestAuthRoundTripper_SetsAuthorizationHeader(t *testing.T) {
+	t.Parallel()
+
+	const token = "test-token-1234567890"
+
+	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()
+
+	rt := &AuthRoundTripper{Base: http.DefaultTransport, Token: token}
+	client := &http.Client{Transport: rt}
+
+	req, err := http.NewRequest(http.MethodGet, srv.URL, nil)
+	require.NoError(t, err)
+
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+	resp.Body.Close()
+
+	assert.Equal(t, "Bearer "+token, gotAuth)
+	// Caller's request must not be mutated by the RoundTripper.
+	assert.Empty(t, req.Header.Get("Authorization"))
+}
+
+func TestAuthRoundTripper_EmptyTokenSkipsHeader(t *testing.T) {
+	t.Parallel()
+
+	var headerSeen bool
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, headerSeen = r.Header["Authorization"]
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer srv.Close()
+
+	rt := &AuthRoundTripper{Base: http.DefaultTransport, Token: ""}
+	client := &http.Client{Transport: rt}
+
+	req, err := http.NewRequest(http.MethodGet, srv.URL, nil)
+	require.NoError(t, err)
+
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+	resp.Body.Close()
+
+	assert.False(t, headerSeen, "Authorization header should not be sent when token is empty")
+}