[CNVM] Sort EC2 instances by volume size (#1133)

deploy/cloudformation/elastic-agent-ec2-cnvm.yml

@@ -78,6 +78,7 @@ Resources:
                   - ec2:CreateSnapshots
                   - ec2:CreateTags
                   - ec2:DeleteSnapshot
+                  - ec2:DescribeVolumes
                 Resource: '*'
               - Effect: Allow
                 Action:

resources/providers/awslib/ec2/ec2_instance.go

@@ -30,6 +30,7 @@ type Ec2Instance struct {
 	types.Instance
 	Region     string
 	awsAccount string
+	RootVolume *Volume
 }
 
 type SecurityGroupInfo struct {

resources/providers/awslib/ec2/provider.go

@@ -52,6 +52,7 @@ type Client interface {
 	DescribeSnapshots(ctx context.Context, params *ec2.DescribeSnapshotsInput, optFns ...func(*ec2.Options)) (*ec2.DescribeSnapshotsOutput, error)
 	DeleteSnapshot(ctx context.Context, params *ec2.DeleteSnapshotInput, optFns ...func(*ec2.Options)) (*ec2.DeleteSnapshotOutput, error)
 	DescribeRouteTables(ctx context.Context, params *ec2.DescribeRouteTablesInput, optFns ...func(*ec2.Options)) (*ec2.DescribeRouteTablesOutput, error)
+	DescribeVolumes(ctx context.Context, params *ec2.DescribeVolumesInput, optFns ...func(*ec2.Options)) (*ec2.DescribeVolumesOutput, error)
 }
 
 func (p *Provider) DescribeNetworkAcl(ctx context.Context) ([]awslib.AwsResource, error) {
@@ -163,8 +164,8 @@ func (p *Provider) GetEbsEncryptionByDefault(ctx context.Context) ([]awslib.AwsR
 	})
 }
 
-func (p *Provider) DescribeInstances(ctx context.Context) ([]Ec2Instance, error) {
-	insances, err := awslib.MultiRegionFetch(ctx, p.clients, func(ctx context.Context, region string, c Client) ([]Ec2Instance, error) {
+func (p *Provider) DescribeInstances(ctx context.Context) ([]*Ec2Instance, error) {
+	insances, err := awslib.MultiRegionFetch(ctx, p.clients, func(ctx context.Context, region string, c Client) ([]*Ec2Instance, error) {
 		input := &ec2.DescribeInstancesInput{}
 		allInstances := []types.Instance{}
 		for {
@@ -181,9 +182,9 @@ func (p *Provider) DescribeInstances(ctx context.Context) ([]Ec2Instance, error)
 			input.NextToken = output.NextToken
 		}
 
-		var result []Ec2Instance
+		var result []*Ec2Instance
 		for _, instance := range allInstances {
-			result = append(result, Ec2Instance{
+			result = append(result, &Ec2Instance{
 				Instance:   instance,
 				awsAccount: p.awsAccountID,
 				Region:     region,
@@ -194,7 +195,7 @@ func (p *Provider) DescribeInstances(ctx context.Context) ([]Ec2Instance, error)
 	return lo.Flatten(insances), err
 }
 
-func (p *Provider) CreateSnapshots(ctx context.Context, ins Ec2Instance) ([]EBSSnapshot, error) {
+func (p *Provider) CreateSnapshots(ctx context.Context, ins *Ec2Instance) ([]EBSSnapshot, error) {
 	client := p.clients[ins.Region]
 	if client == nil {
 		return nil, fmt.Errorf("error in CreateSnapshots no client for region %s", ins.Region)
@@ -220,7 +221,7 @@ func (p *Provider) CreateSnapshots(ctx context.Context, ins Ec2Instance) ([]EBSS
 
 	var result []EBSSnapshot
 	for _, snap := range res.Snapshots {
-		result = append(result, FromSnapshotInfo(snap, ins.Region, p.awsAccountID, ins))
+		result = append(result, FromSnapshotInfo(snap, ins.Region, p.awsAccountID, *ins))
 	}
 	return result, nil
 }
@@ -294,3 +295,49 @@ func (p *Provider) GetRouteTableForSubnet(ctx context.Context, region string, su
 
 	return routeTables.RouteTables[0], nil
 }
+
+func (p *Provider) DescribeVolumes(ctx context.Context, instances []*Ec2Instance) ([]*Volume, error) {
+	instanceFilter := lo.Map(instances, func(ins *Ec2Instance, _ int) string { return *ins.InstanceId })
+	volumes, err := awslib.MultiRegionFetch(ctx, p.clients, func(ctx context.Context, region string, c Client) ([]*Volume, error) {
+		input := &ec2.DescribeVolumesInput{
+			Filters: []types.Filter{
+				{
+					Name:   aws.String("attachment.instance-id"),
+					Values: instanceFilter,
+				},
+			},
+		}
+		allVolumes := []types.Volume{}
+		for {
+			output, err := c.DescribeVolumes(ctx, input)
+			if err != nil {
+				return nil, err
+			}
+			allVolumes = append(allVolumes, output.Volumes...)
+			if output.NextToken == nil {
+				break
+			}
+			input.NextToken = output.NextToken
+		}
+
+		var result []*Volume
+		for _, vol := range allVolumes {
+			if len(vol.Attachments) != 1 {
+				p.log.Errorf("Volume %s has %d attachments", *vol.VolumeId, len(vol.Attachments))
+				continue
+			}
+
+			result = append(result, &Volume{
+				VolumeId:   *vol.VolumeId,
+				Size:       int(*vol.Size),
+				awsAccount: p.awsAccountID,
+				Region:     region,
+				Encrypted:  *vol.Encrypted,
+				InstanceId: *vol.Attachments[0].InstanceId,
+				Device:     *vol.Attachments[0].Device,
+			})
+		}
+		return result, nil
+	})
+	return lo.Flatten(volumes), err
+}

resources/providers/awslib/ec2/provider_test.go

@@ -32,6 +32,7 @@ import (
 	"github.com/stretchr/testify/mock"
 
 	"github.com/elastic/cloudbeat/resources/providers/awslib"
+	"github.com/elastic/cloudbeat/resources/utils/testhelper"
 )
 
 var onlyDefaultRegion = []string{awslib.DefaultRegion}
@@ -314,7 +315,7 @@ func TestProvider_CreateSnapshots(t *testing.T) {
 	}
 	type args struct {
 		ctx context.Context
-		ins Ec2Instance
+		ins *Ec2Instance
 	}
 	tests := []struct {
 		name    string
@@ -472,3 +473,138 @@ func TestProvider_GetRouteTableForSubnet(t *testing.T) {
 		})
 	}
 }
+
+func TestProvider_DescribeVolumes(t *testing.T) {
+	expectToken := func(token string) func(input *ec2.DescribeVolumesInput) bool {
+		return func(input *ec2.DescribeVolumesInput) bool {
+			return *input.NextToken == token
+		}
+	}
+
+	expectInstances := func(ids ...string) func(input *ec2.DescribeVolumesInput) bool {
+		return func(input *ec2.DescribeVolumesInput) bool {
+			if len(input.Filters) != 1 {
+				return false
+			}
+			if *input.Filters[0].Name != "attachment.instance-id" {
+				return false
+			}
+			if len(input.Filters[0].Values) != len(ids) {
+				return false
+			}
+			for i, id := range ids {
+				if input.Filters[0].Values[i] != id {
+					return false
+				}
+			}
+			return true
+		}
+	}
+	mockResult := types.Volume{
+		VolumeId:  aws.String("vol-123456789"),
+		Encrypted: aws.Bool(true),
+		Size:      aws.Int32(8),
+		Attachments: []types.VolumeAttachment{
+			{
+				InstanceId: aws.String("i-123456789"),
+				Device:     aws.String("/dev/sda1"),
+			},
+		},
+	}
+	expectedVolume := &Volume{
+		awsAccount: "aws-account",
+		VolumeId:   "vol-123456789",
+		InstanceId: "i-123456789",
+		Device:     "/dev/sda1",
+		Encrypted:  true,
+		Size:       8,
+		Region:     awslib.DefaultRegion,
+	}
+
+	tests := []struct {
+		name      string
+		client    func() Client
+		instances []*Ec2Instance
+		want      []*Volume
+		wantErr   bool
+		regions   []string
+	}{
+		{
+			name:      "Get 3 volumes from 3 pages",
+			instances: []*Ec2Instance{},
+			client: func() Client {
+				m := &MockClient{}
+				m.EXPECT().DescribeVolumes(mock.Anything, mock.MatchedBy(expectInstances())).Return(&ec2.DescribeVolumesOutput{Volumes: []types.Volume{mockResult}, NextToken: aws.String("1")}, nil).Once()
+				m.EXPECT().DescribeVolumes(mock.Anything, mock.MatchedBy(expectToken("1"))).Return(&ec2.DescribeVolumesOutput{Volumes: []types.Volume{mockResult}, NextToken: aws.String("2")}, nil).Once()
+				m.EXPECT().DescribeVolumes(mock.Anything, mock.MatchedBy(expectToken("2"))).Return(&ec2.DescribeVolumesOutput{Volumes: []types.Volume{mockResult}}, nil).Once()
+				return m
+			},
+			want:    []*Volume{expectedVolume, expectedVolume, expectedVolume},
+			wantErr: false,
+			regions: onlyDefaultRegion,
+		},
+		{
+			name:      "Get 3 volumes from 1 page",
+			instances: []*Ec2Instance{},
+			client: func() Client {
+				m := &MockClient{}
+				m.EXPECT().DescribeVolumes(mock.Anything, mock.Anything).Return(&ec2.DescribeVolumesOutput{Volumes: []types.Volume{mockResult, mockResult, mockResult}}, nil).Once()
+				return m
+			},
+			want:    []*Volume{expectedVolume, expectedVolume, expectedVolume},
+			wantErr: false,
+			regions: onlyDefaultRegion,
+		},
+		{
+			name: "Get volumes filtered by instance id from 2 pages",
+			instances: []*Ec2Instance{
+				{Instance: types.Instance{InstanceId: aws.String("123")}},
+				{Instance: types.Instance{InstanceId: aws.String("456")}},
+			},
+			client: func() Client {
+				m := &MockClient{}
+				m.EXPECT().DescribeVolumes(mock.Anything, mock.MatchedBy(expectInstances("123", "456"))).Return(&ec2.DescribeVolumesOutput{Volumes: []types.Volume{mockResult}, NextToken: aws.String("1")}, nil).Once()
+				m.EXPECT().DescribeVolumes(mock.Anything, mock.MatchedBy(expectInstances("123", "456"))).Return(&ec2.DescribeVolumesOutput{Volumes: []types.Volume{mockResult}}, nil).Once()
+				return m
+			},
+			want:    []*Volume{expectedVolume, expectedVolume},
+			wantErr: false,
+			regions: onlyDefaultRegion,
+		},
+		{
+			name:      "Get error at 3rd page",
+			instances: []*Ec2Instance{},
+			client: func() Client {
+				m := &MockClient{}
+				m.EXPECT().DescribeVolumes(mock.Anything, mock.Anything).Return(&ec2.DescribeVolumesOutput{Volumes: []types.Volume{mockResult}, NextToken: aws.String("1")}, nil).Once()
+				m.EXPECT().DescribeVolumes(mock.Anything, mock.Anything).Return(&ec2.DescribeVolumesOutput{Volumes: []types.Volume{mockResult}, NextToken: aws.String("2")}, nil).Once()
+				m.EXPECT().DescribeVolumes(mock.Anything, mock.Anything).Return(nil, errors.New("bla")).Once()
+				return m
+			},
+			want:    []*Volume{},
+			wantErr: true,
+			regions: onlyDefaultRegion,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			clients := map[string]Client{}
+			for _, r := range tt.regions {
+				clients[r] = tt.client()
+			}
+			p := &Provider{
+				log:          testhelper.NewLogger(t),
+				clients:      clients,
+				awsAccountID: "aws-account",
+			}
+			got, err := p.DescribeVolumes(context.Background(), tt.instances)
+			if tt.wantErr {
+				assert.Error(t, err)
+				return
+			}
+
+			assert.NoError(t, err)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}

resources/providers/awslib/ec2/volume.go

@@ -0,0 +1,28 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package ec2
+
+type Volume struct {
+	VolumeId   string
+	InstanceId string
+	Region     string
+	awsAccount string
+	Size       int
+	Encrypted  bool
+	Device     string
+}