lti-api/internal/common/service/common.document.storage.go

package service

import (
	"context"
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"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/service/s3"
)

type DocumentStorage interface {
	Upload(ctx context.Context, key string, body io.Reader, size int64, contentType string) (DocumentStorageUploadResult, error)
	Delete(ctx context.Context, key string) error
	URL(key string) string
	PresignURL(ctx context.Context, key string, expires time.Duration) (string, error)
}

type DocumentStorageUploadResult struct {
	Key  string
	URL  string
	ETag string
}

type S3DocumentStorageConfig struct {
	Region         string
	Bucket         string
	AccessKey      string
	SecretKey      string
	Endpoint       string
	BaseURL        string
	ForcePathStyle bool
}

type s3DocumentStorage struct {
	client        *s3.Client
	presignClient *s3.PresignClient
	bucket        string
	base          string
}

func NewS3DocumentStorage(ctx context.Context, cfg S3DocumentStorageConfig) (DocumentStorage, error) {
	bucket := strings.TrimSpace(cfg.Bucket)
	if bucket == "" {
		return nil, errors.New("s3 bucket is required")
	}
	region := strings.TrimSpace(cfg.Region)
	if region == "" {
		region = "us-east-1"
	}

	options := []func(*awsconfig.LoadOptions) error{
		awsconfig.WithRegion(region),
	}

	endpoint := strings.TrimSpace(cfg.Endpoint)
	if endpoint != "" {
		resolver := aws.EndpointResolverWithOptionsFunc(func(service, region string, _ ...interface{}) (aws.Endpoint, error) {
			if service == s3.ServiceID {
				return aws.Endpoint{
					URL:               endpoint,
					SigningRegion:     region,
					HostnameImmutable: true,
				}, nil
			}
			return aws.Endpoint{}, &aws.EndpointNotFoundError{}
		})
		options = append(options, awsconfig.WithEndpointResolverWithOptions(resolver))
	}

	accessKey := strings.TrimSpace(cfg.AccessKey)
	secretKey := strings.TrimSpace(cfg.SecretKey)
	if accessKey != "" && secretKey != "" {
		options = append(options, awsconfig.WithCredentialsProvider(
			credentials.NewStaticCredentialsProvider(accessKey, secretKey, ""),
		))
	}

	awsCfg, err := awsconfig.LoadDefaultConfig(ctx, options...)
	if err != nil {
		return nil, err
	}

	client := s3.NewFromConfig(awsCfg, func(o *s3.Options) {
		o.UsePathStyle = cfg.ForcePathStyle
	})
	presignClient := s3.NewPresignClient(client)

	baseURL := strings.TrimSuffix(strings.TrimSpace(cfg.BaseURL), "/")
	if baseURL == "" {
		if endpoint != "" {
			baseURL = fmt.Sprintf("%s/%s", strings.TrimSuffix(endpoint, "/"), bucket)
		} else {
			baseURL = fmt.Sprintf("https://%s.s3.%s.amazonaws.com", bucket, region)
		}
	}

	return &s3DocumentStorage{
		client:        client,
		presignClient: presignClient,
		bucket:        bucket,
		base:          baseURL,
	}, nil
}

func (s *s3DocumentStorage) Upload(ctx context.Context, key string, body io.Reader, size int64, contentType string) (DocumentStorageUploadResult, error) {
	if strings.TrimSpace(key) == "" {
		return DocumentStorageUploadResult{}, errors.New("storage key is required")
	}
	if size < 0 {
		size = 0
	}
	input := &s3.PutObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
		Body:   body,
	}
	input.ContentLength = aws.Int64(size)
	if ct := strings.TrimSpace(contentType); ct != "" {
		input.ContentType = aws.String(ct)
	}

	out, err := s.client.PutObject(ctx, input)
	if err != nil {
		return DocumentStorageUploadResult{}, err
	}

	var etag string
	if out.ETag != nil {
		etag = strings.Trim(*out.ETag, "\"")
	}

	return DocumentStorageUploadResult{
		Key:  key,
		URL:  s.URL(key),
		ETag: etag,
	}, nil
}

func (s *s3DocumentStorage) Delete(ctx context.Context, key string) error {
	if strings.TrimSpace(key) == "" {
		return nil
	}
	_, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
	})
	return err
}

func (s *s3DocumentStorage) URL(key string) string {
	key = strings.TrimPrefix(strings.TrimSpace(key), "/")
	if key == "" {
		return s.base
	}
	if s.base == "" {
		return key
	}
	return fmt.Sprintf("%s/%s", s.base, key)
}

func (s *s3DocumentStorage) PresignURL(ctx context.Context, key string, expires time.Duration) (string, error) {
	key = strings.TrimPrefix(strings.TrimSpace(key), "/")
	if key == "" {
		return "", errors.New("storage key is required")
	}
	if expires <= 0 {
		expires = 15 * time.Minute
	}

	out, err := s.presignClient.PresignGetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
	}, s3.WithPresignExpires(expires))
	if err != nil {
		return "", err
	}

	return out.URL, nil
}