Add a context parameter so we can enforce a timeout for the token exchange.

Signed-off-by: Matt Moyer <moyerm@vmware.com>

cmd/placeholder-name/main.go

@@ -6,10 +6,12 @@ SPDX-License-Identifier: Apache-2.0
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
 	"os"
+	"time"
 
 	"k8s.io/client-go/pkg/apis/clientauthentication"
 
@@ -18,7 +20,7 @@ import (
 )
 
 func main() {
-	err := run(os.LookupEnv, client.ExchangeToken, os.Stdout)
+	err := run(os.LookupEnv, client.ExchangeToken, os.Stdout, 30*time.Second)
 	if err != nil {
 		_, _ = fmt.Fprintf(os.Stderr, "%s", err.Error())
 		os.Exit(1)
@@ -26,11 +28,14 @@ func main() {
 }
 
 type envGetter func(string) (string, bool)
-type tokenExchanger func(token, caBundle, apiEndpoint string) (*clientauthentication.ExecCredential, error)
+type tokenExchanger func(ctx context.Context, token, caBundle, apiEndpoint string) (*clientauthentication.ExecCredential, error)
 
 const ErrMissingEnvVar = constable.Error("failed to login: environment variable not set")
 
-func run(envGetter envGetter, tokenExchanger tokenExchanger, outputWriter io.Writer) error {
+func run(envGetter envGetter, tokenExchanger tokenExchanger, outputWriter io.Writer, timeout time.Duration) error {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
 	token, varExists := envGetter("PLACEHOLDER_NAME_TOKEN")
 	if !varExists {
 		return envVarNotSetError("PLACEHOLDER_NAME_TOKEN")
@@ -46,7 +51,7 @@ func run(envGetter envGetter, tokenExchanger tokenExchanger, outputWriter io.Wri
 		return envVarNotSetError("PLACEHOLDER_NAME_K8S_API_ENDPOINT")
 	}
 
-	execCredential, err := tokenExchanger(token, caBundle, apiEndpoint)
+	execCredential, err := tokenExchanger(ctx, token, caBundle, apiEndpoint)
 	if err != nil {
 		return fmt.Errorf("failed to login: %w", err)
 	}

cmd/placeholder-name/main_test.go

@@ -7,8 +7,10 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"testing"
+	"time"
 
 	"github.com/sclevine/spec"
 	"github.com/sclevine/spec/report"
@@ -49,7 +51,7 @@ func TestRun(t *testing.T) {
 					"PLACEHOLDER_NAME_K8S_API_ENDPOINT": "a",
 					"PLACEHOLDER_NAME_CA_BUNDLE":        "b",
 				}
-				err := run(envGetter, tokenExchanger, buffer)
+				err := run(envGetter, tokenExchanger, buffer, 30*time.Second)
 				require.EqualError(t, err, "failed to login: environment variable not set: PLACEHOLDER_NAME_TOKEN")
 			})
 
@@ -58,7 +60,7 @@ func TestRun(t *testing.T) {
 					"PLACEHOLDER_NAME_K8S_API_ENDPOINT": "a",
 					"PLACEHOLDER_NAME_TOKEN":            "b",
 				}
-				err := run(envGetter, tokenExchanger, buffer)
+				err := run(envGetter, tokenExchanger, buffer, 30*time.Second)
 				require.EqualError(t, err, "failed to login: environment variable not set: PLACEHOLDER_NAME_CA_BUNDLE")
 			})
 
@@ -67,27 +69,27 @@ func TestRun(t *testing.T) {
 					"PLACEHOLDER_NAME_TOKEN":     "a",
 					"PLACEHOLDER_NAME_CA_BUNDLE": "b",
 				}
-				err := run(envGetter, tokenExchanger, buffer)
+				err := run(envGetter, tokenExchanger, buffer, 30*time.Second)
 				require.EqualError(t, err, "failed to login: environment variable not set: PLACEHOLDER_NAME_K8S_API_ENDPOINT")
 			})
 		}, spec.Parallel())
 
 		when("the token exchange fails", func() {
 			it.Before(func() {
-				tokenExchanger = func(token, caBundle, apiEndpoint string) (*clientauthentication.ExecCredential, error) {
+				tokenExchanger = func(ctx context.Context, token, caBundle, apiEndpoint string) (*clientauthentication.ExecCredential, error) {
 					return nil, fmt.Errorf("some error")
 				}
 			})
 
 			it("returns an error", func() {
-				err := run(envGetter, tokenExchanger, buffer)
+				err := run(envGetter, tokenExchanger, buffer, 30*time.Second)
 				require.EqualError(t, err, "failed to login: some error")
 			})
 		}, spec.Parallel())
 
 		when("the JSON encoder fails", func() {
 			it.Before(func() {
-				tokenExchanger = func(token, caBundle, apiEndpoint string) (*clientauthentication.ExecCredential, error) {
+				tokenExchanger = func(ctx context.Context, token, caBundle, apiEndpoint string) (*clientauthentication.ExecCredential, error) {
 					return &clientauthentication.ExecCredential{
 						Status: &clientauthentication.ExecCredentialStatus{Token: "some token"},
 					}, nil
@@ -95,16 +97,36 @@ func TestRun(t *testing.T) {
 			})
 
 			it("returns an error", func() {
-				err := run(envGetter, tokenExchanger, &errWriter{returnErr: fmt.Errorf("some IO error")})
+				err := run(envGetter, tokenExchanger, &errWriter{returnErr: fmt.Errorf("some IO error")}, 30*time.Second)
 				require.EqualError(t, err, "failed to marshal response to stdout: some IO error")
 			})
 		}, spec.Parallel())
 
+		when("the token exchange times out", func() {
+			it.Before(func() {
+				tokenExchanger = func(ctx context.Context, token, caBundle, apiEndpoint string) (*clientauthentication.ExecCredential, error) {
+					select {
+					case <-time.After(100 * time.Millisecond):
+						return &clientauthentication.ExecCredential{
+							Status: &clientauthentication.ExecCredentialStatus{Token: "some token"},
+						}, nil
+					case <-ctx.Done():
+						return nil, ctx.Err()
+					}
+				}
+			})
+
+			it("returns an error", func() {
+				err := run(envGetter, tokenExchanger, buffer, 1*time.Millisecond)
+				require.EqualError(t, err, "failed to login: context deadline exceeded")
+			})
+		}, spec.Parallel())
+
 		when("the token exchange succeeds", func() {
 			var actualToken, actualCaBundle, actualAPIEndpoint string
 
 			it.Before(func() {
-				tokenExchanger = func(token, caBundle, apiEndpoint string) (*clientauthentication.ExecCredential, error) {
+				tokenExchanger = func(ctx context.Context, token, caBundle, apiEndpoint string) (*clientauthentication.ExecCredential, error) {
 					actualToken, actualCaBundle, actualAPIEndpoint = token, caBundle, apiEndpoint
 					return &clientauthentication.ExecCredential{
 						Status: &clientauthentication.ExecCredentialStatus{Token: "some token"},
@@ -113,7 +135,7 @@ func TestRun(t *testing.T) {
 			})
 
 			it("writes the execCredential to the given writer", func() {
-				err := run(envGetter, tokenExchanger, buffer)
+				err := run(envGetter, tokenExchanger, buffer, 30*time.Second)
 				require.NoError(t, err)
 				require.Equal(t, fakeEnv["PLACEHOLDER_NAME_TOKEN"], actualToken)
 				require.Equal(t, fakeEnv["PLACEHOLDER_NAME_CA_BUNDLE"], actualCaBundle)

pkg/client/client.go

@@ -5,11 +5,13 @@ SPDX-License-Identifier: Apache-2.0
 
 package client
 
-import "k8s.io/client-go/pkg/apis/clientauthentication"
+import (
+	"context"
 
-func ExchangeToken(token, caBundle, apiEndpoint string) (*clientauthentication.ExecCredential, error) {
+	"k8s.io/client-go/pkg/apis/clientauthentication"
-	_ = token
+)
-	_ = caBundle
+
-	_ = apiEndpoint
+func ExchangeToken(ctx context.Context, token, caBundle, apiEndpoint string) (*clientauthentication.ExecCredential, error) {
+	_, _, _, _ = ctx, token, caBundle, apiEndpoint
 	return nil, nil
 }