Feature/add metadata credentials provider

go.mod

@@ -30,12 +30,14 @@ require (
 	github.com/go-openapi/swag v0.22.3 // indirect
 	github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/golang-jwt/jwt/v4 v4.4.3 // indirect
 	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/google/gnostic-models v0.6.8 // indirect
 	github.com/google/gofuzz v1.2.0 // indirect
 	github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1 // indirect
 	github.com/imdario/mergo v0.3.6 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
+	github.com/jonboulle/clockwork v0.3.0 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
@@ -44,14 +46,21 @@ require (
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
+	github.com/yandex-cloud/go-genproto v0.0.0-20211115083454-9ca41db5ed9e // indirect
+	github.com/ydb-platform/ydb-go-sdk/v3 v3.47.3 // indirect
+	github.com/ydb-platform/ydb-go-yc v0.12.1 // indirect
+	github.com/ydb-platform/ydb-go-yc-metadata v0.6.1 // indirect
 	golang.org/x/net v0.21.0 // indirect
 	golang.org/x/oauth2 v0.16.0 // indirect
+	golang.org/x/sync v0.6.0 // indirect
 	golang.org/x/sys v0.18.0 // indirect
 	golang.org/x/term v0.17.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
 	golang.org/x/time v0.3.0 // indirect
 	golang.org/x/tools v0.17.0 // indirect
 	google.golang.org/appengine v1.6.8 // indirect
+	google.golang.org/genproto v0.0.0-20240123012728-ef4313101c80 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20240123012728-ef4313101c80 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240228224816-df926f6c8641 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect

pkg/client/auth/credentials/iam_creds.go

@@ -0,0 +1,62 @@
+package credentials
+
+import (
+	"context"
+	"strings"
+	"sync"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/credentials"
+	yc "github.com/ydb-platform/ydb-go-yc"
+	"google.golang.org/grpc/metadata"
+)
+
+type iamCredsProvider struct {
+	keyfileName string
+	iamEndpoint string
+
+	once    sync.Once
+	creds   credentials.Credentials
+	initErr error
+}
+
+// ContextWithAuth implements Provider.
+func (i *iamCredsProvider) ContextWithAuth(ctx context.Context) (context.Context, context.CancelFunc) {
+	// TODO(shmel1k@): add error handling
+	tok, _ := i.GetToken()
+	ctx, cf := context.WithCancel(ctx)
+	return metadata.AppendToOutgoingContext(ctx,
+		"x-ydb-auth-ticket", tok), cf
+}
+
+// ContextWithoutAuth implements Provider.
+func (i *iamCredsProvider) ContextWithoutAuth(ctx context.Context) (context.Context, context.CancelFunc) {
+	// TODO(shmel1k@): remove ContextWithoutAuth method
+	return context.WithCancel(ctx)
+}
+
+// GetToken implements Provider.
+func (i *iamCredsProvider) GetToken() (string, error) {
+	return i.creds.Token(context.TODO())
+}
+
+// Init implements Provider.
+func (i *iamCredsProvider) Init() error {
+	i.once.Do(func() {
+		sp := strings.Split(i.iamEndpoint, ":")
+		if len(sp) == 1 {
+			i.iamEndpoint += ":443"
+		}
+		i.creds, i.initErr = yc.NewClient(
+			yc.WithServiceFile(i.keyfileName),
+			yc.WithEndpoint(i.iamEndpoint),
+		)
+	})
+	return i.initErr
+}
+
+func NewIamCreds(keyfileName, iamEndpoint string) Provider {
+	return &iamCredsProvider{
+		keyfileName: keyfileName,
+		iamEndpoint: iamEndpoint,
+	}
+}

pkg/client/auth/credentials/metadata.go

@@ -0,0 +1,45 @@
+package credentials
+
+import (
+	"context"
+	"sync"
+
+	yc "github.com/ydb-platform/ydb-go-yc-metadata"
+	"go.uber.org/zap"
+	"google.golang.org/grpc/metadata"
+)
+
+type metadataProvider struct {
+	once  sync.Once
+	creds *yc.InstanceServiceAccountCredentials
+}
+
+// ContextWithAuth implements Provider.
+func (m *metadataProvider) ContextWithAuth(ctx context.Context) (context.Context, context.CancelFunc) {
+	token, _ := m.GetToken() // TODO(shmel1k@): add error handling
+	ctx1, cf := context.WithCancel(ctx)
+	return metadata.AppendToOutgoingContext(ctx1, "x-ydb-auth-ticket", token), cf
+}
+
+// ContextWithoutAuth implements Provider.
+func (m *metadataProvider) ContextWithoutAuth(ctx context.Context) (context.Context, context.CancelFunc) {
+	return context.WithCancel(ctx)
+}
+
+// GetToken implements Provider.
+func (m *metadataProvider) GetToken() (string, error) {
+	// TODO(shmel1k@): add contexts.
+	return m.creds.Token(context.TODO())
+}
+
+// Init implements Provider.
+func (m *metadataProvider) Init() error {
+	m.once.Do(func() {
+		m.creds = yc.NewInstanceServiceAccount(yc.WithURL("http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/token"))
+	})
+	return nil
+}
+
+func NewMetadata(logger *zap.SugaredLogger) Provider {
+	return &metadataProvider{}
+}

pkg/client/auth/credentials/provider.go

@@ -74,14 +74,18 @@ func (b *baseProvider) Init() error {
 		case options.IamToken:
 			b.impl = NewIamToken(b.opts.Auth.Creds.(*options.AuthIAMToken).Token)
 		case options.IamCreds:
-			b.initErr = fmt.Errorf("TODO: IAM authorization from SA key not implemented yet")
+			creds := b.opts.Auth.Creds.(*options.AuthIAMCreds)
+			b.impl = NewIamCreds(creds.KeyFilename, creds.Endpoint)
+		case options.IamMetadata:
+			b.impl = NewMetadata(b.logger)
 		case options.None:
 			b.initErr = fmt.Errorf("determined credentials to be anonymous. Anonymous credentials are currently unsupported")
 		default:
 			b.initErr = fmt.Errorf(
 				"internal error: authorization type not recognized after options validation, this should never happen",
 			)
 		}
+		b.initErr = b.impl.Init()
 	})
 	return b.initErr
 }

