GCP CLI support: server-side request handler

api/utils/gcp/endpoint.go

@@ -0,0 +1,29 @@
+// Copyright 2022 Gravitational, Inc
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package gcp
+
+import (
+	"strings"
+)
+
+const (
+	// gcpEndpointSuffix is hostname suffix used to determine if hostname is a GCP endpoint
+	gcpEndpointSuffix = ".googleapis.com"
+)
+
+// IsGCPEndpoint returns true if hostname is a GCP endpoint
+func IsGCPEndpoint(hostname string) bool {
+	return strings.HasSuffix(hostname, gcpEndpointSuffix)
+}

api/utils/gcp/endpoint_test.go

@@ -0,0 +1,60 @@
+// Copyright 2022 Gravitational, Inc
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package gcp
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestIsGCPEndpoint(t *testing.T) {
+	tests := []struct {
+		name     string
+		hostname string
+		want     bool
+	}{
+		{
+			name:     "compute googleapis",
+			hostname: "compute.googleapis.com",
+			want:     true,
+		},
+		{
+			name:     "top level googleapis",
+			hostname: "googleapis.com",
+			want:     false,
+		},
+		{
+			name:     "localhost",
+			hostname: "localhost",
+			want:     false,
+		},
+		{
+			name:     "fake googleapis",
+			hostname: "compute.googleapis.com.fake.com",
+			want:     false,
+		},
+		{
+			name:     "top level-like fake googleapis",
+			hostname: "googleapis.com.fake.com",
+			want:     false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require.Equal(t, tt.want, IsGCPEndpoint(tt.hostname))
+		})
+	}
+}

lib/srv/app/azure/handler.go

@@ -89,8 +89,13 @@ type handler struct {
 	tokenCache *utils.FnCache
 }
 
