diff --git a/pkg/cloud/awscloud/awscloud.go b/pkg/cloud/awscloud/awscloud.go index bab00c1dd5..b76966517b 100644 --- a/pkg/cloud/awscloud/awscloud.go +++ b/pkg/cloud/awscloud/awscloud.go @@ -243,6 +243,8 @@ func WaitUntilImportSnapshotTaskCompletedWithContext(c *ec2.EC2, ctx aws.Context // The caller can also specify the name of the role used to do the import. // If nil is given, the default one from the SDK is used (vmimport). // Returns the image ID and the snapshot ID. +// +// XXX: make this return (string, string, error) instead of pointers func (a *AWS) Register(name, bucket, key string, shareWith []string, rpmArch string, bootMode, importRole *string) (*string, *string, error) { rpmArchToEC2Arch := map[string]string{ "x86_64": "x86_64", @@ -294,11 +296,7 @@ func (a *AWS) Register(name, bucket, key string, shareWith []string, rpmArch str // we no longer need the object in s3, let's just delete it logrus.Infof("[AWS] 🧹 Deleting image from S3: %s/%s", bucket, key) - _, err = a.s3.DeleteObject(&s3.DeleteObjectInput{ - Bucket: aws.String(bucket), - Key: aws.String(key), - }) - if err != nil { + if err = a.DeleteObject(bucket, key); err != nil { return nil, nil, err } @@ -388,6 +386,14 @@ func (a *AWS) Register(name, bucket, key string, shareWith []string, rpmArch str return registerOutput.ImageId, snapshotID, nil } +func (a *AWS) DeleteObject(bucket, key string) error { + _, err := a.s3.DeleteObject(&s3.DeleteObjectInput{ + Bucket: aws.String(bucket), + Key: aws.String(key), + }) + return err +} + // target region is determined by the region configured in the aws session func (a *AWS) CopyImage(name, ami, sourceRegion string) (string, error) { result, err := a.ec2.CopyImage( diff --git a/pkg/cloud/awscloud/export_test.go b/pkg/cloud/awscloud/export_test.go new file mode 100644 index 0000000000..381d5fe3b4 --- /dev/null +++ b/pkg/cloud/awscloud/export_test.go @@ -0,0 +1,11 @@ +package awscloud + +type AwsClient = awsClient + +func MockNewAwsClient(f func(string) (awsClient, error)) (restore func()) { + saved := newAwsClient + newAwsClient = f + return func() { + newAwsClient = saved + } +} diff --git a/pkg/cloud/awscloud/uploader.go b/pkg/cloud/awscloud/uploader.go new file mode 100644 index 0000000000..9a4ca59ef6 --- /dev/null +++ b/pkg/cloud/awscloud/uploader.go @@ -0,0 +1,134 @@ +package awscloud + +import ( + "errors" + "fmt" + "io" + "slices" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/aws/aws-sdk-go/service/s3/s3manager" + "github.com/google/uuid" + + "github.com/osbuild/images/pkg/arch" + "github.com/osbuild/images/pkg/cloud" +) + +type awsUploader struct { + client awsClient + + region string + bucketName string + imageName string + targetArch string +} + +type UploaderOptions struct { + TargetArch string +} + +// testing support +type awsClient interface { + Regions() ([]string, error) + Buckets() ([]string, error) + CheckBucketPermission(string, S3Permission) (bool, error) + UploadFromReader(io.Reader, string, string) (*s3manager.UploadOutput, error) + Register(name, bucket, key string, shareWith []string, rpmArch string, bootMode, importRole *string) (*string, *string, error) + DeleteObject(string, string) error +} + +var newAwsClient = func(region string) (awsClient, error) { + return NewDefault(region) +} + +func NewUploader(region, bucketName, imageName string, opts *UploaderOptions) (cloud.Uploader, error) { + if opts == nil { + opts = &UploaderOptions{} + } + client, err := newAwsClient(region) + if err != nil { + return nil, err + } + + return &awsUploader{ + client: client, + region: region, + bucketName: bucketName, + imageName: imageName, + targetArch: opts.TargetArch, + }, nil +} + +var _ cloud.Uploader = &awsUploader{} + +func (au *awsUploader) Check(status io.Writer) error { + fmt.Fprintf(status, "Checking AWS region access...\n") + regions, err := au.client.Regions() + if err != nil { + return fmt.Errorf("retrieving AWS regions for '%s' failed: %w", au.region, err) + } + + if !slices.Contains(regions, au.region) { + return fmt.Errorf("given AWS region '%s' not found", au.region) + } + + fmt.Fprintf(status, "Checking AWS bucket...\n") + buckets, err := au.client.Buckets() + if err != nil { + return fmt.Errorf("retrieving AWS list of buckets failed: %w", err) + } + if !slices.Contains(buckets, au.bucketName) { + return fmt.Errorf("bucket '%s' not found in the given AWS account", au.bucketName) + } + + fmt.Fprintf(status, "Checking AWS bucket permissions...\n") + writePermission, err := au.client.CheckBucketPermission(au.bucketName, S3PermissionWrite) + if err != nil { + return err + } + if !writePermission { + return fmt.Errorf("you don't have write permissions to bucket '%s' with the given AWS account", au.bucketName) + } + fmt.Fprintf(status, "Upload conditions met.\n") + return nil +} + +func (au *awsUploader) UploadAndRegister(r io.Reader, status io.Writer) (err error) { + keyName := fmt.Sprintf("%s-%s", uuid.New().String(), au.imageName) + fmt.Fprintf(status, "Uploading %s to %s:%s\n", au.imageName, au.bucketName, keyName) + + res, err := au.client.UploadFromReader(r, au.bucketName, keyName) + if err != nil { + return err + } + defer func() { + if err != nil { + aErr := au.client.DeleteObject(au.bucketName, keyName) + fmt.Fprintf(status, "Deleted S3 object %s:%s\n", au.bucketName, keyName) + err = errors.Join(err, aErr) + } + }() + fmt.Fprintf(status, "File uploaded to %s\n", aws.StringValue(&res.Location)) + if au.targetArch == "" { + au.targetArch = arch.Current().String() + } + bootMode := ec2.BootModeValuesUefiPreferred + + fmt.Fprintf(status, "Registering AMI %s\n", au.imageName) + ami, snapshot, err := au.client.Register(au.imageName, au.bucketName, keyName, nil, au.targetArch, &bootMode, nil) + if err != nil { + return err + } + + fmt.Fprintf(status, "Deleted S3 object %s:%s\n", au.bucketName, keyName) + if err := au.client.DeleteObject(au.bucketName, keyName); err != nil { + return err + } + fmt.Fprintf(status, "AMI registered: %s\nSnapshot ID: %s\n", aws.StringValue(ami), aws.StringValue(snapshot)) + if err != nil { + return err + } + + return nil +} diff --git a/pkg/cloud/awscloud/uploader_test.go b/pkg/cloud/awscloud/uploader_test.go new file mode 100644 index 0000000000..184a55e4c8 --- /dev/null +++ b/pkg/cloud/awscloud/uploader_test.go @@ -0,0 +1,195 @@ +package awscloud_test + +import ( + "bytes" + "fmt" + "io" + "testing" + + "github.com/aws/aws-sdk-go/service/s3/s3manager" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + + "github.com/osbuild/images/pkg/cloud/awscloud" +) + +// XXX: put into a new "cloudtest" package? +type fakeAWSClient struct { + regions []string + regionsErr error + regionsCalls int + + buckets []string + bucketsErr error + bucketsCalls int + + checkBucketPermission bool + checkBucketPermissionErr error + checkBucketPermissionCalls int + + uploadFromReader *s3manager.UploadOutput + uploadFromReaderErr error + uploadFromReaderCalls int + + registerErr error + registerImageId string + registerSnapshotId string + registerCalls int + + deleteObjectErr error + deleteObjectCalls int +} + +func (fa *fakeAWSClient) Regions() ([]string, error) { + fa.regionsCalls++ + return fa.regions, fa.regionsErr +} + +func (fa *fakeAWSClient) Buckets() ([]string, error) { + fa.bucketsCalls++ + return fa.buckets, fa.bucketsErr +} + +func (fa *fakeAWSClient) CheckBucketPermission(string, awscloud.S3Permission) (bool, error) { + fa.checkBucketPermissionCalls++ + return fa.checkBucketPermission, fa.checkBucketPermissionErr +} + +func (fa *fakeAWSClient) UploadFromReader(io.Reader, string, string) (*s3manager.UploadOutput, error) { + fa.uploadFromReaderCalls++ + return fa.uploadFromReader, fa.uploadFromReaderErr +} + +func (fa *fakeAWSClient) Register(name, bucket, key string, shareWith []string, rpmArch string, bootMode, importRole *string) (*string, *string, error) { + fa.registerCalls++ + return &fa.registerImageId, &fa.registerSnapshotId, fa.registerErr +} + +func (fa *fakeAWSClient) DeleteObject(string, string) error { + fa.deleteObjectCalls++ + return fa.deleteObjectErr +} + +func TestUploaderCheckHappy(t *testing.T) { + fa := &fakeAWSClient{ + regions: []string{"region"}, + buckets: []string{"bucket"}, + checkBucketPermission: true, + } + restore := awscloud.MockNewAwsClient(func(string) (awscloud.AwsClient, error) { + return fa, nil + }) + defer restore() + + uploader, err := awscloud.NewUploader("region", "bucket", "ami", nil) + assert.NoError(t, err) + var statusLog bytes.Buffer + err = uploader.Check(&statusLog) + assert.NoError(t, err) + assert.Equal(t, 1, fa.regionsCalls) + assert.Equal(t, 1, fa.bucketsCalls) + assert.Equal(t, 1, fa.checkBucketPermissionCalls) + expectedStatusLog := `Checking AWS region access... +Checking AWS bucket... +Checking AWS bucket permissions... +Upload conditions met. +` + assert.Equal(t, expectedStatusLog, statusLog.String()) +} + +type repeatReader struct{} + +func (r *repeatReader) Read(p []byte) (int, error) { + for i := range p { + p[i] = 0x1 + } + return len(p), nil +} + +func TestUploaderUploadHappy(t *testing.T) { + uuid.SetRand(&repeatReader{}) + + fa := &fakeAWSClient{ + uploadFromReader: &s3manager.UploadOutput{ + Location: "some-location", + }, + registerImageId: "image-id", + registerSnapshotId: "snapshot-id", + } + restore := awscloud.MockNewAwsClient(func(string) (awscloud.AwsClient, error) { + return fa, nil + }) + defer restore() + + fakeImage := bytes.NewBufferString("fake-aws-image") + uploader, err := awscloud.NewUploader("region", "bucket", "ami", nil) + assert.NoError(t, err) + var uploadLog bytes.Buffer + err = uploader.UploadAndRegister(fakeImage, &uploadLog) + assert.NoError(t, err) + assert.Equal(t, 1, fa.uploadFromReaderCalls) + assert.Equal(t, 1, fa.registerCalls) + assert.Equal(t, 1, fa.deleteObjectCalls) + expectedUploadLog := `Uploading ami to bucket:01010101-0101-4101-8101-010101010101-ami +File uploaded to some-location +Registering AMI ami +Deleted S3 object bucket:01010101-0101-4101-8101-010101010101-ami +AMI registered: image-id +Snapshot ID: snapshot-id +` + assert.Equal(t, expectedUploadLog, uploadLog.String()) +} + +func TestUploaderUploadButRegisterError(t *testing.T) { + uuid.SetRand(&repeatReader{}) + + fa := &fakeAWSClient{ + uploadFromReader: &s3manager.UploadOutput{ + Location: "some-location", + }, + registerErr: fmt.Errorf("fake-register-err"), + } + restore := awscloud.MockNewAwsClient(func(string) (awscloud.AwsClient, error) { + return fa, nil + }) + defer restore() + + fakeImage := bytes.NewBufferString("fake-aws-image") + uploader, err := awscloud.NewUploader("region", "bucket", "ami", nil) + assert.NoError(t, err) + var uploadLog bytes.Buffer + err = uploader.UploadAndRegister(fakeImage, &uploadLog) + // XXX: this should probably have a context + assert.EqualError(t, err, "fake-register-err") + assert.Equal(t, 1, fa.uploadFromReaderCalls) + assert.Equal(t, 1, fa.registerCalls) + assert.Equal(t, 1, fa.deleteObjectCalls) + expectedUploadLog := `Uploading ami to bucket:01010101-0101-4101-8101-010101010101-ami +File uploaded to some-location +Registering AMI ami +Deleted S3 object bucket:01010101-0101-4101-8101-010101010101-ami +` + assert.Equal(t, expectedUploadLog, uploadLog.String()) +} + +func TestUploaderUploadButRegisterErrorAndDeleteError(t *testing.T) { + fa := &fakeAWSClient{ + uploadFromReader: &s3manager.UploadOutput{ + Location: "some-location", + }, + registerErr: fmt.Errorf("fake-register-err"), + deleteObjectErr: fmt.Errorf("fake-delete-object-err"), + } + restore := awscloud.MockNewAwsClient(func(string) (awscloud.AwsClient, error) { + return fa, nil + }) + defer restore() + + fakeImage := bytes.NewBufferString("fake-aws-image") + uploader, err := awscloud.NewUploader("region", "bucket", "ami", nil) + assert.NoError(t, err) + var uploadLog bytes.Buffer + err = uploader.UploadAndRegister(fakeImage, &uploadLog) + // XXX: this should probably have a context + assert.EqualError(t, err, "fake-register-err\nfake-delete-object-err") +} diff --git a/pkg/cloud/uploader.go b/pkg/cloud/uploader.go new file mode 100644 index 0000000000..2378028b41 --- /dev/null +++ b/pkg/cloud/uploader.go @@ -0,0 +1,24 @@ +package cloud + +import ( + "io" +) + +// Uploader is an interface that is returned from the actual +// cloud implementation. The uploader will be parameterized +// by the actual cloud implemntation, e.g. +// +// awscloud.NewUploader(region, bucket, image) Uploader +// +// which is outside the scope of this interface. +type Uploader interface { + // Check can be called before the actual upload to ensure + // all permissions are correct + Check(status io.Writer) error + + // UploadAndRegister will upload the given image from + // the reader and write status message to the given + // status writer. + // To implement progress a proxy reader can be used. + UploadAndRegister(f io.Reader, status io.Writer) error +} diff --git a/pkg/imagefilter/formatter_test.go b/pkg/imagefilter/formatter_test.go index 9842e21804..f8b4ab5414 100644 --- a/pkg/imagefilter/formatter_test.go +++ b/pkg/imagefilter/formatter_test.go @@ -105,3 +105,11 @@ func TestResultsFormatter(t *testing.T) { assert.Equal(t, tc.expectsOutput, buf.String(), tc) } } + +func TestNewDifferentObj(t *testing.T) { + fmt1, err := imagefilter.NewResultsFormatter(imagefilter.OutputFormatText) + assert.NoError(t, err) + fmt2, err := imagefilter.NewResultsFormatter(imagefilter.OutputFormatText) + assert.NoError(t, err) + assert.NotEqual(t, fmt1, fmt2) +}