support custom AMI (#10)

* support custom AMI * Apply suggestions from code review Co-authored-by: Kyle Brennan <[email protected]> --------- Co-authored-by: Kyle Brennan <[email protected]>
gitpod-io · Jan 17, 2025 · 7bb8cda · 7bb8cda
1 parent bb54dc0
commit 7bb8cda
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 7 deletions.
diff --git a/gitpod-network-check/README.md b/gitpod-network-check/README.md
@@ -43,6 +43,7 @@ A CLI to check if your network setup is suitable for the installation of Gitpod.
    main-subnets: subnet-0554e84f033a64c56, subnet-08584621e7754e505, subnet-094c6fd68aea493b7
    pod-subnets: subnet-028d11dce93b8eefc, subnet-04ec8257d95c434b7,subnet-00a83550ce709f39c
    https-hosts: accounts.google.com, github.com
+   instance-ami: # put your custom ami id here if you want to use it, otherwise it will using latest ubuntu AMI from aws
    ```
 
 2. Run the network diagnosis

diff --git a/gitpod-network-check/cmd/checks.go b/gitpod-network-check/cmd/checks.go
@@ -320,9 +320,19 @@ func launchInstances(ctx context.Context, ec2Client *ec2.Client, subnets []strin
 }
 
 func launchInstanceInSubnet(ctx context.Context, ec2Client *ec2.Client, subnetID, secGroupId string, instanceProfileName *string, instanceType types.InstanceType) (string, error) {
-	regionalAMI, err := findUbuntuAMI(ctx, ec2Client)
-	if err != nil {
-		return "", err
+	amiId := ""
+	if networkConfig.InstanceAMI != "" {
+		customAMIId, err := findCustomAMI(ctx, ec2Client, networkConfig.InstanceAMI)
+		if err != nil {
+			return "", err
+		}
+		amiId = customAMIId
+	} else {
+		regionalAMI, err := findUbuntuAMI(ctx, ec2Client)
+		if err != nil {
+			return "", err
+		}
+		amiId = regionalAMI
 	}
 
 	// Specify the user data script to install the SSM Agent
@@ -335,7 +345,7 @@ func launchInstanceInSubnet(ctx context.Context, ec2Client *ec2.Client, subnetID
 	userDataEncoded := base64.StdEncoding.EncodeToString([]byte(userData))
 
 	input := &ec2.RunInstancesInput{
-		ImageId:          aws.String(regionalAMI), // Example AMI ID, replace with an actual one
+		ImageId:          aws.String(amiId), // Example AMI ID, replace with an actual one
 		InstanceType:     instanceType,
 		MaxCount:         aws.Int32(1),
 		MinCount:         aws.Int32(1),
@@ -359,7 +369,7 @@ func launchInstanceInSubnet(ctx context.Context, ec2Client *ec2.Client, subnetID
 	}
 
 	var result *ec2.RunInstancesOutput
-	err = wait.PollUntilContextTimeout(ctx, 500*time.Millisecond, 10*time.Second, false, func(ctx context.Context) (done bool, err error) {
+	err := wait.PollUntilContextTimeout(ctx, 500*time.Millisecond, 10*time.Second, false, func(ctx context.Context) (done bool, err error) {
 		result, err = ec2Client.RunInstances(ctx, input)
 
 		if err != nil {
@@ -384,6 +394,22 @@ func launchInstanceInSubnet(ctx context.Context, ec2Client *ec2.Client, subnetID
 	return aws.ToString(result.Instances[0].InstanceId), nil
 }
 
+func findCustomAMI(ctx context.Context, client *ec2.Client, amiId string) (string, error) {
+	input := &ec2.DescribeImagesInput{
+		ImageIds: []string{amiId},
+	}
+
+	result, err := client.DescribeImages(ctx, input)
+	if err != nil {
+		return "", err
+	}
+	if len(result.Images) > 0 {
+		return *result.Images[0].ImageId, nil
+	}
+
+	return "", fmt.Errorf("no custom AMI found")
+}
+
 // findUbuntuAMI searches for the latest Ubuntu AMI in the region of the EC2 client.
 func findUbuntuAMI(ctx context.Context, client *ec2.Client) (string, error) {
 	// You may want to update these filters based on your specific requirements
@@ -618,7 +644,7 @@ func instanceTypeExists(ctx context.Context, svc *ec2.Client, instanceType types
 	input := &ec2.DescribeInstanceTypeOfferingsInput{
 		Filters: []types.Filter{
 			{
-				Name:  aws.String("instance-type"),
+				Name:   aws.String("instance-type"),
 				Values: []string{string(instanceType)},
 			},
 		},

diff --git a/gitpod-network-check/cmd/root.go b/gitpod-network-check/cmd/root.go
@@ -23,6 +23,7 @@ type NetworkConfig struct {
 	MainSubnets []string
 	PodSubnets  []string
 	HttpsHosts  []string
+	InstanceAMI string
 }
 
 var networkConfig = NetworkConfig{LogLevel: "INFO"}
@@ -89,6 +90,7 @@ func init() {
 	networkCheckCmd.PersistentFlags().StringSliceVar(&networkConfig.PodSubnets, "pod-subnets", []string{}, "List of pod subnets")
 	networkCheckCmd.PersistentFlags().StringSliceVar(&networkConfig.HttpsHosts, "https-hosts", []string{}, "Hosts to test for outbound HTTPS connectivity")
 	bindFlags(networkCheckCmd, v)
+	networkCheckCmd.PersistentFlags().StringVar(&networkConfig.InstanceAMI, "instance-ami", "", "Custom ec2 instance AMI id, if not set will use latest ubuntu")
 	log.Infof("ℹ️  Running with region `%s`, main subnet `%v`, pod subnet `%v`, and hosts `%v`", networkConfig.AwsRegion, networkConfig.MainSubnets, networkConfig.PodSubnets, networkConfig.HttpsHosts)
 }
 

diff --git a/gitpod-network-check/gitpod-network-check.yaml b/gitpod-network-check/gitpod-network-check.yaml
@@ -2,4 +2,6 @@ log-level: debug # Options: debug, info, warning, error
 region: eu-central-1
 main-subnets: subnet-017c6a80f4879d851, subnet-0215744d52cd1c01f
 pod-subnets: subnet-00a118009d1d572a5, subnet-062288af00ba50d86
-https-hosts: accounts.google.com, https://github.com
+https-hosts: accounts.google.com, https://github.com
+# put your custom ami id here if you want to use it, otherwise it will using latest ubuntu AMI from aws
+instance-ami: