Skip to content

Commit 5916e4b

Browse files
jyxjjjcodex
andcommitted
fix(s3): support multipart object copy
- Route objects larger than 5 GB through multipart copy - Preserve source metadata and abort incomplete multipart copies - Cover size routing, byte ranges, metadata, and failure cleanup Co-authored-by: Codex <[email protected]> Signed-off-by: jyxjjj <[email protected]>
1 parent 0a9e71a commit 5916e4b

3 files changed

Lines changed: 332 additions & 6 deletions

File tree

drivers/s3/driver.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -178,15 +178,15 @@ func (d *S3) Move(ctx context.Context, srcObj, dstDir model.Obj) error {
178178
}
179179

180180
func (d *S3) Rename(ctx context.Context, srcObj model.Obj, newName string) error {
181-
err := d.copy(ctx, srcObj.GetPath(), stdpath.Join(stdpath.Dir(srcObj.GetPath()), newName), srcObj.IsDir())
181+
err := d.copy(ctx, srcObj.GetPath(), stdpath.Join(stdpath.Dir(srcObj.GetPath()), newName), srcObj.GetSize(), srcObj.IsDir())
182182
if err != nil {
183183
return err
184184
}
185185
return d.Remove(ctx, srcObj)
186186
}
187187

188188
func (d *S3) Copy(ctx context.Context, srcObj, dstDir model.Obj) error {
189-
return d.copy(ctx, srcObj.GetPath(), stdpath.Join(dstDir.GetPath(), srcObj.GetName()), srcObj.IsDir())
189+
return d.copy(ctx, srcObj.GetPath(), stdpath.Join(dstDir.GetPath(), srcObj.GetName()), srcObj.GetSize(), srcObj.IsDir())
190190
}
191191

