/****************************************************************************** * Copyright (c) 2022 PingCAP * * 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. * ******************************************************************************/ /******************************************************************************* * @File: sts.go * @Description: * @Author: jiangxunyu@pingcap.com * @Version: 1.0.0 * @Date: 2022/6/22 17:04 *******************************************************************************/ package awsauth import ( "context" "encoding/base64" "encoding/json" "fmt" "os" "strconv" "sync" "time" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/credentials" "github.com/aws/aws-sdk-go-v2/service/iam" "github.com/aws/aws-sdk-go-v2/service/sts" "github.com/aws/aws-sdk-go-v2/service/sts/types" "github.com/patrickmn/go-cache" "github.com/pingcap/log" "github.com/tidbcloud/observability/pkg/common/constants" "github.com/tidbcloud/observability/pkg/common/errors" "github.com/tidbcloud/observability/pkg/common/utils" "go.uber.org/zap" ) // cacheStsTokenKey cache sts token key for global server to regional server authentication const cacheStsTokenKey = constants.ObservabilityGlobalServer // durationSeconds define sts token duration ttl, default minimum value is 900s const durationSeconds = 900 var lock sync.Mutex var stsTokenCache = cache.New(constants.DefaultCacheTimeout, constants.DefaultCacheTimeout*2) var awsCfg *aws.Config func SetAwsConfig(cfg *aws.Config) { awsCfg = cfg } func GetStsTokenCache() *cache.Cache { return stsTokenCache } // GetStsToken // @Description: get or gen sts token // @return string // @return error func GetStsToken(ctx context.Context) (string, error) { log.Debug("get cache sts token", zap.Int("STS_Token_Size", stsTokenCache.ItemCount())) // cache validates if token, ok := stsTokenCache.Get(cacheStsTokenKey); ok { // cache validate success return token.(string), nil } lock.Lock() defer lock.Unlock() // double check lock if once set cache success if token, ok := stsTokenCache.Get(cacheStsTokenKey); ok { // cache validate success return token.(string), nil } // generate sts credentials cred, err := genStsCredentials(ctx) if err != nil { return "", err } b, err := json.Marshal(cred) if err != nil { return "", errors.NewErrorf(errors.MarshalError, "%s", err.Error()) } token := base64.StdEncoding.EncodeToString(b) // set cache, the default cache expiry time is the sts duration seconds minus 60 // letting the cache expire early stsTokenCache.Set(cacheStsTokenKey, token, time.Duration(durationSeconds-60)*time.Second) return token, nil } // ValidateStsToken // @Description: validate sts token // @Parameter stsToken // @return error func ValidateStsToken(ctx context.Context, stsToken string) error { log.Debug("get cache sts token", zap.Int("STS_Token_Size", stsTokenCache.ItemCount())) // cache validate if exp, ok := stsTokenCache.Get(stsToken); ok { if time.Now().Unix() < exp.(int64) { // cache validate success log.Debug("cache validate sts token successfully") return nil } } lock.Lock() defer lock.Unlock() // double check lock if exp, ok := stsTokenCache.Get(stsToken); ok { if time.Now().Unix() < exp.(int64) { // cache validate success return nil } } dec, err := base64.StdEncoding.DecodeString(stsToken) if err != nil { return errors.NewErrorf(errors.B64DecodeError, "%s", err.Error()) } // build sts credentials var cred types.Credentials if err = json.Unmarshal(dec, &cred); err != nil { log.Error("cred unmarshal failed", zap.Error(err), zap.String("Cred", string(dec))) return errors.NewErrorf(errors.UnmarshalError, "%s", err.Error()) } // get region region := os.Getenv("AWS_REGION") log.Debug("load the env `AWS_REGION`", zap.String("AWS_REGION", region)) if region == "" { region = constants.IAMDefaultRegion } if awsCfg == nil { cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region)) if err != nil { return errors.NewErrorf(errors.AWSClientError, "%s", err.Error()) } awsCfg = &cfg } awsCfg.Credentials = credentials.StaticCredentialsProvider{Value: aws.Credentials{ AccessKeyID: *cred.AccessKeyId, SecretAccessKey: *cred.SecretAccessKey, SessionToken: *cred.SessionToken, }} // create iam client iamCli := iam.NewFromConfig(*awsCfg) _, err = iamCli.ListRoles(ctx, &iam.ListRolesInput{}) if err != nil { log.Error("iam list role failed", zap.Error(err)) return err } // set cache, the default cache expiry time is the sts duration seconds minus 60 // letting the cache expire early expire := time.Duration(durationSeconds-60) * time.Second if expire <= 0 { expire = constants.DefaultCacheTimeout } stsTokenCache.Set(stsToken, cred.Expiration.Unix(), expire) return nil } // genStsCredentials // @Description: gen sts credentials // @return *sts.Credentials // @return error func genStsCredentials(ctx context.Context) (*types.Credentials, error) { // the AWS_ROLE_ARN and AWS_WEB_IDENTITY_TOKEN_FILE environment variables are automatically injected by eks when the pod is pulled up. // used to assume role to automatically withdraw rights to obtain sts. roleArn := os.Getenv("AWS_ROLE_ARN") // e.g.: AWS_WEB_IDENTITY_TOKEN_FILE=/var/run/secrets/eks.amazonaws.com/serviceaccount/token tokenFile := os.Getenv("AWS_WEB_IDENTITY_TOKEN_FILE") region := os.Getenv("AWS_REGION") log.Debug("load the env `AWS_ROLE_ARN` and `AWS_WEB_IDENTITY_TOKEN_FILE` and `AWS_REGION`", zap.String("AWS_ROLE_ARN", roleArn), zap.String("AWS_WEB_IDENTITY_TOKEN_FILE", tokenFile), zap.String("AWS_REGION", region)) if roleArn == "" { return nil, errors.NewErrorf(errors.EmptyEnvError, "env `AWS_ROLE_ARN` is empty") } if tokenFile == "" { return nil, errors.NewErrorf(errors.EmptyEnvError, "env `AWS_WEB_IDENTITY_TOKEN_FILE` is empty") } if region == "" { return nil, errors.NewErrorf(errors.EmptyEnvError, "env `AWS_REGION` is empty") } // parse the value of tokenFile, determine path if it exists, and read the contents of the file if !utils.IsFile(tokenFile) { return nil, errors.NewErrorf(errors.TokenFileNotExistsError, "token file `%s` not exists", tokenFile) } webIdentityToken, err := os.ReadFile(tokenFile) if err != nil { return nil, errors.New(errors.IOReadFileError, err.Error()) } if awsCfg == nil { cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region)) if err != nil { return nil, errors.NewErrorf(errors.AWSClientError, "%s", err.Error()) } awsCfg = &cfg } stsCli := sts.NewFromConfig(*awsCfg) cert, err := stsCli.AssumeRoleWithWebIdentity(ctx, &sts.AssumeRoleWithWebIdentityInput{ DurationSeconds: aws.Int32(durationSeconds), RoleArn: aws.String(roleArn), RoleSessionName: aws.String(roleSessionName()), WebIdentityToken: aws.String(string(webIdentityToken)), }) if err != nil { return nil, errors.NewErrorf(errors.GenStsCredentialError, "%s", err.Error()) } log.Info("generate sts credential successfully") return cert.Credentials, nil } // roleSessionName // @Description: gen or get role session name // @return string func roleSessionName() string { suffix, err := os.Hostname() if err != nil { suffix = strconv.FormatInt(time.Now().Unix(), 10) } return fmt.Sprintf("%s_%s", constants.ObservabilityGlobalServer, suffix) }