Datarush/internal/gw/storage/s3.go

package storage

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"mime/multipart"
	"net/url"
	"path/filepath"
	"strings"
	"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/s3"
	"github.com/google/uuid"
)

type S3Storage struct {
	client   *s3.Client
	bucket   string
	region   string
	endpoint string
}

type S3Config struct {
	AccessKeyID     string
	SecretAccessKey string
	Region          string
	Bucket          string
	Endpoint        string
}

func NewS3Storage(cfg S3Config) (*S3Storage, error) {
	var loadOpts []func(*config.LoadOptions) error

	if cfg.Region != "" {
		loadOpts = append(loadOpts, config.WithRegion(cfg.Region))
	}

	if cfg.Endpoint != "" {
		customResolver := aws.EndpointResolverWithOptionsFunc(
			func(service, region string, options ...interface{}) (aws.Endpoint, error) {
				return aws.Endpoint{
					URL:               cfg.Endpoint,
					SigningRegion:     cfg.Region,
					HostnameImmutable: true,
				}, nil
			},
		)
		loadOpts = append(loadOpts, config.WithEndpointResolverWithOptions(customResolver))
	}

	if cfg.AccessKeyID != "" || cfg.SecretAccessKey != "" {
		loadOpts = append(loadOpts, config.WithCredentialsProvider(
			credentials.NewStaticCredentialsProvider(cfg.AccessKeyID, cfg.SecretAccessKey, ""),
		))
	}

	awsCfg, err := config.LoadDefaultConfig(context.TODO(), loadOpts...)
	if err != nil {
		return nil, fmt.Errorf("failed to load AWS config: %w", err)
	}

	var client *s3.Client
	if cfg.Endpoint != "" {
		client = s3.NewFromConfig(awsCfg, func(o *s3.Options) {
			o.UsePathStyle = true
		})
	} else {
		client = s3.NewFromConfig(awsCfg)
	}

	return &S3Storage{
		client:   client,
		bucket:   cfg.Bucket,
		region:   cfg.Region,
		endpoint: cfg.Endpoint,
	}, nil
}

func (s *S3Storage) UploadFile(ctx context.Context, file multipart.File, header *multipart.FileHeader) (string, error) {
	fileBytes, err := io.ReadAll(file)
	if err != nil {
		return "", fmt.Errorf("failed to read file: %w", err)
	}

	ext := filepath.Ext(header.Filename)
	key := fmt.Sprintf("submissions/%s/%s%s",
		time.Now().Format("2006/01/02"),
		uuid.New().String(),
		ext,
	)

	_, err = s.client.PutObject(ctx, &s3.PutObjectInput{
		Bucket:      aws.String(s.bucket),
		Key:         aws.String(key),
		Body:        bytes.NewReader(fileBytes),
		ContentType: aws.String(header.Header.Get("Content-Type")),
	})
	if err != nil {
		return "", fmt.Errorf("failed to upload file to S3: %w", err)
	}

	return s.buildObjectURL(key), nil
}

func (s *S3Storage) UploadFileFromBytes(
	ctx context.Context,
	content []byte,
	filename string,
	contentType string,
) (string, error) {
	ext := filepath.Ext(filename)
	key := fmt.Sprintf("submissions/%s/%s%s",
		time.Now().Format("2006/01/02"),
		uuid.New().String(),
		ext,
	)

	_, err := s.client.PutObject(ctx, &s3.PutObjectInput{
		Bucket:      aws.String(s.bucket),
		Key:         aws.String(key),
		Body:        bytes.NewReader(content),
		ContentType: aws.String(contentType),
	})
	if err != nil {
		return "", fmt.Errorf("failed to upload file to S3: %w", err)
	}

	return s.buildObjectURL(key), nil
}

func (s *S3Storage) DeleteFile(ctx context.Context, fileURL string) error {
	key := s.extractKeyFromURL(fileURL)

	_, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return fmt.Errorf("failed to delete file from S3: %w", err)
	}

	return nil
}

func (s *S3Storage) buildObjectURL(key string) string {
	if s.endpoint != "" {
		ep := strings.TrimRight(s.endpoint, "/")
		return fmt.Sprintf("%s/%s/%s", ep, s.bucket, key)
	}

	return fmt.Sprintf("https://%s.s3.%s.amazonaws.com/%s", s.bucket, s.region, key)
}

func (s *S3Storage) extractKeyFromURL(urlStr string) string {
	u, err := url.Parse(urlStr)
	if err != nil {
		return urlStr
	}

	path := strings.TrimPrefix(u.Path, "/")

	if strings.HasPrefix(u.Host, s.bucket+".") {
		return path
	}

	if strings.HasPrefix(path, s.bucket+"/") {
		return strings.TrimPrefix(path, s.bucket+"/")
	}

	if s.endpoint != "" {
		ep := strings.TrimPrefix(strings.TrimRight(s.endpoint, "/"), "http://")
		ep = strings.TrimPrefix(ep, "https://")
		if strings.HasPrefix(u.Host, ep) {
			if strings.HasPrefix(path, s.bucket+"/") {
				return strings.TrimPrefix(path, s.bucket+"/")
			}
			return path
		}
	}

	return path
}