192192
func (d *S3) Remove(ctx context.Context, obj model.Obj) error {

drivers/s3/util.go

Lines changed: 115 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package s3
33
import (
44
"context"
55
"errors"
6+
"fmt"
67
"net/http"
78
"net/url"
89
"path"
@@ -19,6 +20,13 @@ import (
1920
log "github.com/sirupsen/logrus"
2021
)
2122

23+
const (
24+
maxCopyObjectSize int64 = 5 * 1000 * 1000 * 1000
25+
defaultCopyPartSize int64 = 100 * 1024 * 1024
26+
maxCopyPartSize int64 = 5 * 1024 * 1024 * 1024
27+
maxCopyParts int64 = 10000
28+
)
29+
2230
// do others that not defined in Driver interface
2331

2432
func (d *S3) initSession() error {
@@ -212,17 +220,20 @@ func (d *S3) listV2(dirPath string, args model.ListArgs) ([]model.Obj, error) {
212220
return files, nil
213221
}
214222

215-
func (d *S3) copy(ctx context.Context, src string, dst string, isDir bool) error {
223+
func (d *S3) copy(ctx context.Context, src string, dst string, size int64, isDir bool) error {
216224
if isDir {
217225
return d.copyDir(ctx, src, dst)
218226
}
219-
return d.copyFile(ctx, src, dst)
227+
return d.copyFile(ctx, src, dst, size)
220228
}
221229

222-
func (d *S3) copyFile(ctx context.Context, src string, dst string) error {
230+
func (d *S3) copyFile(ctx context.Context, src string, dst string, size int64) error {
223231
srcKey := getKey(src, false)
224232
dstKey := getKey(dst, false)
225233
encodedKey := strings.ReplaceAll(url.PathEscape(d.Bucket+"/"+srcKey), "+", "%2B")
234+
if size > maxCopyObjectSize {
235+
return d.copyFileMultipart(ctx, srcKey, dstKey, encodedKey, size)
236+
}
226237
input := &s3.CopyObjectInput{
227238
Bucket: &d.Bucket,
228239
CopySource: aws.String(encodedKey),
@@ -232,6 +243,106 @@ func (d *S3) copyFile(ctx context.Context, src string, dst string) error {
232243
return err
233244
}
234245

246+
func (d *S3) copyFileMultipart(ctx context.Context, srcKey, dstKey, encodedKey string, size int64) (err error) {
247+
head, err := d.client.HeadObjectWithContext(ctx, &s3.HeadObjectInput{
248+
Bucket: &d.Bucket,
249+
Key: &srcKey,
250+
})
251+
if err != nil {
252+
return err
253+
}
254+
if head.ContentLength != nil {
255+
size = *head.ContentLength
256+
}
257+
partSize, err := getCopyPartSize(size)
258+
if err != nil {
259+
return err
260+
}
261+
createInput := &s3.CreateMultipartUploadInput{
262+
Bucket: &d.Bucket,
263+
Key: &dstKey,
264+
CacheControl: head.CacheControl,
265+
ContentDisposition: head.ContentDisposition,
266+
ContentEncoding: head.ContentEncoding,
267+
ContentLanguage: head.ContentLanguage,
268+
ContentType: head.ContentType,
269+
Metadata: head.Metadata,
270+
WebsiteRedirectLocation: head.WebsiteRedirectLocation,
271+
}
272+
if head.Expires != nil {
273+
if expires, parseErr := http.ParseTime(*head.Expires); parseErr == nil {
274+
createInput.Expires = &expires
275+
}
276+
}
277+
created, err := d.client.CreateMultipartUploadWithContext(ctx, createInput)
278+
if err != nil {
279+
return err
280+
}
281+
uploadID := aws.StringValue(created.UploadId)
282+
if uploadID == "" {
283+
return errors.New("create multipart upload returned an empty upload ID")
284+
}
285+
completed := false
286+
defer func() {
287+
if completed {
288+
return
289+
}
290+
_, abortErr := d.client.AbortMultipartUploadWithContext(context.WithoutCancel(ctx), &s3.AbortMultipartUploadInput{
291+
Bucket: &d.Bucket,
292+
Key: &dstKey,
293+
UploadId: &uploadID,
294+
})
295+
if abortErr != nil {
296+
err = errors.Join(err, fmt.Errorf("failed to abort multipart copy: %w", abortErr))
297+
}
298+
}()
299+
300+
parts := make([]*s3.CompletedPart, 0, (size+partSize-1)/partSize)
301+
for start, partNumber := int64(0), int64(1); start < size; start, partNumber = start+partSize, partNumber+1 {
302+
end := min(start+partSize, size) - 1
303+
copied, copyErr := d.client.UploadPartCopyWithContext(ctx, &s3.UploadPartCopyInput{
304+
Bucket: &d.Bucket,
305+
CopySource: &encodedKey,
306+
CopySourceRange: aws.String(fmt.Sprintf("bytes=%d-%d", start, end)),
307+
Key: &dstKey,
308+
PartNumber: &partNumber,
309+
UploadId: &uploadID,
310+
})
311+
if copyErr != nil {
312+
return copyErr
313+
}
314+
if copied.CopyPartResult == nil || aws.StringValue(copied.CopyPartResult.ETag) == "" {
315+
return fmt.Errorf("multipart copy part %d returned an empty ETag", partNumber)
316+
}
317+
parts = append(parts, &s3.CompletedPart{
318+
ETag: copied.CopyPartResult.ETag,
319+
PartNumber: &partNumber,
320+
})
321+
}
322+
323+
_, err = d.client.CompleteMultipartUploadWithContext(ctx, &s3.CompleteMultipartUploadInput{
324+
Bucket: &d.Bucket,
325+
Key: &dstKey,
326+
UploadId: &uploadID,
327+
MultipartUpload: &s3.CompletedMultipartUpload{
328+
Parts: parts,
329+
},
330+
})
331+
if err != nil {
332+
return err
333+
}
334+
completed = true
335+
return nil
336+
}
337+
338+
func getCopyPartSize(size int64) (int64, error) {
339+
partSize := max(defaultCopyPartSize, (size-1)/maxCopyParts+1)
340+
if partSize > maxCopyPartSize {
341+
return 0, fmt.Errorf("object size %d exceeds multipart copy limit", size)
342+
}
343+
return partSize, nil
344+
}
345+
235346
func (d *S3) copyDir(ctx context.Context, src string, dst string) error {
236347
objs, err := op.List(ctx, d, src, model.ListArgs{S3ShowPlaceholder: true})
237348
if err != nil {
@@ -243,7 +354,7 @@ func (d *S3) copyDir(ctx context.Context, src string, dst string) error {
243354
if obj.IsDir() {
244355
err = d.copyDir(ctx, cSrc, cDst)
245356
} else {
246-
err = d.copyFile(ctx, cSrc, cDst)
357+
err = d.copyFile(ctx, cSrc, cDst, obj.GetSize())
247358
}
248359
if err != nil {
249360
return err

drivers/s3/util_test.go

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
package s3
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"io"
7+
"net/http"
8+
"net/http/httptest"
9+
"strconv"
10+
"strings"
11+
"testing"
12+
13+
"github.com/aws/aws-sdk-go/aws"
14+
"github.com/aws/aws-sdk-go/aws/credentials"
15+
"github.com/aws/aws-sdk-go/aws/session"
16+
awss3 "github.com/aws/aws-sdk-go/service/s3"
17+
)
18+
19+
func TestCopyFileUsesCopyObjectAtLimit(t *testing.T) {
20+
copyRequests := 0
21+
d := newTestS3Driver(t, func(w http.ResponseWriter, r *http.Request) {
22+
if r.Method != http.MethodPut || r.URL.Query().Get("uploadId") != "" {
23+
t.Errorf("unexpected request: %s %s", r.Method, r.URL.String())
24+
w.WriteHeader(http.StatusBadRequest)
25+
return
26+
}
27+
copyRequests++
28+
writeTestXML(t, w, `<CopyObjectResult><ETag>"copy"</ETag></CopyObjectResult>`)
29+
})
30+
31+
if err := d.copyFile(context.Background(), "source+file", "destination", maxCopyObjectSize); err != nil {
32+
t.Fatalf("copyFile: %v", err)
33+
}
34+
if copyRequests != 1 {
35+
t.Fatalf("copy requests = %d, want 1", copyRequests)
36+
}
37+
}
38+
39+
func TestCopyFileUsesMultipartCopyAboveLimit(t *testing.T) {
40+
size := maxCopyObjectSize + 1
41+
wantParts := int((size + defaultCopyPartSize - 1) / defaultCopyPartSize)
42+
ranges := make(map[int]string, wantParts)
43+
completed := false
44+
aborted := false
45+
46+
d := newTestS3Driver(t, func(w http.ResponseWriter, r *http.Request) {
47+
switch {
48+
case r.Method == http.MethodHead:
49+
w.Header().Set("Content-Length", strconv.FormatInt(size, 10))
50+
w.Header().Set("Content-Type", "application/octet-stream")
51+
w.Header().Set("Cache-Control", "max-age=60")
52+
w.Header().Set("Content-Disposition", "attachment")
53+
w.Header().Set("Expires", "Wed, 21 Oct 2015 07:28:00 GMT")
54+
w.Header().Set("X-Amz-Meta-Source", "preserved")
55+
w.Header().Set("X-Amz-Website-Redirect-Location", "/redirect")
56+
w.WriteHeader(http.StatusOK)
57+
case r.Method == http.MethodPost && r.URL.Query().Has("uploads"):
58+
if got := r.Header.Get("Cache-Control"); got != "max-age=60" {
59+
t.Errorf("Cache-Control = %q, want %q", got, "max-age=60")
60+
}
61+
if got := r.Header.Get("Content-Disposition"); got != "attachment" {
62+
t.Errorf("Content-Disposition = %q, want %q", got, "attachment")
63+
}
64+
if got := r.Header.Get("Content-Type"); got != "application/octet-stream" {
65+
t.Errorf("Content-Type = %q, want %q", got, "application/octet-stream")
66+
}
67+
if got := r.Header.Get("Expires"); got != "Wed, 21 Oct 2015 07:28:00 GMT" {
68+
t.Errorf("Expires = %q, want an unchanged HTTP date", got)
69+
}
70+
if got := r.Header.Get("X-Amz-Meta-Source"); got != "preserved" {
71+
t.Errorf("metadata = %q, want %q", got, "preserved")
72+
}
73+
if got := r.Header.Get("X-Amz-Website-Redirect-Location"); got != "/redirect" {
74+
t.Errorf("website redirect = %q, want %q", got, "/redirect")
75+
}
76+
writeTestXML(t, w, `<InitiateMultipartUploadResult><UploadId>upload-id</UploadId></InitiateMultipartUploadResult>`)
77+
case r.Method == http.MethodPut && r.URL.Query().Get("uploadId") == "upload-id":
78+
partNumber, err := strconv.Atoi(r.URL.Query().Get("partNumber"))
79+
if err != nil {
80+
t.Errorf("invalid part number: %v", err)
81+
w.WriteHeader(http.StatusBadRequest)
82+
return
83+
}
84+
if got := r.Header.Get("X-Amz-Copy-Source"); !strings.Contains(got, "source%2Bfile") {
85+
t.Errorf("copy source = %q, want encoded source key", got)
86+
}
87+
ranges[partNumber] = r.Header.Get("X-Amz-Copy-Source-Range")
88+
writeTestXML(t, w, fmt.Sprintf(`<CopyPartResult><ETag>"part-%d"</ETag></CopyPartResult>`, partNumber))
89+
case r.Method == http.MethodPost && r.URL.Query().Get("uploadId") == "upload-id":
90+
body, err := io.ReadAll(r.Body)
91+
if err != nil {
92+
t.Errorf("read complete body: %v", err)
93+
}
94+
if got := strings.Count(string(body), "<Part>"); got != wantParts {
95+
t.Errorf("completed parts = %d, want %d", got, wantParts)
96+
}
97+
completed = true
98+
writeTestXML(t, w, `<CompleteMultipartUploadResult><ETag>"complete"</ETag></CompleteMultipartUploadResult>`)
99+
case r.Method == http.MethodDelete && r.URL.Query().Get("uploadId") == "upload-id":
100+
aborted = true
101+
w.WriteHeader(http.StatusNoContent)
102+
default:
103+
t.Errorf("unexpected request: %s %s", r.Method, r.URL.String())
104+
w.WriteHeader(http.StatusBadRequest)
105+
}
106+
})
107+
108+
if err := d.copyFile(context.Background(), "source+file", "destination", size); err != nil {
109+
t.Fatalf("copyFile: %v", err)
110+
}
111+
if !completed {
112+
t.Fatal("multipart upload was not completed")
113+
}
114+
if aborted {
115+
t.Fatal("successful multipart upload was aborted")
116+
}
117+
if len(ranges) != wantParts {
118+
t.Fatalf("copied parts = %d, want %d", len(ranges), wantParts)
119+
}
120+
if got := ranges[1]; got != fmt.Sprintf("bytes=0-%d", defaultCopyPartSize-1) {
121+
t.Errorf("first range = %q", got)
122+
}
123+
lastStart := int64(wantParts-1) * defaultCopyPartSize
124+
if got := ranges[wantParts]; got != fmt.Sprintf("bytes=%d-%d", lastStart, size-1) {
125+
t.Errorf("last range = %q", got)
126+
}
127+
}
128+
129+
func TestCopyFileMultipartAbortsOnPartFailure(t *testing.T) {
130+
size := maxCopyObjectSize + 1
131+
aborted := false
132+
completed := false
133+
134+
d := newTestS3Driver(t, func(w http.ResponseWriter, r *http.Request) {
135+
switch {
136+
case r.Method == http.MethodHead:
137+
w.Header().Set("Content-Length", strconv.FormatInt(size, 10))
138+
w.WriteHeader(http.StatusOK)
139+
case r.Method == http.MethodPost && r.URL.Query().Has("uploads"):
140+
writeTestXML(t, w, `<InitiateMultipartUploadResult><UploadId>upload-id</UploadId></InitiateMultipartUploadResult>`)
141+
case r.Method == http.MethodPut && r.URL.Query().Get("uploadId") == "upload-id":
142+
w.WriteHeader(http.StatusInternalServerError)
143+
writeTestXML(t, w, `<Error><Code>InternalError</Code><Message>copy failed</Message></Error>`)
144+
case r.Method == http.MethodDelete && r.URL.Query().Get("uploadId") == "upload-id":
145+
aborted = true
146+
w.WriteHeader(http.StatusNoContent)
147+
case r.Method == http.MethodPost && r.URL.Query().Get("uploadId") == "upload-id":
148+
completed = true
149+
w.WriteHeader(http.StatusOK)
150+
default:
151+
t.Errorf("unexpected request: %s %s", r.Method, r.URL.String())
152+
w.WriteHeader(http.StatusBadRequest)
153+
}
154+
})
155+
156+
if err := d.copyFile(context.Background(), "source", "destination", size); err == nil {
157+
t.Fatal("copyFile returned nil error")
158+
}
159+
if !aborted {
160+
t.Fatal("failed multipart upload was not aborted")
161+
}
162+
if completed {
163+
t.Fatal("failed multipart upload was completed")
164+
}
165+
}
166+
167+
func TestGetCopyPartSize(t *testing.T) {
168+
partSize, err := getCopyPartSize(defaultCopyPartSize * maxCopyParts)
169+
if err != nil {
170+
t.Fatalf("getCopyPartSize: %v", err)
171+
}
172+
if partSize != defaultCopyPartSize {
173+
t.Fatalf("part size = %d, want %d", partSize, defaultCopyPartSize)
174+
}
175+
176+
partSize, err = getCopyPartSize(defaultCopyPartSize*maxCopyParts + 1)
177+
if err != nil {
178+
t.Fatalf("getCopyPartSize: %v", err)
179+
}
180+
if partSize != defaultCopyPartSize+1 {
181+
t.Fatalf("grown part size = %d, want %d", partSize, defaultCopyPartSize+1)
182+
}
183+
184+
if _, err := getCopyPartSize(maxCopyPartSize*maxCopyParts + 1); err == nil {
185+
t.Fatal("getCopyPartSize returned nil error for an oversized object")
186+
}
187+
}
188+
189+
func newTestS3Driver(t *testing.T, handler http.HandlerFunc) *S3 {
190+
t.Helper()
191+
server := httptest.NewServer(handler)
192+
t.Cleanup(server.Close)
193+
sess, err := session.NewSession(&aws.Config{
194+
Credentials: credentials.NewStaticCredentials("access-key", "secret-key", ""),
195+
Endpoint: aws.String(server.URL),
196+
Region: aws.String("us-east-1"),
197+
S3ForcePathStyle: aws.Bool(true),
198+
MaxRetries: aws.Int(0),
199+
})
200+
if err != nil {
201+
t.Fatalf("create AWS session: %v", err)
202+
}
203+
return &S3{
204+
Addition: Addition{Bucket: "bucket"},
205+
client: awss3.New(sess),
206+
}
207+
}
208+
209+
func writeTestXML(t *testing.T, w http.ResponseWriter, body string) {
210+
t.Helper()
211+
w.Header().Set("Content-Type", "application/xml")
212+
if _, err := io.WriteString(w, body); err != nil {
213+
t.Errorf("write response: %v", err)
214+
}
215+
}

0 commit comments

Comments
 (0)