From 92be20f2d54d95ab37c81c05106e0a994782cc8e Mon Sep 17 00:00:00 2001
From: Bharath Vedartham <vedabharath12345@gmail.com>
Date: Fri, 19 Feb 2021 01:20:06 +0530
Subject: [PATCH] Add validation for ami arch to instance type arch

---
 pkg/apis/kops/validation/BUILD.bazel |  3 ++
 pkg/apis/kops/validation/aws.go      | 45 +++++++++++++++++++++++-----
 2 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/pkg/apis/kops/validation/BUILD.bazel b/pkg/apis/kops/validation/BUILD.bazel
index 1b8ce058ba..7738027ad9 100644
--- a/pkg/apis/kops/validation/BUILD.bazel
+++ b/pkg/apis/kops/validation/BUILD.bazel
@@ -49,10 +49,13 @@ go_test(
     ],
     embed = [":go_default_library"],
     deps = [
+        "//cloudmock/aws/mockec2:go_default_library",
         "//pkg/apis/kops:go_default_library",
         "//pkg/nodeidentity/aws:go_default_library",
         "//upup/pkg/fi:go_default_library",
         "//upup/pkg/fi/cloudup/awsup:go_default_library",
+        "//vendor/github.com/aws/aws-sdk-go/aws:go_default_library",
+        "//vendor/github.com/aws/aws-sdk-go/service/ec2:go_default_library",
         "//vendor/k8s.io/apimachinery/pkg/apis/meta/v1:go_default_library",
         "//vendor/k8s.io/apimachinery/pkg/util/intstr:go_default_library",
         "//vendor/k8s.io/apimachinery/pkg/util/sets:go_default_library",
diff --git a/pkg/apis/kops/validation/aws.go b/pkg/apis/kops/validation/aws.go
index 14472c7da8..d458239e2a 100644
--- a/pkg/apis/kops/validation/aws.go
+++ b/pkg/apis/kops/validation/aws.go
@@ -65,7 +65,7 @@ func awsValidateInstanceGroup(ig *kops.InstanceGroup, cloud awsup.AWSCloud) fiel
 
 	allErrs = append(allErrs, awsValidateAdditionalSecurityGroups(field.NewPath("spec", "additionalSecurityGroups"), ig.Spec.AdditionalSecurityGroups)...)
 
-	allErrs = append(allErrs, awsValidateInstanceType(field.NewPath(ig.GetName(), "spec", "machineType"), ig.Spec.MachineType, cloud)...)
+	allErrs = append(allErrs, awsValidateInstanceTypeWithImageArch(field.NewPath(ig.GetName(), "spec", "machineType"), ig.Spec.MachineType, ig.Spec.Image, cloud)...)
 
 	allErrs = append(allErrs, awsValidateSpotDurationInMinute(field.NewPath(ig.GetName(), "spec", "spotDurationInMinutes"), ig)...)
 
@@ -121,12 +121,27 @@ func awsValidateAdditionalSecurityGroups(fieldPath *field.Path, groups []string)
 	return allErrs
 }
 
-func awsValidateInstanceType(fieldPath *field.Path, instanceType string, cloud awsup.AWSCloud) field.ErrorList {
+func awsValidateInstanceTypeWithImageArch(instanceFieldPath *field.Path, instanceType string, image string, cloud awsup.AWSCloud) field.ErrorList {
 	allErrs := field.ErrorList{}
-	if instanceType != "" && cloud != nil {
-		for _, typ := range strings.Split(instanceType, ",") {
-			if _, err := cloud.DescribeInstanceType(typ); err != nil {
-				allErrs = append(allErrs, field.Invalid(fieldPath, typ, "machine type specified is invalid"))
+
+	if cloud != nil && instanceType != "" {
+		imageInfo, err := cloud.ResolveImage(image)
+		if err != nil {
+			allErrs = append(allErrs, field.Invalid(field.NewPath("spec", "image"), image, "Image specified is invalid"))
+		}
+
+		if imageInfo != nil {
+			for _, typ := range strings.Split(instanceType, ",") {
+				machineInfo, err := cloud.DescribeInstanceType(typ)
+				if err != nil {
+					allErrs = append(allErrs, field.Invalid(instanceFieldPath, typ, "machine type specified is invalid"))
+				}
+
+				if machineInfo != nil {
+					if invalidMachineArchitecture(imageInfo, machineInfo) {
+						allErrs = append(allErrs, field.Invalid(instanceFieldPath, typ, "machine type architecture does not match image architecture"))
+					}
+				}
 			}
 		}
 	}
@@ -159,7 +174,7 @@ func awsValidateMixedInstancesPolicy(path *field.Path, spec *kops.MixedInstances
 
 	// @step: check the instance types are valid
 	for i, x := range spec.Instances {
-		errs = append(errs, awsValidateInstanceType(path.Child("instances").Index(i), x, cloud)...)
+		errs = append(errs, awsValidateInstanceTypeWithImageArch(path.Child("instances").Index(i), x, ig.Spec.Image, cloud)...)
 	}
 
 	if spec.OnDemandBase != nil {
@@ -245,3 +260,19 @@ func awsValidateLoadBalancerSubnets(fieldPath *field.Path, spec kops.ClusterSpec
 
 	return allErrs
 }
+
+func invalidMachineArchitecture(imageInfo *ec2.Image, machineInfo *ec2.InstanceTypeInfo) bool {
+	imageArch := fi.StringValue(imageInfo.Architecture)
+
+	if machineInfo.ProcessorInfo == nil {
+		return false
+	}
+
+	for _, arch := range machineInfo.ProcessorInfo.SupportedArchitectures {
+		if imageArch == fi.StringValue(arch) {
+			return false
+		}
+	}
+
+	return true
+}