Add support for fetching keys from a JWKS endpoint

hack/ark/test-e2e.sh

@@ -101,6 +101,7 @@ kubectl apply -f "${root_dir}/hack/ark/cluster-external-secret.yaml"
 
 # We use a non-existent tag and omit the `--version` flag, to work around a Helm
 # v4 bug. See: https://github.com/helm/helm/issues/31600
+# TODO: shouldn't need to set config.sendSecretValues because it will default to true in future
 helm upgrade agent "oci://${ARK_CHART}:NON_EXISTENT_TAG@${ARK_CHART_DIGEST}" \
      --install \
      --wait \
@@ -113,6 +114,7 @@ helm upgrade agent "oci://${ARK_CHART}:NON_EXISTENT_TAG@${ARK_CHART_DIGEST}" \
      --set config.clusterName="e2e-test-cluster" \
      --set config.clusterDescription="A temporary cluster for E2E testing. Contact @wallrj-cyberark." \
      --set config.period=60s \
+     --set config.sendSecretValues=true \
      --set-json "podLabels={\"disco-agent.cyberark.cloud/test-id\": \"${RANDOM}\"}"
 
 kubectl rollout status deployments/disco-agent --namespace "${NAMESPACE}"

internal/cyberark/client_test.go

@@ -32,9 +32,9 @@ func TestCyberArkClient_PutSnapshot_MockAPI(t *testing.T) {
 		Secret:    "somepassword",
 	}
 
-	discoveryClient := servicediscovery.New(httpClient)
+	discoveryClient := servicediscovery.New(httpClient, cfg.Subdomain)
 
-	serviceMap, tenantUUID, err := discoveryClient.DiscoverServices(t.Context(), cfg.Subdomain)
+	serviceMap, tenantUUID, err := discoveryClient.DiscoverServices(t.Context())
 	if err != nil {
 		t.Fatalf("failed to discover mock services: %v", err)
 	}
