From 7fd150b9385e93a3d08012dcae0b20f7b1ed80f2 Mon Sep 17 00:00:00 2001
From: Sanne Raymaekers <sanne.raymaekers@gmail.com>
Date: Wed, 14 Feb 2024 12:06:26 +0100
Subject: [PATCH] cloud/awscloud: specify subnets when creating secure instance

For non-default VPCs, AWS needs the subnets it can launch the instance
in, otherwise it will try to launch the instance in the default VPC,
even if the supplied security groups are attached to a non-default VPC.

Furthermore there can only be 1 subnet specified per availability zone,
so query the subnets in the VPC of the host (as the instance needs to be
launched in the same network), and pick 1 of the VPC's subnets per AZ.
---
 internal/cloud/awscloud/secure-instance.go | 32 ++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/internal/cloud/awscloud/secure-instance.go b/internal/cloud/awscloud/secure-instance.go
index c1ca69105..36cdcf307 100644
--- a/internal/cloud/awscloud/secure-instance.go
+++ b/internal/cloud/awscloud/secure-instance.go
@@ -73,6 +73,37 @@ func (a *AWS) RunSecureInstance(iamProfile string) (*SecureInstance, error) {
 		return nil, err
 	}
 
+	descrSubnetsOutput, err := a.ec2.DescribeSubnets(&ec2.DescribeSubnetsInput{
+		Filters: []*ec2.Filter{
+			&ec2.Filter{
+				Name: aws.String("vpc-id"),
+				Values: []*string{
+					aws.String(vpcID),
+				},
+			},
+		},
+	})
+	if err != nil {
+		return nil, err
+	}
+	if len(descrSubnetsOutput.Subnets) == 0 {
+		return nil, fmt.Errorf("Expected at least 1 subnet in the VPC, got 0")
+	}
+
+	// For creating a fleet in a non-default VPC, AWS needs the subnets, and at most 1 subnet per AZ.
+	// If a VPC has multiple subnets for a single AZ, only pick the first one.
+	overrides := []*ec2.FleetLaunchTemplateOverridesRequest{}
+	availZones := map[string]struct{}{}
+	for _, subnet := range descrSubnetsOutput.Subnets {
+		az := *subnet.AvailabilityZone
+		if _, ok := availZones[az]; !ok {
+			overrides = append(overrides, &ec2.FleetLaunchTemplateOverridesRequest{
+				SubnetId: subnet.SubnetId,
+			})
+			availZones[az] = struct{}{}
+		}
+	}
+
 	createFleetOutput, err := a.ec2.CreateFleet(&ec2.CreateFleetInput{
 		LaunchTemplateConfigs: []*ec2.FleetLaunchTemplateConfigRequest{
 			&ec2.FleetLaunchTemplateConfigRequest{
@@ -80,6 +111,7 @@ func (a *AWS) RunSecureInstance(iamProfile string) (*SecureInstance, error) {
 					LaunchTemplateId: aws.String(secureInstance.LTID),
 					Version:          aws.String("1"),
 				},
+				Overrides: overrides,
 			},
 		},
 		TagSpecifications: []*ec2.TagSpecification{