Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e8a0b0bf10 | ||
|
|
4c8e95fcd8 |
No files matched your search
@@ -665,6 +665,23 @@ func (p *StoragePlugin) GetObject(req *pb.GetObjectRequest, stream pb.StoragePlu
|
||||
return nil
|
||||
}
|
||||
|
||||
const (
|
||||
// multipartThreshold is the file size, in bytes, above which PutObject
|
||||
// uploads an object using a multipart upload instead of a single PutObject
|
||||
// request. 8 MiB matches the AWS CLI and boto3 defaults.
|
||||
multipartThreshold = 8 * 1024 * 1024 // 8 MiB
|
||||
|
||||
// defaultMultipartPartSize is the default size, in bytes, of each part
|
||||
// uploaded during a multipart upload. It is kept equal to
|
||||
// multipartThreshold so a multipart upload always has at least two parts
|
||||
// (S3 requires every part except the last to be at least 5 MiB).
|
||||
defaultMultipartPartSize = 8 * 1024 * 1024 // 8 MiB
|
||||
|
||||
// maxMultipartParts is the maximum number of parts S3 permits for a single
|
||||
// multipart upload.
|
||||
maxMultipartParts = 10000
|
||||
)
|
||||
|
||||
// PutObject is called when putting objects into an s3 bucket.
|
||||
func (p *StoragePlugin) PutObject(ctx context.Context, req *pb.PutObjectRequest) (*pb.PutObjectResponse, error) {
|
||||
if req == nil {
|
||||
@@ -751,6 +768,32 @@ func (p *StoragePlugin) PutObject(ctx context.Context, req *pb.PutObjectRequest)
|
||||
}
|
||||
objectKey := path.Join(bucket.GetBucketPrefix(), req.GetKey())
|
||||
|
||||
// decodedChecksum is the raw SHA-256 digest of the entire object. Boundary
|
||||
// compares this byte-for-byte against the hash it computed locally before
|
||||
// the upload, so it must always be the whole-object checksum regardless of
|
||||
// whether a single PutObject or a multipart upload is used.
|
||||
decodedChecksum, err := base64.StdEncoding.DecodeString(checksum)
|
||||
if err != nil {
|
||||
return nil, errors.UnknownStatus("failed to decode checksum value")
|
||||
}
|
||||
|
||||
// for larger files, use a multipart upload instead of a single PutObject
|
||||
// request.
|
||||
if info.Size() > multipartThreshold {
|
||||
partSize := int64(defaultMultipartPartSize)
|
||||
// S3 allows at most maxMultipartParts parts per upload, so grow the
|
||||
// part size for very large objects to stay within that limit.
|
||||
if sz := info.Size(); sz/partSize >= maxMultipartParts {
|
||||
partSize = (sz + maxMultipartParts - 1) / maxMultipartParts
|
||||
}
|
||||
if err := p.putObjectMultipart(ctx, req, storageState, bucket.GetBucketName(), objectKey, file, partSize, opts...); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &pb.PutObjectResponse{
|
||||
ChecksumSha_256: decodedChecksum,
|
||||
}, nil
|
||||
}
|
||||
|
||||
putCall := func(s3Client S3API) (any, error) {
|
||||
return s3Client.PutObject(ctx, &s3.PutObjectInput{
|
||||
Bucket: aws.String(bucket.GetBucketName()),
|
||||
@@ -772,15 +815,126 @@ func (p *StoragePlugin) PutObject(ctx context.Context, req *pb.PutObjectRequest)
|
||||
if checksum != *resp.ChecksumSHA256 {
|
||||
return nil, checksumMistmatchStatus()
|
||||
}
|
||||
decodedChecksum, err := base64.StdEncoding.DecodeString(*resp.ChecksumSHA256)
|
||||
if err != nil {
|
||||
return nil, errors.UnknownStatus("failed to decode checksum value from aws")
|
||||
}
|
||||
return &pb.PutObjectResponse{
|
||||
ChecksumSha_256: decodedChecksum,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// putObjectMultipart uploads file to the given bucket and key using an S3
|
||||
// multipart upload. The file is split into parts of partSize bytes (the final
|
||||
// part may be smaller) and each part is uploaded with its own SHA-256 checksum
|
||||
// so S3 can verify part integrity on the wire.
|
||||
//
|
||||
// If any step fails the upload is aborted on a best-effort basis to avoid
|
||||
// leaving orphaned parts behind, which would otherwise continue to accrue
|
||||
// storage charges.
|
||||
func (p *StoragePlugin) putObjectMultipart(
|
||||
ctx context.Context,
|
||||
req *pb.PutObjectRequest,
|
||||
storageState *awsStoragePersistedState,
|
||||
bucketName, objectKey string,
|
||||
file *os.File,
|
||||
partSize int64,
|
||||
opts ...s3Option,
|
||||
) error {
|
||||
bucketId := req.GetBucket().GetId()
|
||||
|
||||
createCall := func(s3Client S3API) (any, error) {
|
||||
return s3Client.CreateMultipartUpload(ctx, &s3.CreateMultipartUploadInput{
|
||||
Bucket: aws.String(bucketName),
|
||||
Key: aws.String(objectKey),
|
||||
ChecksumAlgorithm: types.ChecksumAlgorithmSha256,
|
||||
})
|
||||
}
|
||||
createResp, err := p.call(ctx, createCall, bucketId, storageState, opts...)
|
||||
if err != nil {
|
||||
return parseS3Error("create multipart upload", err, req).Err()
|
||||
}
|
||||
uploadId := createResp.(*s3.CreateMultipartUploadOutput).UploadId
|
||||
|
||||
// once the upload has been created, any subsequent failure must abort it so
|
||||
// uploaded parts are not orphaned.
|
||||
abort := func() {
|
||||
abortCall := func(s3Client S3API) (any, error) {
|
||||
return s3Client.AbortMultipartUpload(ctx, &s3.AbortMultipartUploadInput{
|
||||
Bucket: aws.String(bucketName),
|
||||
Key: aws.String(objectKey),
|
||||
UploadId: uploadId,
|
||||
})
|
||||
}
|
||||
// best effort: the failure that triggered the abort is more useful to
|
||||
// the caller than an error from the abort itself.
|
||||
_, _ = p.call(ctx, abortCall, bucketId, storageState, opts...)
|
||||
}
|
||||
|
||||
var completedParts []types.CompletedPart
|
||||
buf := make([]byte, partSize)
|
||||
var partNumber int32 = 1
|
||||
for {
|
||||
n, readErr := io.ReadFull(file, buf)
|
||||
if readErr == io.EOF {
|
||||
// previous iteration consumed the final bytes of the file
|
||||
break
|
||||
}
|
||||
if readErr != nil && readErr != io.ErrUnexpectedEOF {
|
||||
abort()
|
||||
return errors.UnknownStatusf("failed to read file: %s", readErr)
|
||||
}
|
||||
|
||||
partData := buf[:n]
|
||||
sum := sha256.Sum256(partData)
|
||||
partChecksum := base64.StdEncoding.EncodeToString(sum[:])
|
||||
|
||||
pn := partNumber
|
||||
uploadCall := func(s3Client S3API) (any, error) {
|
||||
return s3Client.UploadPart(ctx, &s3.UploadPartInput{
|
||||
Bucket: aws.String(bucketName),
|
||||
Key: aws.String(objectKey),
|
||||
UploadId: uploadId,
|
||||
PartNumber: aws.Int32(pn),
|
||||
Body: bytes.NewReader(partData),
|
||||
ContentLength: aws.Int64(int64(n)),
|
||||
ChecksumAlgorithm: types.ChecksumAlgorithmSha256,
|
||||
ChecksumSHA256: aws.String(partChecksum),
|
||||
})
|
||||
}
|
||||
uploadResp, err := p.call(ctx, uploadCall, bucketId, storageState, opts...)
|
||||
if err != nil {
|
||||
abort()
|
||||
return parseS3Error("upload part", err, req).Err()
|
||||
}
|
||||
part := uploadResp.(*s3.UploadPartOutput)
|
||||
completedParts = append(completedParts, types.CompletedPart{
|
||||
ETag: part.ETag,
|
||||
ChecksumSHA256: part.ChecksumSHA256,
|
||||
PartNumber: aws.Int32(pn),
|
||||
})
|
||||
|
||||
partNumber++
|
||||
if readErr == io.ErrUnexpectedEOF {
|
||||
// the final, short part
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
completeCall := func(s3Client S3API) (any, error) {
|
||||
return s3Client.CompleteMultipartUpload(ctx, &s3.CompleteMultipartUploadInput{
|
||||
Bucket: aws.String(bucketName),
|
||||
Key: aws.String(objectKey),
|
||||
UploadId: uploadId,
|
||||
MultipartUpload: &types.CompletedMultipartUpload{
|
||||
Parts: completedParts,
|
||||
},
|
||||
})
|
||||
}
|
||||
if _, err := p.call(ctx, completeCall, bucketId, storageState, opts...); err != nil {
|
||||
abort()
|
||||
return parseS3Error("complete multipart upload", err, req).Err()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteObjects is used to delete one or many objects from an s3 bucket.
|
||||
func (p *StoragePlugin) DeleteObjects(ctx context.Context, req *pb.DeleteObjectsRequest) (*pb.DeleteObjectsResponse, error) {
|
||||
if req.GetKeyPrefix() == "" {
|
||||
|
||||
@@ -3715,6 +3715,195 @@ func TestStoragePlugin_PutObject(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestStoragePlugin_PutObjectMultipart(t *testing.T) {
|
||||
td := t.TempDir()
|
||||
|
||||
// genFile writes size bytes of deterministic content to a temp file and
|
||||
// returns its path along with the raw SHA-256 digest of the whole file.
|
||||
genFile := func(t *testing.T, name string, size int) (string, []byte) {
|
||||
t.Helper()
|
||||
data := make([]byte, size)
|
||||
for i := range data {
|
||||
data[i] = byte(i % 251)
|
||||
}
|
||||
fp := path.Join(td, name)
|
||||
require.NoError(t, os.WriteFile(fp, data, 0o600))
|
||||
sum := sha256.Sum256(data)
|
||||
return fp, sum[:]
|
||||
}
|
||||
|
||||
validRequest := func(p string) *pb.PutObjectRequest {
|
||||
return &pb.PutObjectRequest{
|
||||
Bucket: &storagebuckets.StorageBucket{
|
||||
BucketName: "external-obj-store",
|
||||
BucketPrefix: "prefix",
|
||||
Secrets: credential.MockStaticCredentialSecrets(),
|
||||
Attributes: &structpb.Struct{
|
||||
Fields: map[string]*structpb.Value{
|
||||
credential.ConstRegion: structpb.NewStringValue("us-west-2"),
|
||||
},
|
||||
},
|
||||
},
|
||||
Key: "mock-object",
|
||||
Path: p,
|
||||
}
|
||||
}
|
||||
|
||||
t.Run("multipart success spanning multiple parts", func(t *testing.T) {
|
||||
require, assert := require.New(t), assert.New(t)
|
||||
// 20 MiB with 8 MiB parts -> 3 parts (8, 8, 4 MiB).
|
||||
fp, wholeChecksum := genFile(t, "large-file", 20*1024*1024)
|
||||
|
||||
state := &testMockS3State{}
|
||||
p := &StoragePlugin{
|
||||
testCredStateOpts: validSTSMock(),
|
||||
testStorageStateOpts: []awsStoragePersistedStateOption{
|
||||
withTestS3APIFunc(newTestMockS3(state)),
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := p.PutObject(context.Background(), validRequest(fp))
|
||||
require.NoError(err)
|
||||
require.NotNil(resp)
|
||||
|
||||
// Boundary must receive the whole-object SHA-256, not the multipart
|
||||
// composite checksum.
|
||||
assert.True(bytes.Equal(wholeChecksum, resp.GetChecksumSha_256()))
|
||||
|
||||
// The single-shot PutObject path must not have been used.
|
||||
assert.False(state.PutObjectCalled)
|
||||
assert.True(state.CreateMultipartUploadCalled)
|
||||
assert.Equal(3, state.UploadPartCallCount)
|
||||
assert.True(state.CompleteMultipartUploadCalled)
|
||||
assert.False(state.AbortMultipartUploadCalled)
|
||||
|
||||
// The reassembled part bodies must equal the original file content.
|
||||
orig, err := os.ReadFile(fp)
|
||||
require.NoError(err)
|
||||
assert.True(bytes.Equal(orig, state.UploadPartBody))
|
||||
|
||||
// The object key must include the configured bucket prefix.
|
||||
require.NotNil(state.CreateMultipartUploadInputParams)
|
||||
assert.Equal("prefix/mock-object", aws.ToString(state.CreateMultipartUploadInputParams.Key))
|
||||
})
|
||||
|
||||
t.Run("file just above threshold uses two parts", func(t *testing.T) {
|
||||
require, assert := require.New(t), assert.New(t)
|
||||
// 8 MiB + 1 byte -> 2 parts (8 MiB, 1 byte).
|
||||
fp, wholeChecksum := genFile(t, "just-over", 8*1024*1024+1)
|
||||
|
||||
state := &testMockS3State{}
|
||||
p := &StoragePlugin{
|
||||
testCredStateOpts: validSTSMock(),
|
||||
testStorageStateOpts: []awsStoragePersistedStateOption{
|
||||
withTestS3APIFunc(newTestMockS3(state)),
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := p.PutObject(context.Background(), validRequest(fp))
|
||||
require.NoError(err)
|
||||
assert.True(bytes.Equal(wholeChecksum, resp.GetChecksumSha_256()))
|
||||
assert.Equal(2, state.UploadPartCallCount)
|
||||
assert.True(state.CompleteMultipartUploadCalled)
|
||||
})
|
||||
|
||||
t.Run("file at threshold uses single PutObject", func(t *testing.T) {
|
||||
require, assert := require.New(t), assert.New(t)
|
||||
// Exactly 8 MiB is not above the threshold, so the single-shot path is used.
|
||||
fp, wholeChecksum := genFile(t, "at-threshold", 8*1024*1024)
|
||||
checksum := base64.StdEncoding.EncodeToString(wholeChecksum)
|
||||
|
||||
state := &testMockS3State{}
|
||||
p := &StoragePlugin{
|
||||
testCredStateOpts: validSTSMock(),
|
||||
testStorageStateOpts: []awsStoragePersistedStateOption{
|
||||
withTestS3APIFunc(newTestMockS3(
|
||||
state,
|
||||
testMockS3WithPutObjectOutput(&s3.PutObjectOutput{
|
||||
ChecksumSHA256: aws.String(checksum),
|
||||
}),
|
||||
)),
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := p.PutObject(context.Background(), validRequest(fp))
|
||||
require.NoError(err)
|
||||
assert.True(bytes.Equal(wholeChecksum, resp.GetChecksumSha_256()))
|
||||
assert.True(state.PutObjectCalled)
|
||||
assert.False(state.CreateMultipartUploadCalled)
|
||||
})
|
||||
|
||||
t.Run("create multipart upload error", func(t *testing.T) {
|
||||
require, assert := require.New(t), assert.New(t)
|
||||
fp, _ := genFile(t, "create-err", 20*1024*1024)
|
||||
|
||||
state := &testMockS3State{}
|
||||
p := &StoragePlugin{
|
||||
testCredStateOpts: validSTSMock(),
|
||||
testStorageStateOpts: []awsStoragePersistedStateOption{
|
||||
withTestS3APIFunc(newTestMockS3(
|
||||
state,
|
||||
testMockS3WithCreateMultipartUploadError(TestAwsS3Error("AccessDenied", "CreateMultipartUpload", "not authorized")),
|
||||
)),
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := p.PutObject(context.Background(), validRequest(fp))
|
||||
require.Error(err)
|
||||
assert.Nil(resp)
|
||||
assert.Equal(codes.PermissionDenied, status.Code(err))
|
||||
// Nothing was created, so nothing should be aborted.
|
||||
assert.False(state.AbortMultipartUploadCalled)
|
||||
})
|
||||
|
||||
t.Run("upload part error aborts upload", func(t *testing.T) {
|
||||
require, assert := require.New(t), assert.New(t)
|
||||
fp, _ := genFile(t, "upload-err", 20*1024*1024)
|
||||
|
||||
state := &testMockS3State{}
|
||||
p := &StoragePlugin{
|
||||
testCredStateOpts: validSTSMock(),
|
||||
testStorageStateOpts: []awsStoragePersistedStateOption{
|
||||
withTestS3APIFunc(newTestMockS3(
|
||||
state,
|
||||
testMockS3WithUploadPartErrorOnPart(2, TestAwsS3Error("AccessDenied", "UploadPart", "not authorized")),
|
||||
)),
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := p.PutObject(context.Background(), validRequest(fp))
|
||||
require.Error(err)
|
||||
assert.Nil(resp)
|
||||
assert.Equal(codes.PermissionDenied, status.Code(err))
|
||||
assert.True(state.CreateMultipartUploadCalled)
|
||||
assert.False(state.CompleteMultipartUploadCalled)
|
||||
// The failed upload must be aborted to avoid orphaned parts.
|
||||
assert.True(state.AbortMultipartUploadCalled)
|
||||
})
|
||||
|
||||
t.Run("complete multipart upload error aborts upload", func(t *testing.T) {
|
||||
require, assert := require.New(t), assert.New(t)
|
||||
fp, _ := genFile(t, "complete-err", 20*1024*1024)
|
||||
|
||||
state := &testMockS3State{}
|
||||
p := &StoragePlugin{
|
||||
testCredStateOpts: validSTSMock(),
|
||||
testStorageStateOpts: []awsStoragePersistedStateOption{
|
||||
withTestS3APIFunc(newTestMockS3(
|
||||
state,
|
||||
testMockS3WithCompleteMultipartUploadError(TestAwsS3Error("InternalError", "CompleteMultipartUpload", "boom")),
|
||||
)),
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := p.PutObject(context.Background(), validRequest(fp))
|
||||
require.Error(err)
|
||||
assert.Nil(resp)
|
||||
assert.Equal(3, state.UploadPartCallCount)
|
||||
assert.True(state.AbortMultipartUploadCalled)
|
||||
})
|
||||
}
|
||||
|
||||
func TestStoragePlugin_DeleteObjects(t *testing.T) {
|
||||
validRequest := func() *pb.DeleteObjectsRequest {
|
||||
return &pb.DeleteObjectsRequest{
|
||||
|
||||
@@ -26,6 +26,10 @@ type S3API interface {
|
||||
GetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error)
|
||||
HeadObject(ctx context.Context, params *s3.HeadObjectInput, optFns ...func(*s3.Options)) (*s3.HeadObjectOutput, error)
|
||||
PutObject(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error)
|
||||
CreateMultipartUpload(ctx context.Context, params *s3.CreateMultipartUploadInput, optFns ...func(*s3.Options)) (*s3.CreateMultipartUploadOutput, error)
|
||||
UploadPart(ctx context.Context, params *s3.UploadPartInput, optFns ...func(*s3.Options)) (*s3.UploadPartOutput, error)
|
||||
CompleteMultipartUpload(ctx context.Context, params *s3.CompleteMultipartUploadInput, optFns ...func(*s3.Options)) (*s3.CompleteMultipartUploadOutput, error)
|
||||
AbortMultipartUpload(ctx context.Context, params *s3.AbortMultipartUploadInput, optFns ...func(*s3.Options)) (*s3.AbortMultipartUploadOutput, error)
|
||||
DeleteObjects(ctx context.Context, params *s3.DeleteObjectsInput, optFns ...func(*s3.Options)) (*s3.DeleteObjectsOutput, error)
|
||||
ListObjectsV2(ctx context.Context, params *s3.ListObjectsV2Input, optFns ...func(*s3.Options)) (*s3.ListObjectsV2Output, error)
|
||||
Credentials() aws.Credentials
|
||||
|
||||
@@ -6,6 +6,8 @@ package storage
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"io"
|
||||
"testing"
|
||||
@@ -39,6 +41,23 @@ type testMockS3State struct {
|
||||
PutObjectInputParams *s3.PutObjectInput
|
||||
PutObjectBody []byte
|
||||
|
||||
CreateMultipartUploadCalled bool
|
||||
CreateMultipartUploadInputParams *s3.CreateMultipartUploadInput
|
||||
|
||||
UploadPartCalled bool
|
||||
UploadPartCallCount int
|
||||
UploadPartInputParams []*s3.UploadPartInput
|
||||
// UploadPartBody is the concatenation, in part-number order, of every
|
||||
// part body received by UploadPart. It can be compared against the full
|
||||
// object that was expected to be uploaded.
|
||||
UploadPartBody []byte
|
||||
|
||||
CompleteMultipartUploadCalled bool
|
||||
CompleteMultipartUploadInputParams *s3.CompleteMultipartUploadInput
|
||||
|
||||
AbortMultipartUploadCalled bool
|
||||
AbortMultipartUploadInputParams *s3.AbortMultipartUploadInput
|
||||
|
||||
HeadObjectCalled bool
|
||||
HeadObjectInputParams *s3.HeadObjectInput
|
||||
|
||||
@@ -57,6 +76,16 @@ func (s *testMockS3State) Reset() {
|
||||
s.PutObjectCalled = false
|
||||
s.PutObjectInputParams = nil
|
||||
s.PutObjectBody = nil
|
||||
s.CreateMultipartUploadCalled = false
|
||||
s.CreateMultipartUploadInputParams = nil
|
||||
s.UploadPartCalled = false
|
||||
s.UploadPartCallCount = 0
|
||||
s.UploadPartInputParams = nil
|
||||
s.UploadPartBody = nil
|
||||
s.CompleteMultipartUploadCalled = false
|
||||
s.CompleteMultipartUploadInputParams = nil
|
||||
s.AbortMultipartUploadCalled = false
|
||||
s.AbortMultipartUploadInputParams = nil
|
||||
s.HeadObjectCalled = false
|
||||
s.HeadObjectInputParams = nil
|
||||
s.ListObjectsV2Called = false
|
||||
@@ -81,6 +110,19 @@ type testMockS3 struct {
|
||||
PutObjectOutput *s3.PutObjectOutput
|
||||
PutObjectErr error
|
||||
|
||||
// mocked responses for the multipart upload calls
|
||||
CreateMultipartUploadOutput *s3.CreateMultipartUploadOutput
|
||||
CreateMultipartUploadErr error
|
||||
UploadPartOutput *s3.UploadPartOutput
|
||||
UploadPartErr error
|
||||
// UploadPartErrOnPart, when greater than zero, causes UploadPart to return
|
||||
// UploadPartErr (or a generic error) only for the matching part number.
|
||||
UploadPartErrOnPart int32
|
||||
CompleteMultipartUploadOutput *s3.CompleteMultipartUploadOutput
|
||||
CompleteMultipartUploadErr error
|
||||
AbortMultipartUploadOutput *s3.AbortMultipartUploadOutput
|
||||
AbortMultipartUploadErr error
|
||||
|
||||
// mocked responses for headObject
|
||||
HeadObjectOutput *s3.HeadObjectOutput
|
||||
HeadObjectErr error
|
||||
@@ -114,6 +156,63 @@ func testMockS3WithPutObjectError(e error) testMockS3Option {
|
||||
}
|
||||
}
|
||||
|
||||
func testMockS3WithCreateMultipartUploadOutput(o *s3.CreateMultipartUploadOutput) testMockS3Option {
|
||||
return func(m *testMockS3) error {
|
||||
m.CreateMultipartUploadOutput = o
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func testMockS3WithCreateMultipartUploadError(e error) testMockS3Option {
|
||||
return func(m *testMockS3) error {
|
||||
m.CreateMultipartUploadErr = e
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func testMockS3WithUploadPartOutput(o *s3.UploadPartOutput) testMockS3Option {
|
||||
return func(m *testMockS3) error {
|
||||
m.UploadPartOutput = o
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func testMockS3WithUploadPartError(e error) testMockS3Option {
|
||||
return func(m *testMockS3) error {
|
||||
m.UploadPartErr = e
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func testMockS3WithUploadPartErrorOnPart(partNumber int32, e error) testMockS3Option {
|
||||
return func(m *testMockS3) error {
|
||||
m.UploadPartErr = e
|
||||
m.UploadPartErrOnPart = partNumber
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func testMockS3WithCompleteMultipartUploadOutput(o *s3.CompleteMultipartUploadOutput) testMockS3Option {
|
||||
return func(m *testMockS3) error {
|
||||
m.CompleteMultipartUploadOutput = o
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func testMockS3WithCompleteMultipartUploadError(e error) testMockS3Option {
|
||||
return func(m *testMockS3) error {
|
||||
m.CompleteMultipartUploadErr = e
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func testMockS3WithAbortMultipartUploadError(e error) testMockS3Option {
|
||||
return func(m *testMockS3) error {
|
||||
m.AbortMultipartUploadErr = e
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func testMockS3WithGetObjectOutput(o *s3.GetObjectOutput) testMockS3Option {
|
||||
return func(m *testMockS3) error {
|
||||
m.GetObjectOutput = o
|
||||
@@ -255,6 +354,95 @@ func (m *testMockS3) PutObject(ctx context.Context, params *s3.PutObjectInput, o
|
||||
return m.PutObjectOutput, nil
|
||||
}
|
||||
|
||||
func (m *testMockS3) CreateMultipartUpload(ctx context.Context, params *s3.CreateMultipartUploadInput, optFns ...func(*s3.Options)) (*s3.CreateMultipartUploadOutput, error) {
|
||||
if m.State != nil {
|
||||
m.State.CreateMultipartUploadCalled = true
|
||||
m.State.CreateMultipartUploadInputParams = params
|
||||
}
|
||||
|
||||
if m.CreateMultipartUploadErr != nil {
|
||||
return nil, m.CreateMultipartUploadErr
|
||||
}
|
||||
|
||||
if m.CreateMultipartUploadOutput != nil {
|
||||
return m.CreateMultipartUploadOutput, nil
|
||||
}
|
||||
return &s3.CreateMultipartUploadOutput{
|
||||
UploadId: aws.String("test-upload-id"),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (m *testMockS3) UploadPart(ctx context.Context, params *s3.UploadPartInput, optFns ...func(*s3.Options)) (*s3.UploadPartOutput, error) {
|
||||
if m.State != nil {
|
||||
m.State.UploadPartCalled = true
|
||||
m.State.UploadPartCallCount++
|
||||
m.State.UploadPartInputParams = append(m.State.UploadPartInputParams, params)
|
||||
}
|
||||
|
||||
// always drain the body so the part data can be asserted on and so the
|
||||
// reader is not left dangling.
|
||||
var data []byte
|
||||
if params.Body != nil {
|
||||
var err error
|
||||
data, err = io.ReadAll(params.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if m.State != nil {
|
||||
m.State.UploadPartBody = append(m.State.UploadPartBody, data...)
|
||||
}
|
||||
}
|
||||
|
||||
if m.UploadPartErr != nil {
|
||||
if m.UploadPartErrOnPart == 0 || (params.PartNumber != nil && *params.PartNumber == m.UploadPartErrOnPart) {
|
||||
return nil, m.UploadPartErr
|
||||
}
|
||||
}
|
||||
|
||||
if m.UploadPartOutput != nil {
|
||||
return m.UploadPartOutput, nil
|
||||
}
|
||||
// echo the supplied part checksum so CompleteMultipartUpload can be built
|
||||
// the way the real service would allow.
|
||||
sum := sha256.Sum256(data)
|
||||
return &s3.UploadPartOutput{
|
||||
ETag: aws.String(fmt.Sprintf("etag-%d", aws.ToInt32(params.PartNumber))),
|
||||
ChecksumSHA256: aws.String(base64.StdEncoding.EncodeToString(sum[:])),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (m *testMockS3) CompleteMultipartUpload(ctx context.Context, params *s3.CompleteMultipartUploadInput, optFns ...func(*s3.Options)) (*s3.CompleteMultipartUploadOutput, error) {
|
||||
if m.State != nil {
|
||||
m.State.CompleteMultipartUploadCalled = true
|
||||
m.State.CompleteMultipartUploadInputParams = params
|
||||
}
|
||||
|
||||
if m.CompleteMultipartUploadErr != nil {
|
||||
return nil, m.CompleteMultipartUploadErr
|
||||
}
|
||||
|
||||
if m.CompleteMultipartUploadOutput != nil {
|
||||
return m.CompleteMultipartUploadOutput, nil
|
||||
}
|
||||
return &s3.CompleteMultipartUploadOutput{}, nil
|
||||
}
|
||||
|
||||
func (m *testMockS3) AbortMultipartUpload(ctx context.Context, params *s3.AbortMultipartUploadInput, optFns ...func(*s3.Options)) (*s3.AbortMultipartUploadOutput, error) {
|
||||
if m.State != nil {
|
||||
m.State.AbortMultipartUploadCalled = true
|
||||
m.State.AbortMultipartUploadInputParams = params
|
||||
}
|
||||
|
||||
if m.AbortMultipartUploadErr != nil {
|
||||
return nil, m.AbortMultipartUploadErr
|
||||
}
|
||||
|
||||
if m.AbortMultipartUploadOutput != nil {
|
||||
return m.AbortMultipartUploadOutput, nil
|
||||
}
|
||||
return &s3.AbortMultipartUploadOutput{}, nil
|
||||
}
|
||||
|
||||
func (m *testMockS3) ListObjectsV2(ctx context.Context, params *s3.ListObjectsV2Input, optFns ...func(*s3.Options)) (*s3.ListObjectsV2Output, error) {
|
||||
if m.State != nil {
|
||||
m.State.ListObjectsV2Called = true
|
||||
|
||||
Reference in new issue
Block a user