spark-server/internal/infrastructure/objectstore/ops.go

package objectstore

import (
	"bytes"
	"context"
	"errors"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/aws/smithy-go"
	smithyhttp "github.com/aws/smithy-go/transport/http"

	"spark-server/internal/pkg/errs"
)

const (
	maxGetBytes   = 32 << 20
	maxPutBytes   = 4 << 20
	maxPutLarge   = 50 << 20 // admin music audio etc.
)

// PutLarge writes an object up to maxBytes (capped at 50 MiB).
func (s *Store) PutLarge(ctx context.Context, key string, body []byte, contentType string, maxBytes int) error {
	if !s.Ready() {
		return errs.ErrAvailable
	}
	if maxBytes <= 0 || maxBytes > maxPutLarge {
		maxBytes = maxPutLarge
	}
	if len(body) == 0 || len(body) > maxBytes {
		return errs.ErrInvalidArgument
	}
	in := &s3.PutObjectInput{
		Bucket:        aws.String(s.bucket),
		Key:           aws.String(key),
		Body:          bytes.NewReader(body),
		ContentLength: aws.Int64(int64(len(body))),
	}
	if contentType != "" {
		in.ContentType = aws.String(contentType)
	}
	if _, err := s.client.PutObject(ctx, in); err != nil {
		return mapObjectErr(err)
	}
	return nil
}

func (s *Store) PresignPut(ctx context.Context, key, contentType string, size int64, ttl time.Duration) (string, error) {
	if !s.Ready() {
		return "", errs.ErrAvailable
	}
	if size < 1 {
		return "", errs.ErrInvalidArgument
	}
	in := &s3.PutObjectInput{
		Bucket:        aws.String(s.bucket),
		Key:           aws.String(key),
		ContentLength: aws.Int64(size),
		IfNoneMatch:   aws.String("*"),
	}
	if contentType != "" {
		in.ContentType = aws.String(contentType)
	}
	out, err := s.presign.PresignPutObject(ctx, in, s3.WithPresignExpires(ttl))
	if err != nil {
		return "", mapObjectErr(err)
	}
	return out.URL, nil
}

func (s *Store) PresignGet(ctx context.Context, key string, ttl time.Duration) (string, error) {
	if !s.Ready() {
		return "", errs.ErrAvailable
	}
	out, err := s.presign.PresignGetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
	}, s3.WithPresignExpires(ttl))
	if err != nil {
		return "", mapObjectErr(err)
	}
	return out.URL, nil
}

func (s *Store) Get(ctx context.Context, key string) ([]byte, error) {
	if !s.Ready() {
		return nil, errs.ErrAvailable
	}
	out, err := s.client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return nil, mapObjectErr(err)
	}
	defer func() { _ = out.Body.Close() }()
	data, err := io.ReadAll(io.LimitReader(out.Body, maxGetBytes+1))
	if err != nil {
		return nil, mapObjectErr(err)
	}
	if len(data) > maxGetBytes {
		return nil, errs.ErrInvalidArgument
	}
	return data, nil
}

func (s *Store) Put(ctx context.Context, key string, body []byte, contentType string) error {
	if !s.Ready() {
		return errs.ErrAvailable
	}
	if len(body) == 0 || len(body) > maxPutBytes {
		return errs.ErrInvalidArgument
	}
	in := &s3.PutObjectInput{
		Bucket:        aws.String(s.bucket),
		Key:           aws.String(key),
		Body:          bytes.NewReader(body),
		ContentLength: aws.Int64(int64(len(body))),
	}
	if contentType != "" {
		in.ContentType = aws.String(contentType)
	}
	if _, err := s.client.PutObject(ctx, in); err != nil {
		return mapObjectErr(err)
	}
	return nil
}

func (s *Store) Copy(ctx context.Context, src, dst string) error {
	if !s.Ready() {
		return errs.ErrAvailable
	}
	_, err := s.client.CopyObject(ctx, &s3.CopyObjectInput{
		Bucket:     aws.String(s.bucket),
		Key:        aws.String(dst),
		CopySource: aws.String(s.bucket + "/" + src),
	})
	if err != nil {
		return mapObjectErr(err)
	}
	return nil
}

func (s *Store) Delete(ctx context.Context, key string) error {
	if !s.Ready() {
		return errs.ErrAvailable
	}
	_, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return mapObjectErr(err)
	}
	return nil
}

func mapObjectErr(err error) error {
	if err == nil {
		return nil
	}
	if app, ok := errs.AsAppError(err); ok {
		return app
	}
	if isMissingObject(err) {
		return errs.ErrMediaExpired.WithError("objectstore", err)
	}
	return errs.ErrOperatFail.WithError("objectstore", err)
}

func isMissingObject(err error) bool {
	var nf *types.NotFound
	var nk *types.NoSuchKey
	var nb *types.NoSuchBucket
	if errors.As(err, &nf) || errors.As(err, &nk) || errors.As(err, &nb) {
		return true
	}
	var api smithy.APIError
	if errors.As(err, &api) {
		switch strings.ToLower(api.ErrorCode()) {
		case "nosuchkey", "nosuchobject", "notfound", "nosuchbucket":
			return true
		}
	}
	var httpErr *smithyhttp.ResponseError
	if errors.As(err, &httpErr) && httpErr.HTTPStatusCode() == http.StatusNotFound {
		return true
	}
	return false
}