helm-external-secrets/providers/v1/nebius/mysterybox/token_getter.go

/*
Copyright © The ESO Authors

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

    https://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 mysterybox

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"

	lru "github.com/hashicorp/golang-lru"
	"github.com/nebius/gosdk/auth"
	"golang.org/x/sync/singleflight"
	"k8s.io/utils/clock"

	"github.com/external-secrets/external-secrets/providers/v1/nebius/common/sdk/iam"
)

const (
	errInvalidSubjectCreds = "invalid subject credentials: malformed JSON"
)

// TokenGetter is an interface for generating and retrieving authentication tokens.
type TokenGetter interface {
	GetToken(ctx context.Context, apiDomain, subjectCreds string, caCert []byte) (string, error)
}

type tokenCacheKey struct {
	APIDomain        string
	PublicKeyID      string
	ServiceAccountID string
	PrivateKeyHash   string
}

func (k *tokenCacheKey) String() string {
	return k.APIDomain + "|" + k.PublicKeyID + "|" + k.ServiceAccountID + "|" + k.PrivateKeyHash
}

// CachedTokenGetter is responsible for managing Nebius IAM token caching and token exchange processes.
type CachedTokenGetter struct {
	TokenExchanger iam.TokenExchanger
	Clock          clock.Clock
	tokenCache     *lru.Cache
	sf             singleflight.Group
}

// NewCachedTokenGetter initializes a CachedTokenGetter with the specified cache size, token exchanger, and clock.
// Returns a CachedTokenGetter instance and an error if LRU cache creation fails.
func NewCachedTokenGetter(cacheSize int, tokenExchanger iam.TokenExchanger, clock clock.Clock) (*CachedTokenGetter, error) {
	cache, err := lru.New(cacheSize)
	if err != nil {
		return nil, err
	}
	return &CachedTokenGetter{
		tokenCache:     cache,
		TokenExchanger: tokenExchanger,
		Clock:          clock,
	}, nil
}

func isTokenExpired(token *iam.Token, clk clock.Clock) bool {
	now := clk.Now()
	if token.ExpiresAt.After(now) {
		total := token.ExpiresAt.Sub(token.IssuedAt)
		remaining := token.ExpiresAt.Sub(now)
		if remaining > total/10 {
			return false
		}
	}
	return true
}

// GetToken retrieves an IAM token for the given API domain and subject credentials, using a cache to optimize requests.
// It exchanges credentials for a new token if no valid cached token exists or the cached token is nearing expiration.
func (c *CachedTokenGetter) GetToken(ctx context.Context, apiDomain, subjectCreds string, caCert []byte) (string, error) {
	byteCreds := []byte(subjectCreds)
	cacheKey, err := buildTokenCacheKey(byteCreds, apiDomain)

	if err != nil {
		return "", err
	}

	value, ok := c.tokenCache.Get(*cacheKey)
	if ok {
		token := value.(*iam.Token)
		tokenExpired := isTokenExpired(token, c.Clock)
		if !tokenExpired {
			return token.Token, nil
		}
	}

	tokenCacheKeyString := cacheKey.String()

	token, err, _ := c.sf.Do(tokenCacheKeyString, func() (any, error) {
		if v, ok := c.tokenCache.Get(*cacheKey); ok {
			tok := v.(*iam.Token)
			if !isTokenExpired(tok, c.Clock) {
				return tok.Token, nil
			}
		}

		newToken, err := c.TokenExchanger.ExchangeIamToken(ctx, apiDomain, subjectCreds, c.Clock.Now(), caCert)
		if err != nil {
			return "", fmt.Errorf("could not exchange creds to iam token: %w", MapGrpcErrors("create token", err))
		}

		c.tokenCache.Add(*cacheKey, newToken)
		return newToken.Token, nil
	})

	if err != nil {
		return "", err
	}
	return token.(string), nil
}

func buildTokenCacheKey(subjectCreds []byte, apiDomain string) (*tokenCacheKey, error) {
	parsedSubjectCreds := &auth.ServiceAccountCredentials{}
	err := json.Unmarshal(subjectCreds, parsedSubjectCreds)
	if err != nil {
		return nil, errors.New(errInvalidSubjectCreds)
	}
	return &tokenCacheKey{
		APIDomain:        apiDomain,
		PublicKeyID:      parsedSubjectCreds.SubjectCredentials.KeyID,
		ServiceAccountID: parsedSubjectCreds.SubjectCredentials.Subject,
		PrivateKeyHash:   HashBytes([]byte(parsedSubjectCreds.SubjectCredentials.PrivateKey)),
	}, nil
}

var _ TokenGetter = &CachedTokenGetter{}