-// NewAzureHandler creates a new instance of an http.Handler for Azure requests.
+// NewAzureHandler creates a new instance of a http.Handler for Azure requests.
 func NewAzureHandler(ctx context.Context, config HandlerConfig) (http.Handler, error) {
+	return newAzureHandler(ctx, config)
+}
+
+// newAzureHandler creates a new instance of a handler for Azure requests. Used by NewAzureHandler and in tests.
+func newAzureHandler(ctx context.Context, config HandlerConfig) (*handler, error) {
 	if err := config.CheckAndSetDefaults(); err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -282,10 +287,23 @@ const getTokenTimeout = time.Second * 5
 func (s *handler) getToken(ctx context.Context, managedIdentity string, scope string) (*azcore.AccessToken, error) {
 	key := cacheKey{managedIdentity, scope}
 
-	timeoutCtx, cancel := context.WithTimeout(ctx, getTokenTimeout)
+	cancelCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	return utils.FnCacheGet(timeoutCtx, s.tokenCache, key, func(ctx context.Context) (*azcore.AccessToken, error) {
-		return s.getAccessToken(ctx, managedIdentity, scope)
-	})
+	var tokenResult *azcore.AccessToken
+	var errorResult error
+
+	go func() {
+		tokenResult, errorResult = utils.FnCacheGet(cancelCtx, s.tokenCache, key, func(ctx context.Context) (*azcore.AccessToken, error) {
+			return s.getAccessToken(ctx, managedIdentity, scope)
+		})
+		cancel()
+	}()
+
+	select {
+	case <-s.Clock.After(getTokenTimeout):
+		return nil, trace.Wrap(context.DeadlineExceeded, "timeout waiting for access token for %v", getTokenTimeout)
+	case <-cancelCtx.Done():
+		return tokenResult, errorResult
+	}
 }

lib/srv/app/azure/handler_test.go

@@ -0,0 +1,148 @@
+// Copyright 2022 Gravitational, Inc
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package azure
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/gravitational/trace"
+	"github.com/jonboulle/clockwork"
+	"github.com/stretchr/testify/require"
+)
+
+func TestForwarder_getToken(t *testing.T) {
+	t.Parallel()
+
+	type testCase struct {
+		name string
+
+		config HandlerConfig
+
+		managedIdentity string
+		scope           string
+
+		wantToken *azcore.AccessToken
+		checkErr  require.ErrorAssertionFunc
+	}
+
+	var tests []testCase
+
+	tests = []testCase{
+		{
+			name: "base case",
+			config: HandlerConfig{
+				getAccessToken: func(ctx context.Context, managedIdentity string, scope string) (*azcore.AccessToken, error) {
+					if managedIdentity != "MY_IDENTITY" {
+						return nil, trace.BadParameter("wrong managedIdentity")
+					}
+					if scope != "MY_SCOPE" {
+						return nil, trace.BadParameter("wrong scope")
+					}
+					return &azcore.AccessToken{Token: "foobar"}, nil
+				},
+			},
+			managedIdentity: "MY_IDENTITY",
+			scope:           "MY_SCOPE",
+			wantToken:       &azcore.AccessToken{Token: "foobar"},
+			checkErr:        require.NoError,
+		},
+		{
+			name: "timeout",
+			config: HandlerConfig{
+				Clock: clockwork.NewFakeClock(),
+				getAccessToken: func(ctx context.Context, managedIdentity string, scope string) (*azcore.AccessToken, error) {
+					// find the fake clock from above
+					var clock clockwork.FakeClock
+					for _, test := range tests {
+						if test.name == "timeout" {
+							clock = test.config.Clock.(clockwork.FakeClock)
+						}
+					}
+
+					clock.Advance(getTokenTimeout)
+					clock.Sleep(getTokenTimeout * 2)
+					return &azcore.AccessToken{Token: "foobar"}, nil
+				},
+			},
+			checkErr: func(t require.TestingT, err error, i ...interface{}) {
+				require.ErrorContains(t, err, "timeout waiting for access token for 5s")
+				require.ErrorIs(t, err, context.DeadlineExceeded)
+			},
+		},
+		{
+			name: "non-timeout error",
+			config: HandlerConfig{
+				getAccessToken: func(ctx context.Context, managedIdentity string, scope string) (*azcore.AccessToken, error) {
+					return nil, trace.BadParameter("bad param foo")
+				},
+			},
+			checkErr: func(t require.TestingT, err error, i ...interface{}) {
+				require.ErrorContains(t, err, "bad param foo")
+				require.True(t, trace.IsBadParameter(err))
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := context.Background()
+
+			hnd, err := newAzureHandler(ctx, tt.config)
+			require.NoError(t, err)
+
+			token, err := hnd.getToken(ctx, tt.managedIdentity, tt.scope)
+
+			require.Equal(t, tt.wantToken, token)
+			tt.checkErr(t, err)
+		})
+	}
+}
+
+func TestForwarder_getToken_cache(t *testing.T) {
+	ctx := context.Background()
+
+	clock := clockwork.NewFakeClock()
+
+	calls := 0
+	hnd, err := newAzureHandler(ctx, HandlerConfig{
+		Clock: clock,
+		getAccessToken: func(ctx context.Context, managedIdentity string, scope string) (*azcore.AccessToken, error) {
+			calls++
+			return &azcore.AccessToken{Token: "OK"}, nil
+		},
+	})
+	require.NoError(t, err)
+
+	// first call goes through
+	_, err = hnd.getToken(ctx, "", "")
+	require.NoError(t, err)
+	require.Equal(t, 1, calls)
+
+	// second call is cached
+	_, err = hnd.getToken(ctx, "", "")
+	require.NoError(t, err)
+	require.Equal(t, 1, calls)
+
+	// advance past cache expiry
+	clock.Advance(time.Second * 60 * 2)
+
+	// third call goes through
+	_, err = hnd.getToken(ctx, "", "")
+	require.NoError(t, err)
+	require.Equal(t, 2, calls)
+}