pkg/client/auth/credentials/static.go

@@ -41,8 +41,6 @@ func (s *staticCredentialsProvider) Init() error {
 // GetToken implements Provider.
 func (s *staticCredentialsProvider) GetToken() (string, error) {
 	// TODO(shmel1k@): probably, token can change time to time.
-	_ = s.Init()
-
 	s.tokenOnce.Do(func() {
 		s.token, s.tokenErr = s.authClient.Auth(s.params.user, s.params.password)
 	})
@@ -51,8 +49,6 @@ func (s *staticCredentialsProvider) GetToken() (string, error) {
 
 // ContextWithAuth implements Provider.
 func (s *staticCredentialsProvider) ContextWithAuth(ctx context.Context) (context.Context, context.CancelFunc) {
-	_ = s.Init()
-
 	tok, _ := s.GetToken() // TODO(shmel1k@): return err as params
 	ctx, cf := context.WithCancel(ctx)
 	return metadata.AppendToOutgoingContext(ctx,
@@ -61,8 +57,6 @@ func (s *staticCredentialsProvider) ContextWithAuth(ctx context.Context) (contex
 
 // ContextWithoutAuth implements Provider.
 func (s *staticCredentialsProvider) ContextWithoutAuth(ctx context.Context) (context.Context, context.CancelFunc) {
-	_ = s.Init()
-
 	return context.WithCancel(ctx)
 }
 

pkg/options/auth.go

@@ -31,14 +31,16 @@ const (
 	Static         AuthType = "static"
 	IamToken       AuthType = "iam-token"
 	IamCreds       AuthType = "iam-creds"
+	IamMetadata    AuthType = "iam-metadata"
 	MultipleAtOnce AuthType = "multiple-at-once"
 )
 
 var Auths = map[AuthType]Options{
-	None:     &AuthNone{},
-	Static:   &AuthStatic{},
-	IamToken: &AuthIAMToken{},
-	IamCreds: &AuthIAMCreds{},
+	None:        &AuthNone{},
+	Static:      &AuthStatic{},
+	IamToken:    &AuthIAMToken{},
+	IamCreds:    &AuthIAMCreds{},
+	IamMetadata: &AuthIAMMetadataCredentials{},
 	// TODO support OAuth
 }
 
@@ -57,6 +59,10 @@ type (
 		Token     string
 	}
 
+	AuthIAMMetadataCredentials struct {
+		Enabled bool
+	}
+
 	AuthIAMCreds struct {
 		KeyFilename string
 		Endpoint    string
@@ -180,6 +186,21 @@ func (a *AuthIAMCreds) Validate() error {
 	return nil
 }
 
+func (a *AuthIAMMetadataCredentials) DefineFlags(fs *pflag.FlagSet) {
+	fs.BoolVar(&a.Enabled, "use-metadata-credentials", false,
+		`Use metadata service on a virtual machine to get credentials
+                          For more info go to: cloud.yandex.ru/docs/compute/operations/vm-connect/auth-inside-vm
+                          Definition priority:
+                            1. This option
+                            2. Profile specified with --profile option
+                            3. "USE_METADATA_CREDENTIALS" environment variable
+                            4. Active configuration profile (default: 0)`)
+}
+
+func (a *AuthIAMMetadataCredentials) Validate() error {
+	return nil
+}
+
 func determineExplicitAuthType() AuthType {
 	authType := map[AuthType]bool{}
 
@@ -190,6 +211,10 @@ func determineExplicitAuthType() AuthType {
 		}
 	}
 
+	if iamMetadata, ok := Auths[IamMetadata]; ok && iamMetadata.(*AuthIAMMetadataCredentials).Enabled {
+		authType[IamMetadata] = true
+	}
+
 	if static, ok := Auths[IamToken]; ok && static.(*AuthIAMToken).TokenFile != "" {
 		authType[IamToken] = true
 	}