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 ) 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 }