diff options
Diffstat (limited to 'workhorse/internal/objectstore/test/s3_stub.go')
-rw-r--r-- | workhorse/internal/objectstore/test/s3_stub.go | 142 |
1 files changed, 142 insertions, 0 deletions
diff --git a/workhorse/internal/objectstore/test/s3_stub.go b/workhorse/internal/objectstore/test/s3_stub.go new file mode 100644 index 00000000000..36514b3b887 --- /dev/null +++ b/workhorse/internal/objectstore/test/s3_stub.go @@ -0,0 +1,142 @@ +package test + +import ( + "io/ioutil" + "net/http/httptest" + "os" + "strings" + "testing" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/aws/credentials" + "github.com/aws/aws-sdk-go/aws/session" + "github.com/stretchr/testify/require" + + "gitlab.com/gitlab-org/gitlab-workhorse/internal/config" + + "github.com/aws/aws-sdk-go/service/s3" + "github.com/aws/aws-sdk-go/service/s3/s3manager" + + "github.com/johannesboyne/gofakes3" + "github.com/johannesboyne/gofakes3/backend/s3mem" +) + +func SetupS3(t *testing.T, encryption string) (config.S3Credentials, config.S3Config, *session.Session, *httptest.Server) { + return SetupS3WithBucket(t, "test-bucket", encryption) +} + +func SetupS3WithBucket(t *testing.T, bucket string, encryption string) (config.S3Credentials, config.S3Config, *session.Session, *httptest.Server) { + backend := s3mem.New() + faker := gofakes3.New(backend) + ts := httptest.NewServer(faker.Server()) + + creds := config.S3Credentials{ + AwsAccessKeyID: "YOUR-ACCESSKEYID", + AwsSecretAccessKey: "YOUR-SECRETACCESSKEY", + } + + config := config.S3Config{ + Bucket: bucket, + Endpoint: ts.URL, + Region: "eu-central-1", + PathStyle: true, + } + + if encryption != "" { + config.ServerSideEncryption = encryption + + if encryption == s3.ServerSideEncryptionAwsKms { + config.SSEKMSKeyID = "arn:aws:1234" + } + } + + sess, err := session.NewSession(&aws.Config{ + Credentials: credentials.NewStaticCredentials(creds.AwsAccessKeyID, creds.AwsSecretAccessKey, ""), + Endpoint: aws.String(ts.URL), + Region: aws.String(config.Region), + DisableSSL: aws.Bool(true), + S3ForcePathStyle: aws.Bool(true), + }) + require.NoError(t, err) + + // Create S3 service client + svc := s3.New(sess) + + _, err = svc.CreateBucket(&s3.CreateBucketInput{ + Bucket: aws.String(bucket), + }) + require.NoError(t, err) + + return creds, config, sess, ts +} + +// S3ObjectExists will fail the test if the file does not exist. +func S3ObjectExists(t *testing.T, sess *session.Session, config config.S3Config, objectName string, expectedBytes string) { + downloadObject(t, sess, config, objectName, func(tmpfile *os.File, numBytes int64, err error) { + require.NoError(t, err) + require.Equal(t, int64(len(expectedBytes)), numBytes) + + output, err := ioutil.ReadFile(tmpfile.Name()) + require.NoError(t, err) + + require.Equal(t, []byte(expectedBytes), output) + }) +} + +func CheckS3Metadata(t *testing.T, sess *session.Session, config config.S3Config, objectName string) { + // In a real S3 provider, s3crypto.NewDecryptionClient should probably be used + svc := s3.New(sess) + result, err := svc.GetObject(&s3.GetObjectInput{ + Bucket: aws.String(config.Bucket), + Key: aws.String(objectName), + }) + require.NoError(t, err) + + if config.ServerSideEncryption != "" { + require.Equal(t, aws.String(config.ServerSideEncryption), result.ServerSideEncryption) + + if config.ServerSideEncryption == s3.ServerSideEncryptionAwsKms { + require.Equal(t, aws.String(config.SSEKMSKeyID), result.SSEKMSKeyId) + } else { + require.Nil(t, result.SSEKMSKeyId) + } + } else { + require.Nil(t, result.ServerSideEncryption) + require.Nil(t, result.SSEKMSKeyId) + } +} + +// S3ObjectDoesNotExist returns true if the object has been deleted, +// false otherwise. The return signature is different from +// S3ObjectExists because deletion may need to be retried since deferred +// clean up callsinternal/objectstore/test/s3_stub.go may cause the actual deletion to happen after the +// initial check. +func S3ObjectDoesNotExist(t *testing.T, sess *session.Session, config config.S3Config, objectName string) bool { + deleted := false + + downloadObject(t, sess, config, objectName, func(tmpfile *os.File, numBytes int64, err error) { + if err != nil && strings.Contains(err.Error(), "NoSuchKey") { + deleted = true + } + }) + + return deleted +} + +func downloadObject(t *testing.T, sess *session.Session, config config.S3Config, objectName string, handler func(tmpfile *os.File, numBytes int64, err error)) { + tmpDir, err := ioutil.TempDir("", "workhorse-test-") + require.NoError(t, err) + defer os.Remove(tmpDir) + + tmpfile, err := ioutil.TempFile(tmpDir, "s3-output") + require.NoError(t, err) + defer os.Remove(tmpfile.Name()) + + downloadSvc := s3manager.NewDownloader(sess) + numBytes, err := downloadSvc.Download(tmpfile, &s3.GetObjectInput{ + Bucket: aws.String(config.Bucket), + Key: aws.String(objectName), + }) + + handler(tmpfile, numBytes, err) +} |