distributor/internal/adapters/s3/backend.go

package s3

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"path"
	"sort"
	"strings"

	"gitea.maximumdirect.net/eric/distributor/internal/storage"
	"github.com/aws/aws-sdk-go-v2/aws"
	awscfg "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	awss3 "github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/aws/smithy-go"
)

type Client interface {
	HeadObject(ctx context.Context, input *awss3.HeadObjectInput, optFns ...func(*awss3.Options)) (*awss3.HeadObjectOutput, error)
	GetObject(ctx context.Context, input *awss3.GetObjectInput, optFns ...func(*awss3.Options)) (*awss3.GetObjectOutput, error)
	PutObject(ctx context.Context, input *awss3.PutObjectInput, optFns ...func(*awss3.Options)) (*awss3.PutObjectOutput, error)
	ListObjectsV2(ctx context.Context, input *awss3.ListObjectsV2Input, optFns ...func(*awss3.Options)) (*awss3.ListObjectsV2Output, error)
	DeleteObject(ctx context.Context, input *awss3.DeleteObjectInput, optFns ...func(*awss3.Options)) (*awss3.DeleteObjectOutput, error)
}

type Backend struct {
	client Client
	bucket string
	prefix string
}

func New(ctx context.Context, options Options) (*Backend, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	options, err := options.normalized()
	if err != nil {
		return nil, storage.NewError(storage.OpOpenBackend, BackendName, options.Prefix, storage.ErrInvalidPath, err)
	}
	loadOptions := []func(*awscfg.LoadOptions) error{
		awscfg.WithRegion(options.Region),
	}
	if options.AccessKeyID != "" {
		loadOptions = append(loadOptions, awscfg.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(options.AccessKeyID, options.SecretAccessKey, "")))
	}
	cfg, err := awscfg.LoadDefaultConfig(ctx, loadOptions...)
	if err != nil {
		return nil, storage.NewError(storage.OpOpenBackend, BackendName, options.Endpoint, storage.ErrUnknown, err)
	}
	client := awss3.NewFromConfig(cfg, func(o *awss3.Options) {
		o.BaseEndpoint = aws.String(options.Endpoint)
		o.UsePathStyle = options.ForcePathStyle
	})
	return NewWithClient(client, options)
}

func NewWithClient(client Client, options Options) (*Backend, error) {
	options, err := options.normalized()
	if err != nil {
		return nil, storage.NewError(storage.OpOpenBackend, BackendName, options.Prefix, storage.ErrInvalidPath, err)
	}
	if client == nil {
		return nil, storage.NewError(storage.OpOpenBackend, BackendName, options.Bucket, storage.ErrInvalidPath, fmt.Errorf("client is required"))
	}
	return &Backend{
		client: client,
		bucket: options.Bucket,
		prefix: options.Prefix,
	}, nil
}

func (b *Backend) ReadFile(ctx context.Context, logicalPath string) ([]byte, error) {
	reader, err := b.OpenReader(ctx, logicalPath)
	if err != nil {
		return nil, err
	}
	defer reader.Close()
	data, err := io.ReadAll(reader)
	if err != nil {
		return nil, storage.NewError(storage.OpReadFile, BackendName, logicalPath, storage.ErrUnknown, err)
	}
	return data, nil
}

