aws.go

package utils

import (
	"context"
	"crypto/tls"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	smithyendpoints "github.com/aws/smithy-go/endpoints"
	"github.com/aws/smithy-go/ptr"
	"github.com/google/uuid"

	"github.com/PeerDB-io/peerdb/flow/generated/protos"
	"github.com/PeerDB-io/peerdb/flow/internal"
	"github.com/PeerDB-io/peerdb/flow/shared"
)

const (
	_peerDBCheck = "peerdb_check"
)

var s3CompatibleServiceEndpointPattern = regexp.MustCompile(`^https?://[a-zA-Z0-9.-]+(:\d+)?$`)

type AWSSecrets struct {
	AccessKeyID     string
	SecretAccessKey string
	AwsRoleArn      string
	Region          string
	Endpoint        string
	SessionToken    string
}

type PeerAWSCredentials struct {
	Credentials    aws.Credentials
	RoleArn        *string
	ChainedRoleArn *string
	EndpointUrl    *string
	Region         string
	RootCAs        *string
	TlsHost        string
}

func NewPeerAWSCredentials(s3 *protos.S3Config) PeerAWSCredentials {
	if s3 == nil {
		return PeerAWSCredentials{}
	}
	return PeerAWSCredentials{
		Credentials: aws.Credentials{
			AccessKeyID:     s3.GetAccessKeyId(),
			SecretAccessKey: s3.GetSecretAccessKey(),
		},
		RoleArn:        s3.RoleArn,
		ChainedRoleArn: nil,
		EndpointUrl:    s3.Endpoint,
		Region:         s3.GetRegion(),
		RootCAs:        s3.RootCa,
		TlsHost:        s3.TlsHost,
	}
}

type ClickHouseS3Credentials struct {
	Provider   AWSCredentialsProvider
	BucketPath string
}

type AWSCredentials struct {
	EndpointUrl *string
	AWS         aws.Credentials
}

type AWSCredentialsProvider interface {
	Retrieve(ctx context.Context) (AWSCredentials, error)
	GetUnderlyingProvider() aws.CredentialsProvider
	GetRegion() string
	GetEndpointURL() string
	GetTlsConfig() (*string, string)
}

type ConfigBasedAWSCredentialsProvider struct {
	config aws.Config
}

func NewConfigBasedAWSCredentialsProvider(config aws.Config) *ConfigBasedAWSCredentialsProvider {
	return &ConfigBasedAWSCredentialsProvider{config: config}
}

func (r *ConfigBasedAWSCredentialsProvider) GetUnderlyingProvider() aws.CredentialsProvider {
	return r.config.Credentials
}

func (r *ConfigBasedAWSCredentialsProvider) GetRegion() string {
	return r.config.Region
}

func (r *ConfigBasedAWSCredentialsProvider) GetEndpointURL() string {
	endpoint := ""
	if r.config.BaseEndpoint != nil {
		endpoint = *r.config.BaseEndpoint
	}

	return endpoint
}

func (r *ConfigBasedAWSCredentialsProvider) GetTlsConfig() (*string, string) {
	return nil, ""
}

// Retrieve should be called as late as possible in order to have credentials with latest expiry
func (r *ConfigBasedAWSCredentialsProvider) Retrieve(ctx context.Context) (AWSCredentials, error) {
	retrieved, err := r.config.Credentials.Retrieve(ctx)
	if err != nil {
		return AWSCredentials{}, err
	}
	return AWSCredentials{
		AWS:         retrieved,
		EndpointUrl: r.config.BaseEndpoint,
	}, nil
}

type StaticAWSCredentialsProvider struct {
	credentials AWSCredentials
	region      string
	rootCAs     *string
	tlsHost     string
}

func NewStaticAWSCredentialsProvider(credentials AWSCredentials, region string, rootCAs *string, tlsHost string) *StaticAWSCredentialsProvider {
	return &StaticAWSCredentialsProvider{
		credentials: credentials,
		region:      region,
		rootCAs:     rootCAs,
		tlsHost:     tlsHost,
	}
}

func (s *StaticAWSCredentialsProvider) GetUnderlyingProvider() aws.CredentialsProvider {
	return credentials.NewStaticCredentialsProvider(s.credentials.AWS.AccessKeyID, s.credentials.AWS.SecretAccessKey,
		s.credentials.AWS.SessionToken)
}

func (s *StaticAWSCredentialsProvider) GetRegion() string {
	return s.region
}

