aliyun / credentials-go

import (

	"bytes"

	"context"

	"encoding/json"

	"errors"

	"fmt"

	"os"

	"os/exec"

	"strings"

	"sync"

	"time"

)

type ExternalOptions struct {

	// Timeout, in milliseconds.

	Timeout int

}

// ExternalCredentialUpdateCallback 定义External凭证更新回调函数类型

type ExternalCredentialUpdateCallback func(accessKeyId, accessKeySecret, securityToken string, expiration int64) error

type externalCredentialResponse struct {

	Mode            string `json:"mode"`

	AccessKeyId     string `json:"access_key_id"`

	AccessKeySecret string `json:"access_key_secret"`

	SecurityToken   string `json:"sts_token"`

	Expiration      string `json:"expiration,omitempty"`

}

type ExternalCredentialsProvider struct {

	processCommand string

	options        *ExternalOptions

	lastUpdateTimestamp int64

	expirationTimestamp int64

	sessionCredentials  *sessionCredentials

	// External credential call back

	credentialUpdateCallback ExternalCredentialUpdateCallback

	// 互斥锁，用于并发安全

	mu sync.RWMutex

}

type ExternalCredentialsProviderBuilder struct {

	provider *ExternalCredentialsProvider

}

func NewExternalCredentialsProviderBuilder() *ExternalCredentialsProviderBuilder {

	return &ExternalCredentialsProviderBuilder{

		provider: &ExternalCredentialsProvider{},

	}

}

func (b *ExternalCredentialsProviderBuilder) WithProcessCommand(processCommand string) *ExternalCredentialsProviderBuilder {

	b.provider.processCommand = processCommand

	return b

}

func (b *ExternalCredentialsProviderBuilder) WithExternalOptions(options *ExternalOptions) *ExternalCredentialsProviderBuilder {

	b.provider.options = options

	return b

}

func (b *ExternalCredentialsProviderBuilder) WithCredentialUpdateCallback(callback ExternalCredentialUpdateCallback) *ExternalCredentialsProviderBuilder {

	b.provider.credentialUpdateCallback = callback

	return b

}

func (b *ExternalCredentialsProviderBuilder) Build() (provider *ExternalCredentialsProvider, err error) {

	if b.provider.processCommand == "" {

		err = errors.New("process_command is empty")

		return

	}

	provider = b.provider

	return

}

func (provider *ExternalCredentialsProvider) getCredentials() (session *sessionCredentials, err error) {

	args := strings.Fields(provider.processCommand)

	if len(args) == 0 {

		err = errors.New("process_command is empty")

		return

	}

	// 确保 options 不为 nil，并设置默认超时时间

	timeout := 60 * 1000 // 默认 60 秒

	if provider.options != nil && provider.options.Timeout > 0 {

		timeout = provider.options.Timeout

	}

	var cancelFunc func()

	ctx, cancelFunc := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Millisecond)

	defer cancelFunc()

	cmd := exec.CommandContext(ctx, args[0], args[1:]...)

	cmd.Env = os.Environ()

	// 创建一个buffer来捕获标准输出

	var stdoutBuf bytes.Buffer

	cmd.Stdout = &stdoutBuf

	// 创建一个buffer来捕获标准错误输出

	var stderrBuf bytes.Buffer

	cmd.Stderr = &stderrBuf

	// Start the command

	if err := cmd.Start(); err != nil {

		return nil, fmt.Errorf("failed to execute external command: %w\nstderr: %s", err, stderrBuf.String())

	}

	done := make(chan error, 1)

	go func() {

		done <- cmd.Wait()

	}()

	select {

	case <-ctx.Done():

		// 超时了，context 会自动终止命令

		<-done

		return nil, fmt.Errorf("command process timed out after %d milliseconds", timeout)

	case execError := <-done:

		if execError != nil {

			// 检查是否是超时导致的错误

			if errors.Is(execError, context.DeadlineExceeded) {

				return nil, fmt.Errorf("command process timed out after %d milliseconds", timeout)

			}

			return nil, fmt.Errorf("failed to execute external command: %w\nstderr: %s", execError, stderrBuf.String())

		}

	}

	// 只解析标准输出

	buf := stdoutBuf.Bytes()

	// 解析得到凭证响应

	var resp externalCredentialResponse

	err = json.Unmarshal(buf, &resp)

	if err != nil {

		fmt.Println(provider.processCommand)

		fmt.Println(string(buf))

		return nil, fmt.Errorf("failed to parse external command output: %w", err)

	}

	// 验证返回的凭证数据

	if resp.AccessKeyId == "" || resp.AccessKeySecret == "" {

		return nil, fmt.Errorf("invalid credential response: access_key_id or access_key_secret is empty")

	}

	// 根据 mode 验证 SecurityToken

	if resp.Mode == "StsToken" && resp.SecurityToken == "" {

		return nil, fmt.Errorf("invalid StsToken credential response: sts_token is empty")

	}

	session = &sessionCredentials{

		AccessKeyId:     resp.AccessKeyId,

		AccessKeySecret: resp.AccessKeySecret,

		SecurityToken:   resp.SecurityToken,

		Expiration:      resp.Expiration,

	}

	return

}

func (provider *ExternalCredentialsProvider) needUpdateCredential() (result bool) {

	provider.mu.RLock()

	defer provider.mu.RUnlock()

	// 如果没有缓存凭证，需要更新

	if provider.sessionCredentials == nil {

		return true

	}

	// 如果没有过期时间，每次都更新（因为外部命令可能返回动态凭证）

	if provider.expirationTimestamp == 0 {

		return true

	}

	// 如果凭证即将过期（提前180秒），需要更新

	return provider.expirationTimestamp-time.Now().Unix() <= 180

}

func (provider *ExternalCredentialsProvider) GetProviderName() string {

	return "external"

}