func (b *Backend) OpenReader(ctx context.Context, logicalPath string) (io.ReadCloser, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	key, err := b.objectKey(logicalPath, false)
	if err != nil {
		return nil, err
	}
	output, err := b.client.GetObject(ctx, &awss3.GetObjectInput{
		Bucket: aws.String(b.bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return nil, b.translateError(storage.OpOpenReader, logicalPath, err)
	}
	return output.Body, nil
}

func (b *Backend) WriteFile(ctx context.Context, logicalPath string, data []byte, opts storage.WriteOptions) (storage.Entry, error) {
	opts.Size = int64(len(data))
	opts.SizeKnown = true
	return b.WriteFrom(ctx, logicalPath, bytes.NewReader(data), opts)
}

func (b *Backend) WriteFrom(ctx context.Context, logicalPath string, r io.Reader, opts storage.WriteOptions) (storage.Entry, error) {
	if err := ctx.Err(); err != nil {
		return storage.Entry{}, err
	}
	key, err := b.objectKey(logicalPath, false)
	if err != nil {
		return storage.Entry{}, err
	}
	if !opts.Overwrite {
		_, err := b.client.HeadObject(ctx, &awss3.HeadObjectInput{
			Bucket: aws.String(b.bucket),
			Key:    aws.String(key),
		})
		if err == nil {
			return storage.Entry{}, storage.NewError(storage.OpWriteFrom, BackendName, logicalPath, storage.ErrAlreadyExist, nil)
		}
		if !isNotFound(err) {
			return storage.Entry{}, b.translateError(storage.OpWriteFrom, logicalPath, err)
		}
	}
	data, err := io.ReadAll(r)
	if err != nil {
		return storage.Entry{}, storage.NewError(storage.OpWriteFrom, BackendName, logicalPath, storage.ErrUnknown, err)
	}
	if opts.SizeKnown && int64(len(data)) != opts.Size {
		return storage.Entry{}, storage.NewError(storage.OpWriteFrom, BackendName, logicalPath, storage.ErrConflict, fmt.Errorf("stream size %d does not match expected size %d", len(data), opts.Size))
	}
	contentType := opts.ContentType
	if contentType == "" {
		contentType = ContentType(logicalPath)
	}
	_, err = b.client.PutObject(ctx, &awss3.PutObjectInput{
		Bucket:        aws.String(b.bucket),
		Key:           aws.String(key),
		Body:          bytes.NewReader(data),
		ContentLength: aws.Int64(int64(len(data))),
		ContentType:   aws.String(contentType),
	})
	if err != nil {
		return storage.Entry{}, b.translateError(storage.OpWriteFrom, logicalPath, err)
	}
	return storage.Entry{Path: logicalPath, Type: storage.EntryTypeFile, Size: int64(len(data))}, nil
}

func (b *Backend) Stat(ctx context.Context, logicalPath string) (storage.Entry, error) {
	if err := ctx.Err(); err != nil {
		return storage.Entry{}, err
	}
	key, err := b.objectKey(logicalPath, false)
	if err != nil {
		return storage.Entry{}, err
	}
	output, err := b.client.HeadObject(ctx, &awss3.HeadObjectInput{
		Bucket: aws.String(b.bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return storage.Entry{}, b.translateError(storage.OpStat, logicalPath, err)
	}
	size := int64(0)
	if output.ContentLength != nil {
		size = *output.ContentLength
	}
	return storage.Entry{Path: logicalPath, Type: storage.EntryTypeFile, Size: size}, nil
}

func (b *Backend) Walk(ctx context.Context, logicalPrefix string, opts storage.WalkOptions, fn storage.WalkFunc) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	if err := storage.ValidatePrefix(logicalPrefix); err != nil {
		return err
	}
	emitter := storage.NewWalkEmitter(ctx, BackendName, opts, fn)
	if logicalPrefix != "" {
		entry, err := b.Stat(ctx, logicalPrefix)
		if err == nil {
			if err := emitter.Emit(entry); err != nil {
				return storage.FinishWalk(err)
			}
			if emitter.LimitReached() {
				return nil
			}
		} else if !storage.IsNotFound(err) {
			return err
		}
	}
	return storage.FinishWalk(b.walkObjects(ctx, logicalPrefix, opts, emitter.Emit))
}

func (b *Backend) HasAny(ctx context.Context, logicalPrefix string) (bool, error) {
	return storage.HasAny(ctx, b, logicalPrefix)
}

func (b *Backend) DeleteManagedBundle(ctx context.Context, bundlePath string, managedOutputPaths []string, opts storage.DeleteOptions) error {
	return b.deleteManagedTargets(ctx, storage.OpDeleteManagedBundle, func() ([]string, error) {
		return storage.ManagedBundleTargets(bundlePath, managedOutputPaths)
	}, opts)
}

func (b *Backend) DeleteManagedOutputs(ctx context.Context, bundlePath string, managedOutputPaths []string, opts storage.DeleteOptions) error {
	return b.deleteManagedTargets(ctx, storage.OpDeleteManagedOutputs, func() ([]string, error) {
		return storage.ManagedOutputTargets(bundlePath, managedOutputPaths)
	}, opts)
}

func (b *Backend) deleteManagedTargets(ctx context.Context, op string, targetsFunc func() ([]string, error), opts storage.DeleteOptions) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	targets, err := targetsFunc()
	if err != nil {
		return err
	}
	for _, target := range targets {
		if err := b.deleteObject(ctx, op, target, opts); err != nil {
			return err
		}
	}
	return nil
}

func (b *Backend) DeletePrefix(ctx context.Context, logicalPrefix string, opts storage.DeleteOptions) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	if err := storage.ValidatePrefix(logicalPrefix); err != nil {
		return err
	}
	found := false
	if logicalPrefix != "" {
		key, err := b.objectKey(logicalPrefix, false)
		if err != nil {
			return err
		}
		_, err = b.client.HeadObject(ctx, &awss3.HeadObjectInput{
			Bucket: aws.String(b.bucket),
			Key:    aws.String(key),
		})
		if err == nil {
			found = true
			if err := b.deleteObject(ctx, storage.OpDeletePrefix, logicalPrefix, storage.DeleteOptions{IgnoreMissing: false}); err != nil {
				return err
			}
		} else if !isNotFound(err) {
			return b.translateError(storage.OpDeletePrefix, logicalPrefix, err)
		}
	}
	var entries []storage.Entry
	if err := b.walkObjects(ctx, logicalPrefix, storage.WalkOptions{Recursive: true}, func(entry storage.Entry) error {
		if entry.Type == storage.EntryTypeFile {
			entries = append(entries, entry)
		}
		return nil
	}); err != nil {
		return err
	}
	for _, entry := range entries {
		found = true
		if err := b.deleteObject(ctx, storage.OpDeletePrefix, entry.Path, storage.DeleteOptions{IgnoreMissing: true}); err != nil {
			return err
		}
	}
	if !found && !opts.IgnoreMissing {
		return storage.NewError(storage.OpDeletePrefix, BackendName, logicalPrefix, storage.ErrNotFound, nil)
	}
	return nil
}

func (b *Backend) deleteObject(ctx context.Context, op, logicalPath string, opts storage.DeleteOptions) error {
	key, err := b.objectKey(logicalPath, false)
	if err != nil {
		return err
	}
	if key == b.prefix {
		return storage.NewError(op, BackendName, logicalPath, storage.ErrInvalidPath, nil)
	}
	if !opts.IgnoreMissing {
		_, err := b.client.HeadObject(ctx, &awss3.HeadObjectInput{
			Bucket: aws.String(b.bucket),
			Key:    aws.String(key),
		})
		if err != nil {
			return b.translateError(op, logicalPath, err)
		}
	}
	_, err = b.client.DeleteObject(ctx, &awss3.DeleteObjectInput{
		Bucket: aws.String(b.bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		if opts.IgnoreMissing && isNotFound(err) {
			return nil
		}
		return b.translateError(op, logicalPath, err)
	}
	return nil
}

func (b *Backend) walkObjects(ctx context.Context, logicalPrefix string, opts storage.WalkOptions, emit func(storage.Entry) error) error {
	listPrefix, err := b.listPrefix(logicalPrefix)
	if err != nil {
		return err
	}
	delimiter := ""
	if !opts.Recursive {
		delimiter = "/"
	}
	var token *string
	for {
		output, err := b.client.ListObjectsV2(ctx, &awss3.ListObjectsV2Input{
			Bucket:            aws.String(b.bucket),
			Prefix:            aws.String(listPrefix),
			Delimiter:         aws.String(delimiter),
			ContinuationToken: token,
		})
		if err != nil {
			return b.translateError(storage.OpWalk, logicalPrefix, err)
		}
		entries := entriesFromList(logicalPrefix, b.prefix, output)
		sort.Slice(entries, func(i, j int) bool { return entries[i].Path < entries[j].Path })
		for _, entry := range entries {
			if entry.Path == "" {
				continue
			}
			if err := emit(entry); err != nil {
				return err
			}
		}
		if output.IsTruncated == nil || !*output.IsTruncated {
			return nil
		}
		token = output.NextContinuationToken
	}
}

func entriesFromList(logicalPrefix, rootPrefix string, output *awss3.ListObjectsV2Output) []storage.Entry {
	seen := make(map[string]storage.Entry)
	for _, object := range output.Contents {
		if object.Key == nil {
			continue
		}
		logicalPath := logicalPathFromKey(rootPrefix, *object.Key)
		if logicalPath == "" || logicalPath == logicalPrefix {
			continue
		}
		size := int64(0)
		if object.Size != nil {
			size = *object.Size
		}
		seen[logicalPath] = storage.Entry{Path: logicalPath, Type: storage.EntryTypeFile, Size: size}
	}
	for _, commonPrefix := range output.CommonPrefixes {
		if commonPrefix.Prefix == nil {
			continue
		}
		logicalPath := strings.TrimSuffix(logicalPathFromKey(rootPrefix, *commonPrefix.Prefix), "/")
		if logicalPath == "" || logicalPath == logicalPrefix {
			continue
		}
		seen[logicalPath] = storage.Entry{Path: logicalPath, Type: storage.EntryTypeDirectory}
	}
	entries := make([]storage.Entry, 0, len(seen))
	for _, entry := range seen {
		entries = append(entries, entry)
	}
	return entries
}

func (b *Backend) objectKey(logicalPath string, allowEmpty bool) (string, error) {
	if logicalPath == "" {
		if !allowEmpty {
			return "", storage.NewError(storage.OpValidatePath, BackendName, logicalPath, storage.ErrInvalidPath, nil)
		}
		return b.prefix, nil
	}
	if err := storage.ValidatePath(logicalPath); err != nil {
		return "", err
	}
	if b.prefix == "" {
		return logicalPath, nil
	}
	return b.prefix + "/" + logicalPath, nil
}

func (b *Backend) listPrefix(logicalPrefix string) (string, error) {
	key, err := b.objectKey(logicalPrefix, true)
	if err != nil {
		return "", err
	}
	if key != "" {
		key = strings.TrimSuffix(key, "/") + "/"
	}
	return key, nil
}

func logicalPathFromKey(rootPrefix, key string) string {
	if rootPrefix == "" {
		return key
	}
	if key == rootPrefix {
		return ""
	}
	return strings.TrimPrefix(key, rootPrefix+"/")
}

func ContentType(logicalPath string) string {
	switch strings.ToLower(path.Ext(logicalPath)) {
	case ".md":
		return "text/markdown; charset=utf-8"
	case ".html":
		return "text/html; charset=utf-8"
	case ".json":
		return "application/json"
	case ".txt":
		return "text/plain; charset=utf-8"
	default:
		return "application/octet-stream"
	}
}

func isNotFound(err error) bool {
	var notFound *types.NotFound
	if errors.As(err, &notFound) {
		return true
	}
	var apiErr smithy.APIError
	if errors.As(err, &apiErr) {
		switch apiErr.ErrorCode() {
		case "NotFound", "NoSuchKey", "404":
			return true
		}
	}
	return false
}

func (b *Backend) translateError(op, logicalPath string, err error) error {
	kind := storage.ErrUnknown
	if isNotFound(err) {
		kind = storage.ErrNotFound
	} else {
		var apiErr smithy.APIError
		if errors.As(err, &apiErr) {
			switch apiErr.ErrorCode() {
			case "AccessDenied", "InvalidAccessKeyId", "SignatureDoesNotMatch":
				kind = storage.ErrPermission
			case "SlowDown", "RequestTimeout", "ServiceUnavailable", "InternalError":
				kind = storage.ErrTemporary
			case "InvalidBucketName", "NoSuchBucket":
				kind = storage.ErrInvalidPath
			}
		}
	}
	return storage.NewError(op, BackendName, logicalPath, kind, err)
}