func (s *StaticAWSCredentialsProvider) Retrieve(ctx context.Context) (AWSCredentials, error) {
	return s.credentials, nil
}

func (s *StaticAWSCredentialsProvider) GetEndpointURL() string {
	if s.credentials.EndpointUrl != nil {
		return *s.credentials.EndpointUrl
	}
	return ""
}

func (s *StaticAWSCredentialsProvider) GetTlsConfig() (*string, string) {
	return s.rootCAs, s.tlsHost
}

type AssumeRoleBasedAWSCredentialsProvider struct {
	Provider aws.CredentialsProvider // New Credentials
	config   aws.Config              // Initial Config
}

func NewAssumeRoleBasedAWSCredentialsProvider(
	ctx context.Context,
	config aws.Config,
	roleArn string,
	sessionName string,
) (*AssumeRoleBasedAWSCredentialsProvider, error) {
	provider := stscreds.NewAssumeRoleProvider(sts.NewFromConfig(config), roleArn, func(o *stscreds.AssumeRoleOptions) {
		o.RoleSessionName = sessionName
	})
	_, err := provider.Retrieve(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to retrieve chained AWS credentials: %w", err)
	}
	return &AssumeRoleBasedAWSCredentialsProvider{
		config:   config,
		Provider: aws.NewCredentialsCache(provider),
	}, nil
}

func (a *AssumeRoleBasedAWSCredentialsProvider) Retrieve(ctx context.Context) (AWSCredentials, error) {
	retrieved, err := a.Provider.Retrieve(ctx)
	if err != nil {
		return AWSCredentials{}, err
	}
	return AWSCredentials{
		AWS:         retrieved,
		EndpointUrl: ptr.String(a.GetEndpointURL()),
	}, nil
}

func (a *AssumeRoleBasedAWSCredentialsProvider) GetUnderlyingProvider() aws.CredentialsProvider {
	return a.Provider
}

func (a *AssumeRoleBasedAWSCredentialsProvider) GetRegion() string {
	return a.config.Region
}

func (a *AssumeRoleBasedAWSCredentialsProvider) GetEndpointURL() string {
	endpoint := ""
	if a.config.BaseEndpoint != nil {
		endpoint = *a.config.BaseEndpoint
	}

	return endpoint
}

func (a *AssumeRoleBasedAWSCredentialsProvider) GetTlsConfig() (*string, string) {
	return nil, ""
}

func getPeerDBAWSEnv(connectorName string, awsKey string) string {
	return os.Getenv(fmt.Sprintf("PEERDB_%s_AWS_CREDENTIALS_%s", strings.ToUpper(connectorName), awsKey))
}

func LoadPeerDBAWSEnvConfigProvider(connectorName string) *StaticAWSCredentialsProvider {
	accessKeyId := getPeerDBAWSEnv(connectorName, "AWS_ACCESS_KEY_ID")
	secretAccessKey := getPeerDBAWSEnv(connectorName, "AWS_SECRET_ACCESS_KEY")
	region := getPeerDBAWSEnv(connectorName, "AWS_REGION")
	endpointUrl := getPeerDBAWSEnv(connectorName, "AWS_ENDPOINT_URL_S3")
	rootCa := getPeerDBAWSEnv(connectorName, "ROOT_CA")
	tlsHost := getPeerDBAWSEnv(connectorName, "TLS_HOST")
	var endpointUrlPtr *string
	if endpointUrl != "" {
		endpointUrlPtr = &endpointUrl
	}

	if accessKeyId == "" && secretAccessKey == "" && region == "" && endpointUrl == "" {
		return nil
	}

	var rootCAs *string
	if rootCa != "" {
		rootCAs = &rootCa
	}

	return NewStaticAWSCredentialsProvider(AWSCredentials{
		AWS: aws.Credentials{
			AccessKeyID:     accessKeyId,
			SecretAccessKey: secretAccessKey,
		},
		EndpointUrl: endpointUrlPtr,
	}, region, rootCAs, tlsHost)
}

