forked from keybase/client
-
Notifications
You must be signed in to change notification settings - Fork 0
/
s3.go
398 lines (350 loc) · 11.7 KB
/
s3.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
package attachments
import (
"bytes"
"context"
"crypto/md5"
"encoding/hex"
"errors"
"fmt"
"io"
"sync"
"github.com/keybase/client/go/chat/s3"
"github.com/keybase/client/go/chat/types"
"github.com/keybase/client/go/protocol/chat1"
"golang.org/x/sync/errgroup"
)
const s3PipelineMaxWidth = 10
type s3UploadPipeliner struct {
sync.Mutex
width int
waiters []chan struct{}
}
func (s *s3UploadPipeliner) QueueForTakeoff(ctx context.Context) error {
s.Lock()
if s.width >= s3PipelineMaxWidth {
ch := make(chan struct{})
s.waiters = append(s.waiters, ch)
s.Unlock()
select {
case <-ch:
case <-ctx.Done():
return ctx.Err()
}
s.Lock()
s.width++
s.Unlock()
return nil
}
s.width++
s.Unlock()
return nil
}
func (s *s3UploadPipeliner) Complete() {
s.Lock()
defer s.Unlock()
if len(s.waiters) > 0 {
close(s.waiters[0])
if len(s.waiters) > 1 {
s.waiters = s.waiters[1:]
} else {
s.waiters = nil
}
}
if s.width > 0 {
s.width--
}
}
var s3UploadPipeline = &s3UploadPipeliner{}
const minMultiSize = 5 * 1024 * 1024 // can't use Multi API with parts less than 5MB
const blockSize = 5 * 1024 * 1024 // 5MB is the minimum Multi part size
// ErrAbortOnPartMismatch is returned when there is a mismatch between a current
// part and a previous attempt part. If ErrAbortOnPartMismatch is returned,
// the caller should abort the upload attempt and start from scratch.
var ErrAbortOnPartMismatch = errors.New("local part mismatch, aborting upload")
// PutS3Result is the success result of calling PutS3.
type PutS3Result struct {
Region string
Endpoint string
Bucket string
Path string
Size int64
}
// PutS3 uploads the data in Reader r to S3. It chooses whether to use
// putSingle or putMultiPipeline based on the size of the object.
func (a *S3Store) PutS3(ctx context.Context, r io.Reader, size int64, task *UploadTask, previous *AttachmentInfo) (res *PutS3Result, err error) {
defer a.Trace(ctx, func() error { return err }, "PutS3")()
region := a.regionFromParams(task.S3Params)
b := a.s3Conn(task.S3Signer, region, task.S3Params.AccessKey).Bucket(task.S3Params.Bucket)
multiPartUpload := size > minMultiSize
if multiPartUpload && a.G().Env.GetAttachmentDisableMulti() {
a.Debug(ctx, "PutS3: multi part upload manually disabled, overriding for size: %v", size)
multiPartUpload = false
}
if !multiPartUpload {
if err := a.putSingle(ctx, r, size, task.S3Params, b, task.Progress); err != nil {
return nil, err
}
} else {
objectKey, err := a.putMultiPipeline(ctx, r, size, task, b, previous)
if err != nil {
return nil, err
}
task.S3Params.ObjectKey = objectKey
}
s3res := PutS3Result{
Region: task.S3Params.RegionName,
Endpoint: task.S3Params.RegionEndpoint,
Bucket: task.S3Params.Bucket,
Path: task.S3Params.ObjectKey,
Size: size,
}
return &s3res, nil
}
// putSingle uploads data in r to S3 with the Put API. It has to be
// used for anything less than 5MB. It can be used for anything up
// to 5GB, but putMultiPipeline best for anything over 5MB.
func (a *S3Store) putSingle(ctx context.Context, r io.Reader, size int64, params chat1.S3Params,
b s3.BucketInt, progress types.ProgressReporter) (err error) {
defer a.Trace(ctx, func() error { return err }, fmt.Sprintf("putSingle(size=%d)", size))()
progWriter := newProgressWriter(progress, size)
tee := io.TeeReader(r, progWriter)
if err := b.PutReader(ctx, params.ObjectKey, tee, size, "application/octet-stream", s3.ACL(params.Acl),
s3.Options{}); err != nil {
a.Debug(ctx, "putSingle: failed: %s", err)
return NewErrorWrapper("failed putSingle", err)
}
progWriter.Finish()
return nil
}
// putMultiPipeline uploads data in r to S3 using the Multi API. It uses a
// pipeline to upload 10 blocks of data concurrently.
// Each block is 5MB. It returns the object key if no errors. putMultiPipeline
// will return a different object key from params.ObjectKey if a previous Put is
// successfully resumed and completed.
func (a *S3Store) putMultiPipeline(ctx context.Context, r io.Reader, size int64, task *UploadTask, b s3.BucketInt, previous *AttachmentInfo) (res string, err error) {
defer a.Trace(ctx, func() error { return err }, fmt.Sprintf("putMultiPipeline(size=%d)", size))()
var multi s3.MultiInt
if previous != nil {
a.Debug(ctx, "putMultiPipeline: previous exists. Changing object key from %q to %q",
task.S3Params.ObjectKey, previous.ObjectKey)
task.S3Params.ObjectKey = previous.ObjectKey
}
multi, err = b.Multi(ctx, task.S3Params.ObjectKey, "application/octet-stream", s3.ACL(task.S3Params.Acl))
if err != nil {
a.Debug(ctx, "putMultiPipeline: b.Multi error: %s", err.Error())
return "", NewErrorWrapper("s3 Multi error", err)
}
var previousParts map[int]s3.Part
if previous != nil {
previousParts = make(map[int]s3.Part)
list, err := multi.ListParts(ctx)
if err != nil {
a.Debug(ctx, "putMultiPipeline: ignoring multi.ListParts error: %s", err)
// dump previous since we can't check it anymore
previous = nil
} else {
for _, p := range list {
previousParts[p.N] = p
}
}
}
// need to use ectx in everything in eg.Go() funcs since eg
// will cancel ectx in eg.Wait().
a.Debug(ctx, "putMultiPipeline: beginning parts uploader process")
eg, ectx := errgroup.WithContext(ctx)
blockCh := make(chan job)
retCh := make(chan s3.Part)
eg.Go(func() error {
defer close(blockCh)
return a.makeBlockJobs(ectx, r, blockCh, task.stashKey(), previous)
})
eg.Go(func() error {
for lb := range blockCh {
if err := s3UploadPipeline.QueueForTakeoff(ectx); err != nil {
return err
}
b := lb
eg.Go(func() error {
defer s3UploadPipeline.Complete()
if err := a.uploadPart(ectx, task, b, previous, previousParts, multi, retCh); err != nil {
return err
}
return nil
})
}
return nil
})
go func() {
eg.Wait()
close(retCh)
}()
var parts []s3.Part
progWriter := newProgressWriter(task.Progress, size)
for p := range retCh {
parts = append(parts, p)
progWriter.Update(int(p.Size))
}
if err := eg.Wait(); err != nil {
return "", err
}
if a.blockLimit > 0 {
return "", errors.New("block limit hit, not completing multi upload")
}
a.Debug(ctx, "putMultiPipeline: all parts uploaded, completing request")
if err := multi.Complete(ctx, parts); err != nil {
a.Debug(ctx, "putMultiPipeline: Complete() failed: %s", err)
return "", err
}
a.Debug(ctx, "putMultiPipeline: success, %d parts", len(parts))
// Just to make sure the UI gets the 100% call
progWriter.Finish()
return task.S3Params.ObjectKey, nil
}
type job struct {
block []byte
index int
hash string
}
func (j job) etag() string {
return `"` + j.hash + `"`
}
// makeBlockJobs reads ciphertext chunks from r and creates jobs that it puts onto blockCh.
// If this is a resumed upload, it verifies the blocks against the local stash before
// creating jobs.
func (a *S3Store) makeBlockJobs(ctx context.Context, r io.Reader, blockCh chan job, stashKey StashKey, previous *AttachmentInfo) error {
var partNumber int
for {
partNumber++
block := make([]byte, blockSize)
// must call io.ReadFull to ensure full block read
n, err := io.ReadFull(r, block)
// io.ErrUnexpectedEOF will be returned for last partial block,
// which is ok.
if err != nil && err != io.ErrUnexpectedEOF && err != io.EOF {
return err
}
if n < blockSize {
block = block[:n]
}
if n > 0 {
md5sum := md5.Sum(block)
md5hex := hex.EncodeToString(md5sum[:])
if previous != nil {
// resuming an upload, so check local stash record
// and abort on mismatch before adding a job for this block
// because if we don't it amounts to nonce reuse
lhash, found := previous.Parts[partNumber]
if found && lhash != md5hex {
a.Debug(ctx, "makeBlockJobs: part %d failed local part record verification", partNumber)
return ErrAbortOnPartMismatch
}
}
if err := a.addJob(ctx, blockCh, block, partNumber, md5hex); err != nil {
return err
}
}
if err == io.EOF || err == io.ErrUnexpectedEOF {
break
}
if a.blockLimit > 0 && partNumber >= a.blockLimit {
a.Debug(ctx, "makeBlockJobs: hit blockLimit of %d", a.blockLimit)
break
}
}
return nil
}
// addJob creates a job and puts it on blockCh, unless the blockCh isn't ready and the context has been canceled.
func (a *S3Store) addJob(ctx context.Context, blockCh chan job, block []byte, partNumber int, hash string) error {
// Create a job, unless the context has been canceled.
select {
case blockCh <- job{block: block, index: partNumber, hash: hash}:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
// uploadPart handles uploading a job to S3. The job `b` has already passed local stash verification.
// If this is a resumed upload, it checks the previous parts reported by S3 and will skip uploading
// any that already exist.
func (a *S3Store) uploadPart(ctx context.Context, task *UploadTask, b job, previous *AttachmentInfo, previousParts map[int]s3.Part, multi s3.MultiInt, retCh chan s3.Part) (err error) {
defer a.Trace(ctx, func() error { return err }, fmt.Sprintf("uploadPart(%d)", b.index))()
// check to see if this part has already been uploaded.
// for job `b` to be here, it has already passed local stash verification.
if previous != nil {
// check s3 previousParts for this block
p, ok := previousParts[b.index]
if ok && int(p.Size) == len(b.block) && p.ETag == b.etag() {
a.Debug(ctx, "uploadPart: part %d already uploaded to s3", b.index)
// part already uploaded, so put it in the retCh unless the context
// has been canceled
select {
case retCh <- p:
case <-ctx.Done():
return ctx.Err()
}
// nothing else to do
return nil
}
if p.Size > 0 {
// only abort if the part size from s3 is > 0.
a.Debug(ctx, "uploadPart: part %d s3 mismatch: size %d != expected %d or etag %s != expected %s",
b.index, p.Size, len(b.block), p.ETag, b.etag())
return ErrAbortOnPartMismatch
}
// this part doesn't exist on s3, so it needs to be uploaded
a.Debug(ctx, "uploadPart: part %d not uploaded to s3 by previous upload attempt", b.index)
}
// stash part info locally before attempting S3 put
// doing this before attempting the S3 put is important
// for security concerns.
if err := a.stash.RecordPart(task.stashKey(), b.index, b.hash); err != nil {
a.Debug(ctx, "uploadPart: StashRecordPart error: %s", err)
}
part, putErr := multi.PutPart(ctx, b.index, bytes.NewReader(b.block))
if putErr != nil {
return NewErrorWrapper(fmt.Sprintf("failed to put part %d", b.index), putErr)
}
// put the successfully uploaded part information in the retCh
// unless the context has been canceled.
select {
case retCh <- part:
case <-ctx.Done():
a.Debug(ctx, "uploadPart: upload part %d, context canceled", b.index)
return ctx.Err()
}
return nil
}
type ErrorWrapper struct {
prefix string
err error
}
func NewErrorWrapper(prefix string, err error) *ErrorWrapper {
return &ErrorWrapper{prefix: prefix, err: err}
}
func (e *ErrorWrapper) Error() string {
return fmt.Sprintf("%s: %s (%T)", e.prefix, e.err, e.err)
}
func (e *ErrorWrapper) Details() string {
switch err := e.err.(type) {
case *s3.Error:
return fmt.Sprintf("%s: error %q, status code: %d, code: %s, message: %s, bucket: %s", e.prefix, e.err, err.StatusCode, err.Code, err.Message, err.BucketName)
default:
return fmt.Sprintf("%s: error %q, no details for type %T", e.prefix, e.err, e.err)
}
}
type S3Signer struct {
ri func() chat1.RemoteInterface
}
func NewS3Signer(ri func() chat1.RemoteInterface) *S3Signer {
return &S3Signer{
ri: ri,
}
}
// Sign implements github.com/keybase/go/chat/s3.Signer interface.
func (s *S3Signer) Sign(payload []byte) ([]byte, error) {
arg := chat1.S3SignArg{
Payload: payload,
Version: 1,
}
return s.ri().S3Sign(context.Background(), arg)
}