@@ -76,9 +76,9 @@ func TestCyberArkClient_PutSnapshot_RealAPI(t *testing.T) {
 	cfg, err := cyberark.LoadClientConfigFromEnvironment()
 	require.NoError(t, err)
 
-	discoveryClient := servicediscovery.New(httpClient)
+	discoveryClient := servicediscovery.New(httpClient, cfg.Subdomain)
 
-	serviceMap, tenantUUID, err := discoveryClient.DiscoverServices(t.Context(), cfg.Subdomain)
+	serviceMap, tenantUUID, err := discoveryClient.DiscoverServices(t.Context())
 	if err != nil {
 		t.Fatalf("failed to discover services: %v", err)
 	}

internal/cyberark/identity/cmd/testidentity/main.go

@@ -50,8 +50,8 @@ func run(ctx context.Context) error {
 	var rootCAs *x509.CertPool
 	httpClient := http_client.NewDefaultClient(version.UserAgent(), rootCAs)
 
-	sdClient := servicediscovery.New(httpClient)
-	services, _, err := sdClient.DiscoverServices(ctx, subdomain)
+	sdClient := servicediscovery.New(httpClient, subdomain)
+	services, _, err := sdClient.DiscoverServices(ctx)
 	if err != nil {
 		return fmt.Errorf("while performing service discovery: %s", err)
 	}

internal/cyberark/identity/identity.go

@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"net/url"
 	"sync"
+	"time"
 
 	"k8s.io/klog/v2"
 
@@ -180,6 +181,7 @@ type Client struct {
 
 	tokenCached      token
 	tokenCachedMutex sync.Mutex
+	tokenCachedTime  time.Time
 }
 
 // token is a wrapper type for holding auth tokens we want to cache.
@@ -205,12 +207,23 @@ func New(httpClient *http.Client, baseURL string, subdomain string) *Client {
 // Tokens are cached internally and are not directly accessible to code; use Client.AuthenticateRequest to add credentials
 // to an *http.Request.
 func (c *Client) LoginUsernamePassword(ctx context.Context, username string, password []byte) error {
+	// note: we hold the mutex for the whole login attempt to ensure that only one login attempt can be in flight at once,
+	// and to ensure that the token cache is correctly updated
+	c.tokenCachedMutex.Lock()
+	defer c.tokenCachedMutex.Unlock()
+
 	defer func() {
 		for i := range password {
 			password[i] = 0x00
 		}
 	}()
 
+	if time.Since(c.tokenCachedTime) < 15*time.Minute && c.tokenCached.Username == username {
+		// If the cached token is recent and for the same username, we can reuse it.
+		klog.FromContext(ctx).V(2).Info("reusing cached token for user", "username", username)
+		return nil
+	}
+
 	advanceRequestBody, err := c.doStartAuthentication(ctx, username)
 	if err != nil {
 		return err
@@ -405,15 +418,13 @@ func (c *Client) doAdvanceAuthentication(ctx context.Context, username string, p
 
 	klog.FromContext(ctx).Info("successfully completed AdvanceAuthentication request to CyberArk Identity; login complete", "username", username)
 
-	c.tokenCachedMutex.Lock()
-
+	// NB: This assumes we already hold the token cache mutex, which we do in LoginUsernamePassword, so this is safe.
+	c.tokenCachedTime = time.Now()
 	c.tokenCached = token{
 		Username: username,
 		Token:    advanceAuthResponse.Result.Token,
 	}
 
-	c.tokenCachedMutex.Unlock()
-
 	return nil
 }
 

internal/cyberark/identity/identity_test.go

@@ -53,7 +53,7 @@ func TestLoginUsernamePassword_RealAPI(t *testing.T) {
 	arktesting.SkipIfNoEnv(t)
 	subdomain := os.Getenv("ARK_SUBDOMAIN")
 	httpClient := http.DefaultClient
-	services, _, err := servicediscovery.New(httpClient).DiscoverServices(t.Context(), subdomain)
+	services, _, err := servicediscovery.New(httpClient, subdomain).DiscoverServices(t.Context())
 	require.NoError(t, err)
 
 	loginUsernamePasswordTests(t, func(t testing.TB) inputs {

internal/cyberark/servicediscovery/discovery.go

@@ -9,6 +9,8 @@ import (
 	"net/url"
 	"os"
 	"path"
+	"sync"
+	"time"
 
 	arkapi "github.com/jetstack/preflight/internal/cyberark/api"
 	"github.com/jetstack/preflight/pkg/version"
@@ -35,21 +37,34 @@ const (
 // users to fetch URLs for various APIs available in CyberArk. This client is specialised to
 // fetch only API endpoints, since only API endpoints are required by the Venafi Kubernetes Agent currently.
 type Client struct {
-	client  *http.Client
-	baseURL string
+	client    *http.Client
+	baseURL   string
+	subdomain string
+
+	cachedResponse      *Services
+	cachedTenantID      string
+	cachedResponseTime  time.Time
+	cachedResponseMutex sync.Mutex
 }
 
 // New creates a new CyberArk Service Discovery client. If the ARK_DISCOVERY_API
 // environment variable is set, it is used as the base URL for the service
 // discovery API. Otherwise, the production URL is used.
-func New(httpClient *http.Client) *Client {
+func New(httpClient *http.Client, subdomain string) *Client {
 	baseURL := os.Getenv("ARK_DISCOVERY_API")
 	if baseURL == "" {
 		baseURL = ProdDiscoveryAPIBaseURL
 	}
+
 	client := &Client{
-		client:  httpClient,
-		baseURL: baseURL,
+		client:    httpClient,
+		baseURL:   baseURL,
+		subdomain: subdomain,
+
+		cachedResponse:      nil,
+		cachedTenantID:      "",
+		cachedResponseTime:  time.Time{},
+		cachedResponseMutex: sync.Mutex{},
 	}
 
 	return client
@@ -93,17 +108,24 @@ type Services struct {
 	DiscoveryContext ServiceEndpoint
 }
 
-// DiscoverServices fetches from the service discovery service for a given subdomain
+// DiscoverServices fetches from the service discovery service for the configured subdomain
 // and parses the CyberArk Identity API URL and Inventory API URL.
 // It also returns the Tenant ID UUID corresponding to the subdomain.
-func (c *Client) DiscoverServices(ctx context.Context, subdomain string) (*Services, string, error) {
+func (c *Client) DiscoverServices(ctx context.Context) (*Services, string, error) {
+	c.cachedResponseMutex.Lock()
+	defer c.cachedResponseMutex.Unlock()
+
+	if c.cachedResponse != nil && time.Since(c.cachedResponseTime) < 1*time.Hour {
+		return c.cachedResponse, c.cachedTenantID, nil
+	}
+
 	u, err := url.Parse(c.baseURL)
 	if err != nil {
 		return nil, "", fmt.Errorf("invalid base URL for service discovery: %w", err)
 	}
 
 	u.Path = path.Join(u.Path, "api/public/tenant-discovery")
-	u.RawQuery = url.Values{"bySubdomain": []string{subdomain}}.Encode()
+	u.RawQuery = url.Values{"bySubdomain": []string{c.subdomain}}.Encode()
 
 	endpoint := u.String()
 
@@ -127,7 +149,7 @@ func (c *Client) DiscoverServices(ctx context.Context, subdomain string) (*Servi
 		// a 404 error is returned with an empty JSON body "{}" if the subdomain is unknown; at the time of writing, we haven't observed
 		// any other errors and so we can't special case them
 		if resp.StatusCode == http.StatusNotFound {
-			return nil, "", fmt.Errorf("got an HTTP 404 response from service discovery; maybe the subdomain %q is incorrect or does not exist?", subdomain)
+			return nil, "", fmt.Errorf("got an HTTP 404 response from service discovery; maybe the subdomain %q is incorrect or does not exist?", c.subdomain)
 		}
 
 		return nil, "", fmt.Errorf("got unexpected status code %s from request to service discovery API", resp.Status)
@@ -167,8 +189,14 @@ func (c *Client) DiscoverServices(ctx context.Context, subdomain string) (*Servi
 	}
 	//TODO: Should add a check for discoveryContextAPI too?
 
-	return &Services{
+	services := &Services{
 		Identity:         ServiceEndpoint{API: identityAPI},
 		DiscoveryContext: ServiceEndpoint{API: discoveryContextAPI},
-	}, discoveryResp.TenantID, nil
+	}
+
+	c.cachedResponse = services
+	c.cachedTenantID = discoveryResp.TenantID
+	c.cachedResponseTime = time.Now()
+
+	return services, discoveryResp.TenantID, nil
 }

internal/cyberark/servicediscovery/discovery_test.go

@@ -64,9 +64,9 @@ func Test_DiscoverIdentityAPIURL(t *testing.T) {
 				},
 			})
 
-			client := New(httpClient)
+			client := New(httpClient, testSpec.subdomain)
 
-			services, _, err := client.DiscoverServices(ctx, testSpec.subdomain)
+			services, _, err := client.DiscoverServices(ctx)
 			if testSpec.expectedError != nil {
 				assert.EqualError(t, err, testSpec.expectedError.Error())
 				assert.Nil(t, services)