func GetAWSCredentialsProvider(ctx context.Context, connectorName string, peerCredentials PeerAWSCredentials) (AWSCredentialsProvider, error) {
	logger := internal.LoggerFromCtx(ctx)
	if peerCredentials.Credentials.AccessKeyID != "" || peerCredentials.Credentials.SecretAccessKey != "" ||
		peerCredentials.Region != "" || (peerCredentials.RoleArn != nil && *peerCredentials.RoleArn != "") ||
		(peerCredentials.ChainedRoleArn != nil && *peerCredentials.ChainedRoleArn != "") ||
		(peerCredentials.EndpointUrl != nil && *peerCredentials.EndpointUrl != "") {
		staticProvider := NewStaticAWSCredentialsProvider(AWSCredentials{
			AWS:         peerCredentials.Credentials,
			EndpointUrl: peerCredentials.EndpointUrl,
		}, peerCredentials.Region, peerCredentials.RootCAs, peerCredentials.TlsHost)
		if peerCredentials.RoleArn == nil || *peerCredentials.RoleArn == "" {
			logger.Info("Received AWS credentials from peer for connector: " + connectorName)
			return staticProvider, nil
		}
		awsConfig, err := config.LoadDefaultConfig(ctx)
		if err != nil {
			return nil, err
		}
		awsConfig.Credentials = stscreds.NewAssumeRoleProvider(sts.NewFromConfig(awsConfig), *peerCredentials.RoleArn,
			func(options *stscreds.AssumeRoleOptions) {
				options.RoleSessionName = getAssumedRoleSessionName()
			},
		)
		if peerCredentials.ChainedRoleArn != nil && *peerCredentials.ChainedRoleArn != "" {
			logger.Info("Received AWS credentials with chained role from peer for connector: " + connectorName)
			return NewAssumeRoleBasedAWSCredentialsProvider(ctx, awsConfig, *peerCredentials.ChainedRoleArn, getChainedRoleSessionName())
		}
		logger.Info("Received AWS credentials from peer for connector: " + connectorName)
		return NewConfigBasedAWSCredentialsProvider(awsConfig), nil
	}
	envCredentialsProvider := LoadPeerDBAWSEnvConfigProvider(connectorName)
	if envCredentialsProvider != nil {
		logger.Info("Received AWS credentials from PeerDB Env for connector: " + connectorName)
		return envCredentialsProvider, nil
	}

	awsConfig, err := config.LoadDefaultConfig(ctx, func(options *config.LoadOptions) error {
		return nil
	})
	if err != nil {
		return nil, err
	}
	logger.Info("Received AWS credentials from SDK config for connector: " + connectorName)
	return NewConfigBasedAWSCredentialsProvider(awsConfig), nil
}

const MaxAWSSessionNameLength = 63 // Docs mention 64 as limit, but always good to stay under

func getAssumedRoleSessionName() string {
	defaultSessionName := "peeraws"
	if deployUid := internal.PeerDBDeploymentUID(); deployUid != "" {
		defaultSessionName += "-" + deployUid
	}
	sessionName := internal.GetEnvString("PEERDB_AWS_ASSUMED_ROLE_SESSION_NAME", defaultSessionName)
	if len(sessionName) > MaxAWSSessionNameLength {
		sessionName = sessionName[:MaxAWSSessionNameLength-1]
	}
	return sessionName
}

func getChainedRoleSessionName() string {
	defaultSessionName := "peerchain"
	if deployUid := internal.PeerDBDeploymentUID(); deployUid != "" {
		defaultSessionName += "-" + deployUid
	}
	sessionName := internal.GetEnvString("PEERDB_AWS_CHAINED_ROLE_SESSION_NAME", defaultSessionName)
	if len(sessionName) > MaxAWSSessionNameLength {
		sessionName = sessionName[:MaxAWSSessionNameLength-1]
	}
	return sessionName
}

func FileURLForS3Service(endpoint string, region string, bucket string, filePath string) string {
	if s3CompatibleServiceEndpointPattern.MatchString(endpoint) {
		return fmt.Sprintf("%s/%s/%s", endpoint, bucket, filePath)
	}
	return fmt.Sprintf("https://%s.s3.%s.amazonaws.com/%s", bucket, region, filePath)
}

type S3BucketAndPrefix struct {
	Bucket string
	Prefix string
}

// path would be something like s3://bucket/prefix
func NewS3BucketAndPrefix(s3Path string) (*S3BucketAndPrefix, error) {
	// Remove s3:// prefix
	stagingPath := strings.TrimPrefix(s3Path, "s3://")

	// Split into bucket and prefix
	bucket, prefix, _ := strings.Cut(stagingPath, "/")

	return &S3BucketAndPrefix{
		Bucket: bucket,
		Prefix: strings.Trim(prefix, "/"),
	}, nil
}

