Add support for oauth2.TokenSource in GCP KMS

gcpkms/keysource.go

@@ -12,6 +12,7 @@ import (
 	kms "cloud.google.com/go/kms/apiv1"
 	"cloud.google.com/go/kms/apiv1/kmspb"
 	"github.com/sirupsen/logrus"
+	"golang.org/x/oauth2"
 	"google.golang.org/api/option"
 	"google.golang.org/grpc"
 
@@ -50,6 +51,11 @@ type MasterKey struct {
 	// for NeedsRotation.
 	CreationDate time.Time
 
+	// tokenSource contains the oauth2.TokenSource used by the GCP client.
+	// It can be injected by a (local) keyservice.KeyServiceServer using
+	// TokenSource.ApplyToMasterKey.
+	// If nil, the remaining authentication methods are attempted.
+	tokenSource oauth2.TokenSource
 	// credentialJSON is the Service Account credentials JSON used for
 	// authenticating towards the GCP KMS service.
 	credentialJSON []byte
@@ -82,6 +88,22 @@ func MasterKeysFromResourceIDString(resourceID string) []*MasterKey {
 	return keys
 }
 
+// TokenSource is an oauth2.TokenSource used for authenticating towards the
+// GCP KMS service.
+type TokenSource struct {
+	source oauth2.TokenSource
+}
+
+// NewTokenSource creates a new TokenSource from the provided oauth2.TokenSource.
+func NewTokenSource(source oauth2.TokenSource) TokenSource {
+	return TokenSource{source: source}
+}
+
+// ApplyToMasterKey configures the TokenSource on the provided key.
+func (t TokenSource) ApplyToMasterKey(key *MasterKey) {
+	key.tokenSource = t.source
+}
+
 // CredentialJSON is the Service Account credentials JSON used for authenticating
 // towards the GCP KMS service.
 type CredentialJSON []byte
@@ -203,8 +225,8 @@ func (key *MasterKey) TypeToIdentifier() string {
 	return KeyTypeIdentifier
 }
 
-// newKMSClient returns a GCP KMS client configured with the credentialJSON
-// and/or grpcConn, falling back to environmental defaults.
+// newKMSClient returns a GCP KMS client configured with the tokenSource
+// or credentialJSON, and/or grpcConn, falling back to environmental defaults.
 // It returns an error if the ResourceID is invalid, or if the setup of the
 // client fails.
 func (key *MasterKey) newKMSClient() (*kms.KeyManagementClient, error) {
@@ -216,6 +238,8 @@ func (key *MasterKey) newKMSClient() (*kms.KeyManagementClient, error) {
 
 	var opts []option.ClientOption
 	switch {
+	case key.tokenSource != nil:
+		opts = append(opts, option.WithTokenSource(key.tokenSource))
 	case key.credentialJSON != nil:
 		opts = append(opts, option.WithCredentialsJSON(key.credentialJSON))
 	default:

gcpkms/keysource_test.go

@@ -9,6 +9,7 @@ import (
 
 	"cloud.google.com/go/kms/apiv1/kmspb"
 	"github.com/stretchr/testify/assert"
+	"golang.org/x/oauth2"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 )
@@ -38,6 +39,13 @@ func TestMasterKeysFromResourceIDString(t *testing.T) {
 	}
 }
 
+func TestTokenSource_ApplyToMasterKey(t *testing.T) {
+	src := NewTokenSource(oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "some-token"}))
+	key := &MasterKey{}
+	src.ApplyToMasterKey(key)
+	assert.Equal(t, src.source, key.tokenSource)
+}
+
 func TestCredentialJSON_ApplyToMasterKey(t *testing.T) {
 	key := &MasterKey{}
 	credential := CredentialJSON("mock")

go.mod

@@ -36,6 +36,7 @@ require (
 	github.com/urfave/cli v1.22.16
 	golang.org/x/crypto v0.36.0
 	golang.org/x/net v0.37.0
+	golang.org/x/oauth2 v0.28.0
 	golang.org/x/sys v0.31.0
 	golang.org/x/term v0.30.0
 	google.golang.org/api v0.227.0
@@ -138,7 +139,6 @@ require (
 	go.opentelemetry.io/otel/sdk v1.34.0 // indirect
 	go.opentelemetry.io/otel/sdk/metric v1.34.0 // indirect
 	go.opentelemetry.io/otel/trace v1.34.0 // indirect
-	golang.org/x/oauth2 v0.28.0 // indirect
 	golang.org/x/sync v0.12.0 // indirect
 	golang.org/x/text v0.23.0 // indirect
 	golang.org/x/time v0.11.0 // indirect