1package integration
2
3import (
4	"context"
5	"crypto/rand"
6	"errors"
7	"fmt"
8	"io/ioutil"
9	"net/http"
10	"os"
11	"time"
12
13	"github.com/aws/aws-sdk-go-v2/aws"
14	"github.com/aws/aws-sdk-go-v2/service/s3"
15	"github.com/aws/aws-sdk-go-v2/service/s3/types"
16	smithyrand "github.com/aws/smithy-go/rand"
17)
18
19var uuid = smithyrand.NewUUID(rand.Reader)
20
21// MustUUID returns an UUID string or panics
22func MustUUID() string {
23	uuid, err := uuid.GetUUID()
24	if err != nil {
25		panic(err)
26	}
27	return uuid
28}
29
30// CreateFileOfSize will return an *os.File that is of size bytes
31func CreateFileOfSize(dir string, size int64) (*os.File, error) {
32	file, err := ioutil.TempFile(dir, "s3integration")
33	if err != nil {
34		return nil, err
35	}
36
37	err = file.Truncate(size)
38	if err != nil {
39		file.Close()
40		os.Remove(file.Name())
41		return nil, err
42	}
43
44	return file, nil
45}
46
47// SizeToName returns a human-readable string for the given size bytes
48func SizeToName(size int) string {
49	units := []string{"B", "KB", "MB", "GB"}
50	i := 0
51	for size >= 1024 {
52		size /= 1024
53		i++
54	}
55
56	if i > len(units)-1 {
57		i = len(units) - 1
58	}
59
60	return fmt.Sprintf("%d%s", size, units[i])
61}
62
63// BucketPrefix is the root prefix of integration test buckets.
64const BucketPrefix = "aws-sdk-go-v2-integration"
65
66// GenerateBucketName returns a unique bucket name.
67func GenerateBucketName() string {
68	var id [16]byte
69	_, err := rand.Read(id[:])
70	if err != nil {
71		panic(err)
72	}
73
74	return fmt.Sprintf("%s-%x",
75		BucketPrefix, id)
76}
77
78// SetupBucket returns a test bucket created for the integration tests.
79func SetupBucket(client *s3.Client, bucketName, region string) (err error) {
80	fmt.Println("Setup: Creating test bucket,", bucketName)
81	_, err = client.CreateBucket(context.Background(), &s3.CreateBucketInput{
82		Bucket: &bucketName,
83		CreateBucketConfiguration: &types.CreateBucketConfiguration{
84			LocationConstraint: types.BucketLocationConstraint(region),
85		},
86	})
87	if err != nil {
88		return fmt.Errorf("failed to create bucket %s, %v", bucketName, err)
89	}
90
91	fmt.Println("Setup: Waiting for bucket to exist,", bucketName)
92	err = waitUntilBucketExists(context.Background(), client, &s3.HeadBucketInput{Bucket: &bucketName})
93	if err != nil {
94		return fmt.Errorf("failed waiting for bucket %s to be created, %v",
95			bucketName, err)
96	}
97
98	return nil
99}
100
101func waitUntilBucketExists(ctx context.Context, client *s3.Client, params *s3.HeadBucketInput) error {
102	for i := 0; i < 20; i++ {
103		_, err := client.HeadBucket(ctx, params)
104		if err == nil {
105			return nil
106		}
107
108		var httpErr interface{ HTTPStatusCode() int }
109
110		if !errors.As(err, &httpErr) {
111			return err
112		}
113
114		if httpErr.HTTPStatusCode() == http.StatusMovedPermanently || httpErr.HTTPStatusCode() == http.StatusForbidden {
115			return nil
116		}
117
118		if httpErr.HTTPStatusCode() != http.StatusNotFound {
119			return err
120		}
121
122		time.Sleep(5 * time.Second)
123	}
124	return nil
125}
126
127// CleanupBucket deletes the contents of a S3 bucket, before deleting the bucket
128// it self.
129func CleanupBucket(client *s3.Client, bucketName string) error {
130	var errs []error
131
132	{
133		fmt.Println("TearDown: Deleting objects from test bucket,", bucketName)
134		input := &s3.ListObjectsV2Input{Bucket: &bucketName}
135		for {
136			listObjectsV2, err := client.ListObjectsV2(context.Background(), input)
137			if err != nil {
138				return fmt.Errorf("failed to list objects, %w", err)
139			}
140
141			var delete types.Delete
142			for _, content := range listObjectsV2.Contents {
143				obj := content
144				delete.Objects = append(delete.Objects, types.ObjectIdentifier{Key: obj.Key})
145			}
146
147			deleteObjects, err := client.DeleteObjects(context.Background(), &s3.DeleteObjectsInput{
148				Bucket: &bucketName,
149				Delete: &delete,
150			})
151			if err != nil {
152				errs = append(errs, err)
153				break
154			}
155			for _, deleteError := range deleteObjects.Errors {
156				errs = append(errs, fmt.Errorf("failed to delete %s, %s", aws.ToString(deleteError.Key), aws.ToString(deleteError.Message)))
157			}
158
159			if listObjectsV2.IsTruncated {
160				input.ContinuationToken = listObjectsV2.NextContinuationToken
161			} else {
162				break
163			}
164		}
165	}
166
167	{
168		fmt.Println("TearDown: Deleting partial uploads from test bucket,", bucketName)
169
170		input := &s3.ListMultipartUploadsInput{Bucket: &bucketName}
171		for {
172			uploads, err := client.ListMultipartUploads(context.Background(), input)
173			if err != nil {
174				return fmt.Errorf("failed to list multipart objects, %w", err)
175			}
176
177			for _, upload := range uploads.Uploads {
178				client.AbortMultipartUpload(context.Background(), &s3.AbortMultipartUploadInput{
179					Bucket:   &bucketName,
180					Key:      upload.Key,
181					UploadId: upload.UploadId,
182				})
183			}
184
185			if uploads.IsTruncated {
186				input.KeyMarker = uploads.NextKeyMarker
187				input.UploadIdMarker = uploads.NextUploadIdMarker
188			} else {
189				break
190			}
191		}
192	}
193
194	if len(errs) != 0 {
195		return fmt.Errorf("failed to delete objects, %s", errs)
196	}
197
198	fmt.Println("TearDown: Deleting test bucket,", bucketName)
199	if _, err := client.DeleteBucket(context.Background(), &s3.DeleteBucketInput{Bucket: &bucketName}); err != nil {
200		return fmt.Errorf("failed to delete test bucket %s, %w", bucketName, err)
201	}
202
203	return nil
204}
205