type resolverV2 struct {
	url.URL
}

func (r *resolverV2) ResolveEndpoint(ctx context.Context, params s3.EndpointParameters) (
	smithyendpoints.Endpoint, error,
) {
	uri := r.URL
	uri.Path += "/" + *params.Bucket
	return smithyendpoints.Endpoint{
		URI: uri,
	}, nil
}

func CreateS3Client(ctx context.Context, credsProvider AWSCredentialsProvider) (*s3.Client, error) {
	awsCredentials, err := credsProvider.Retrieve(ctx)
	if err != nil {
		return nil, err
	}

	options := s3.Options{
		Region:      credsProvider.GetRegion(),
		Credentials: credsProvider.GetUnderlyingProvider(),
	}
	if awsCredentials.EndpointUrl != nil && *awsCredentials.EndpointUrl != "" {
		options.BaseEndpoint = awsCredentials.EndpointUrl
		options.UsePathStyle = true
		url, err := url.Parse(*awsCredentials.EndpointUrl)
		if err != nil {
			return nil, err
		}
		options.EndpointResolverV2 = &resolverV2{
			URL: *url,
		}

		if strings.Contains(*awsCredentials.EndpointUrl, "storage.googleapis.com") {
			// Assign custom client with our own transport
			options.HTTPClient = &http.Client{
				Transport: &RecalculateV4Signature{
					next:        http.DefaultTransport,
					signer:      v4.NewSigner(),
					credentials: credsProvider.GetUnderlyingProvider(),
					region:      options.Region,
				},
			}
		} else {
			rootCAs, tlsHost := credsProvider.GetTlsConfig()
			if rootCAs != nil || tlsHost != "" {
				// start with a clone of DefaultTransport so we keep http2, idle-conns, etc.
				tlsConfig, err := shared.CreateTlsConfig(tls.VersionTLS13, rootCAs, tlsHost, tlsHost, tlsHost == "")
				if err != nil {
					return nil, err
				}

				tr := http.DefaultTransport.(*http.Transport).Clone()
				tr.TLSClientConfig = tlsConfig
				options.HTTPClient = &http.Client{Transport: tr}
			}
		}
	}

	return s3.New(options), nil
}

// RecalculateV4Signature allow GCS over S3, removing Accept-Encoding header from sign
// https://stackoverflow.com/a/74382598/1204665
// https://github.com/aws/aws-sdk-go-v2/issues/1816
type RecalculateV4Signature struct {
	next        http.RoundTripper
	signer      *v4.Signer
	credentials aws.CredentialsProvider
	region      string
}

func (lt *RecalculateV4Signature) RoundTrip(req *http.Request) (*http.Response, error) {
	// store for later use
	acceptEncodingValue := req.Header.Get("Accept-Encoding")

	// delete the header so the header doesn't account for in the signature
	req.Header.Del("Accept-Encoding")

	// sign with the same date
	timeString := req.Header.Get("X-Amz-Date")
	timeDate, _ := time.Parse("20060102T150405Z", timeString)

	creds, err := lt.credentials.Retrieve(req.Context())
	if err != nil {
		return nil, err
	}
	if err := lt.signer.SignHTTP(req.Context(), creds, req, v4.GetPayloadHash(req.Context()), "s3", lt.region, timeDate); err != nil {
		return nil, err
	}
	// Reset Accept-Encoding if desired
	req.Header.Set("Accept-Encoding", acceptEncodingValue)

	// follows up the original round tripper
	return lt.next.RoundTrip(req)
}

// Write an empty file and then delete it
// to check if we have write permissions
func PutAndRemoveS3(ctx context.Context, client *s3.Client, bucket string, prefix string) error {
	reader := strings.NewReader(time.Now().Format(time.RFC3339))
	bucketName := aws.String(bucket)
	temporaryObjectPath := prefix + "/" + _peerDBCheck + uuid.New().String()
	key := aws.String(strings.TrimPrefix(temporaryObjectPath, "/"))

	if _, putErr := client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: bucketName,
		Key:    key,
		Body:   reader,
	}); putErr != nil {
		return fmt.Errorf("failed to write to bucket: %w", putErr)
	}

	if _, delErr := client.DeleteObject(ctx, &s3.DeleteObjectInput{
		Bucket: bucketName,
		Key:    key,
	}); delErr != nil {
		return fmt.Errorf("failed to delete from bucket: %w", delErr)
	}

	return nil
}