kra-oa/internal/data/storage/aws_storage.go

package storage

import (
	"context"
	"fmt"
	"io"
	"path"
	"strings"

	"github.com/aws/aws-sdk-go-v2/aws"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"kra/internal/biz"
	"kra/internal/conf"
)

type awsStorage struct {
	client                  *s3.Client
	bucket, baseURL, prefix string
}

func newAWSStorage(provider string, config *conf.AdminBackend_ObjectStore) (biz.FileStorage, error) {
	if config == nil || config.Bucket == "" || config.AccessKey == "" || config.SecretKey == "" {
		return nil, fmt.Errorf("%s storage configuration is incomplete", provider)
	}
	region := config.Region
	endpoint := config.Endpoint
	if provider == "cloudflare-r2" {
		region = "auto"
		if endpoint == "" {
			if config.AccountId == "" {
				return nil, fmt.Errorf("cloudflare-r2 account id is required")
			}
			endpoint = fmt.Sprintf("https://%s.r2.cloudflarestorage.com", config.AccountId)
		}
	}
	if region == "" {
		return nil, fmt.Errorf("%s region is required", provider)
	}
	loaded, err := awsconfig.LoadDefaultConfig(context.Background(), awsconfig.WithRegion(region), awsconfig.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(config.AccessKey, config.SecretKey, "")))
	if err != nil {
		return nil, err
	}
	client := s3.NewFromConfig(loaded, func(options *s3.Options) {
		if endpoint != "" {
			if !strings.HasPrefix(endpoint, "http://") && !strings.HasPrefix(endpoint, "https://") {
				if config.UseSsl {
					endpoint = "https://" + endpoint
				} else {
					endpoint = "http://" + endpoint
				}
			}
			options.BaseEndpoint = aws.String(endpoint)
		}
		options.UsePathStyle = config.ForcePathStyle
	})
	return &awsStorage{client: client, bucket: config.Bucket, baseURL: strings.TrimSuffix(config.BaseUrl, "/"), prefix: strings.Trim(config.PathPrefix, "/")}, nil
}

func (s *awsStorage) key(name string) string {
	if s.prefix == "" {
		return strings.TrimPrefix(name, "/")
	}
	return path.Join(s.prefix, strings.TrimPrefix(name, "/"))
}
func (s *awsStorage) unkey(key string) string {
	return strings.TrimPrefix(strings.TrimPrefix(key, s.prefix), "/")
}
func (s *awsStorage) file(key string, size int64) *biz.StoredFile {
	name := s.unkey(key)
	url := s.baseURL + "/" + key
	return &biz.StoredFile{Name: path.Base(name), Path: name, URL: url, Size: size}
}
func (s *awsStorage) Put(ctx context.Context, name string, reader io.Reader) (*biz.StoredFile, error) {
	key := s.key(name)
	result, err := manager.NewUploader(s.client).Upload(ctx, &s3.PutObjectInput{Bucket: aws.String(s.bucket), Key: aws.String(key), Body: reader})
	if err != nil {
		return nil, err
	}
	_ = result
	head, err := s.client.HeadObject(ctx, &s3.HeadObjectInput{Bucket: aws.String(s.bucket), Key: aws.String(key)})
	if err != nil {
		return nil, err
	}
	size := int64(0)
	if head.ContentLength != nil {
		size = *head.ContentLength
	}
	return s.file(key, size), nil
}
func (s *awsStorage) Open(ctx context.Context, name string) (io.ReadCloser, error) {
	result, err := s.client.GetObject(ctx, &s3.GetObjectInput{Bucket: aws.String(s.bucket), Key: aws.String(s.key(name))})
	if err != nil {
		return nil, err
	}
	return result.Body, nil
}
func (s *awsStorage) Delete(ctx context.Context, name string) error {
	_, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{Bucket: aws.String(s.bucket), Key: aws.String(s.key(name))})
	return err
}
func (s *awsStorage) Compose(ctx context.Context, names []string, destination string) (*biz.StoredFile, string, error) {
	return composeFiles(ctx, s, names, destination)
}
func (s *awsStorage) DeletePrefix(ctx context.Context, prefix string) error {
	for {
		items, _, more, err := s.List(ctx, prefix, "", 1000)
		if err != nil {
			return err
		}
		for _, item := range items {
			if err = s.Delete(ctx, item.Path); err != nil {
				return err
			}
		}
		if !more || len(items) == 0 {
			return nil
		}
	}
}
func (s *awsStorage) List(ctx context.Context, prefix, cursor string, limit int) ([]*biz.StoredFile, string, bool, error) {
	if limit < 1 || limit > 1000 {
		limit = 100
	}
	input := &s3.ListObjectsV2Input{Bucket: aws.String(s.bucket), Prefix: aws.String(s.key(prefix)), MaxKeys: aws.Int32(int32(limit))}
	if cursor != "" {
		input.ContinuationToken = aws.String(cursor)
	}
	result, err := s.client.ListObjectsV2(ctx, input)
	if err != nil {
		return nil, "", false, err
	}
	items := make([]*biz.StoredFile, 0, len(result.Contents))
	for _, object := range result.Contents {
		if object.Key == nil {
			continue
		}
		size := int64(0)
		if object.Size != nil {
			size = *object.Size
		}
		item := s.file(*object.Key, size)
		if object.LastModified != nil {
			item.LastModified = *object.LastModified
		}
		items = append(items, item)
	}
	next := ""
	if result.NextContinuationToken != nil {
		next = *result.NextContinuationToken
	}
	more := result.IsTruncated != nil && *result.IsTruncated
	return items, next, more, nil
}