Support for Security Group specification/override

pkg/cloudprovider/aws/ami.go

@@ -0,0 +1,51 @@
+package aws
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/ssm"
+	"github.com/aws/aws-sdk-go/service/ssm/ssmiface"
+	"github.com/patrickmn/go-cache"
+	"k8s.io/client-go/kubernetes"
+)
+
+type AMIProvider struct {
+	cache     *cache.Cache
+	ssm       ssmiface.SSMAPI
+	clientSet *kubernetes.Clientset
+}
+
+func NewAMIProvider(ssm ssmiface.SSMAPI, clientSet *kubernetes.Clientset) *AMIProvider {
+	return &AMIProvider{
+		ssm:       ssm,
+		clientSet: clientSet,
+		cache:     cache.New(CacheTTL, CacheCleanupInterval),
+	}
+}
+
+func (p *AMIProvider) Get(ctx context.Context, constraints *Constraints) (string, error) {
+	version, err := p.kubeServerVersion()
+	if err != nil {
+		return "", fmt.Errorf("kube server version, %w", err)
+	}
+	name := fmt.Sprintf("/aws/service/bottlerocket/aws-k8s-%s/%s/latest/image_id", version, KubeToAWSArchitectures[*constraints.Architecture])
+	if id, ok := p.cache.Get(name); ok {
+		return id.(string), nil
+	}
+	output, err := p.ssm.GetParameterWithContext(ctx, &ssm.GetParameterInput{Name: aws.String(name)})
+	if err != nil {
+		return "", fmt.Errorf("getting ssm parameter, %w", err)
+	}
+	return aws.StringValue(output.Parameter.Value), nil
+}
+
+func (p *AMIProvider) kubeServerVersion() (string, error) {
+	version, err := p.clientSet.Discovery().ServerVersion()
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprintf("%s.%s", version.Major, strings.TrimSuffix(version.Minor, "+")), nil
+}

pkg/cloudprovider/aws/cloudprovider.go

@@ -28,7 +28,6 @@ import (
 	"github.com/awslabs/karpenter/pkg/cloudprovider"
 	"github.com/awslabs/karpenter/pkg/cloudprovider/aws/utils"
 	"github.com/awslabs/karpenter/pkg/utils/project"
-	"github.com/patrickmn/go-cache"
 	"go.uber.org/zap"
 	v1 "k8s.io/api/core/v1"
 	"knative.dev/pkg/apis"
@@ -76,13 +75,11 @@ func NewCloudProvider(options cloudprovider.Options) *CloudProvider {
 	ec2api := ec2.New(sess)
 	return &CloudProvider{
 		nodeAPI: &NodeFactory{ec2api: ec2api},
-		launchTemplateProvider: &LaunchTemplateProvider{
-			ec2api:                ec2api,
-			cache:                 cache.New(CacheTTL, CacheCleanupInterval),
-			securityGroupProvider: NewSecurityGroupProvider(ec2api),
-			ssm:                   ssm.New(sess),
-			clientSet:             options.ClientSet,
-		},
+		launchTemplateProvider: NewLaunchTemplateProvider(
+			ec2api,
+			NewAMIProvider(ssm.New(sess), options.ClientSet),
+			NewSecurityGroupProvider(ec2api),
+		),
 		subnetProvider:       NewSubnetProvider(ec2api),
 		instanceTypeProvider: NewInstanceTypeProvider(ec2api),
 		instanceProvider:     &InstanceProvider{ec2api: ec2api},

pkg/cloudprovider/aws/constraints.go

@@ -38,12 +38,16 @@ var (
 	LaunchTemplateVersionLabel = AWSLabelPrefix + "launch-template-version"
 	SubnetNameLabel            = AWSLabelPrefix + "subnet-name"
 	SubnetTagKeyLabel          = AWSLabelPrefix + "subnet-tag-key"
+	SecurityGroupNameLabel     = AWSLabelPrefix + "security-group-name"
+	SecurityGroupTagKeyLabel   = AWSLabelPrefix + "security-group-tag-key"
 	AllowedLabels              = []string{
 		CapacityTypeLabel,
 		LaunchTemplateIdLabel,
 		LaunchTemplateVersionLabel,
 		SubnetNameLabel,
 		SubnetTagKeyLabel,
+		SecurityGroupNameLabel,
+		SecurityGroupTagKeyLabel,
 	}
 	AWSToKubeArchitectures = map[string]string{
 		"x86_64":                   v1alpha1.ArchitectureAmd64,
@@ -66,8 +70,8 @@ func (c *Constraints) GetCapacityType() string {
 }
 
 type LaunchTemplate struct {
-	Id      *string
-	Version *string
+	Id      string
+	Version string
 }
 
 func (c *Constraints) GetLaunchTemplate() *LaunchTemplate {
@@ -80,25 +84,41 @@ func (c *Constraints) GetLaunchTemplate() *LaunchTemplate {
 		version = DefaultLaunchTemplateVersion
 	}
 	return &LaunchTemplate{
-		Id:      &id,
-		Version: &version,
+		Id:      id,
+		Version: version,
 	}
 }
 
 func (c *Constraints) GetSubnetName() *string {
-	subnetName, ok := c.Labels[SubnetNameLabel]
+	name, ok := c.Labels[SubnetNameLabel]
 	if !ok {
 		return nil
 	}
-	return aws.String(subnetName)
+	return aws.String(name)
 }
 
 func (c *Constraints) GetSubnetTagKey() *string {
-	subnetTag, ok := c.Labels[SubnetTagKeyLabel]
+	tag, ok := c.Labels[SubnetTagKeyLabel]
 	if !ok {
 		return nil
 	}
-	return aws.String(subnetTag)
+	return aws.String(tag)
+}
+
+func (c *Constraints) GetSecurityGroupName() *string {
+	name, ok := c.Labels[SecurityGroupNameLabel]
+	if !ok {
+		return nil
+	}
+	return aws.String(name)
+}
+
+func (c *Constraints) GetSecurityGroupTagKey() *string {
+	tag, ok := c.Labels[SecurityGroupTagKeyLabel]
+	if !ok {
+		return nil
+	}
+	return aws.String(tag)
 }
 
 func (c *Constraints) Validate(ctx context.Context) (errs *apis.FieldError) {

pkg/cloudprovider/aws/fake/ec2api.go

@@ -20,6 +20,7 @@ import (
 
 	"github.com/Pallinder/go-randomdata"
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
@@ -29,17 +30,17 @@ import (
 // EC2Behavior must be reset between tests otherwise tests will
 // pollute each other.
 type EC2Behavior struct {
-	CreateFleetOutput                   *ec2.CreateFleetOutput
 	DescribeInstancesOutput             *ec2.DescribeInstancesOutput
 	DescribeLaunchTemplatesOutput       *ec2.DescribeLaunchTemplatesOutput
 	DescribeSubnetsOutput               *ec2.DescribeSubnetsOutput
 	DescribeSecurityGroupsOutput        *ec2.DescribeSecurityGroupsOutput
 	DescribeInstanceTypesOutput         *ec2.DescribeInstanceTypesOutput
 	DescribeInstanceTypeOfferingsOutput *ec2.DescribeInstanceTypeOfferingsOutput
 	DescribeAvailabilityZonesOutput     *ec2.DescribeAvailabilityZonesOutput
-	WantErr                             error
-	CalledWithCreateFleetInput          []ec2.CreateFleetInput
+	CalledWithCreateFleetInput          []*ec2.CreateFleetInput
+	CalledWithCreateLaunchTemplateInput []*ec2.CreateLaunchTemplateInput
 	Instances                           []*ec2.Instance
+	LaunchTemplates                     []*ec2.LaunchTemplate
 }
 
 type EC2API struct {
@@ -54,13 +55,7 @@ func (e *EC2API) Reset() {
 }
 
 func (e *EC2API) CreateFleetWithContext(ctx context.Context, input *ec2.CreateFleetInput, options ...request.Option) (*ec2.CreateFleetOutput, error) {
-	e.CalledWithCreateFleetInput = append(e.CalledWithCreateFleetInput, *input)
-	if e.WantErr != nil {
-		return nil, e.WantErr
-	}
-	if e.CreateFleetOutput != nil {
-		return e.CreateFleetOutput, nil
-	}
+	e.CalledWithCreateFleetInput = append(e.CalledWithCreateFleetInput, input)
 	if input.LaunchTemplateConfigs[0].LaunchTemplateSpecification.LaunchTemplateId == nil &&
 		input.LaunchTemplateConfigs[0].LaunchTemplateSpecification.LaunchTemplateName == nil {
 		return nil, fmt.Errorf("missing launch template id or name")
@@ -74,10 +69,14 @@ func (e *EC2API) CreateFleetWithContext(ctx context.Context, input *ec2.CreateFl
 	return &ec2.CreateFleetOutput{Instances: []*ec2.CreateFleetInstance{{InstanceIds: []*string{instance.InstanceId}}}}, nil
 }
 
+func (e *EC2API) CreateLaunchTemplateWithContext(ctx context.Context, input *ec2.CreateLaunchTemplateInput, options ...request.Option) (*ec2.CreateLaunchTemplateOutput, error) {
+	e.CalledWithCreateLaunchTemplateInput = append(e.CalledWithCreateLaunchTemplateInput, input)
+	launchTemplate := &ec2.LaunchTemplate{LaunchTemplateName: input.LaunchTemplateName, LaunchTemplateId: aws.String("test-launch-template-id")}
+	e.LaunchTemplates = append(e.LaunchTemplates, launchTemplate)
+	return &ec2.CreateLaunchTemplateOutput{LaunchTemplate: launchTemplate}, nil
+}
+
 func (e *EC2API) DescribeInstancesWithContext(context.Context, *ec2.DescribeInstancesInput, ...request.Option) (*ec2.DescribeInstancesOutput, error) {
-	if e.WantErr != nil {
-		return nil, e.WantErr
-	}
 	if e.DescribeInstancesOutput != nil {
 		return e.DescribeInstancesOutput, nil
 	}
@@ -86,23 +85,25 @@ func (e *EC2API) DescribeInstancesWithContext(context.Context, *ec2.DescribeInst
 	}, nil
 }
 
-func (e *EC2API) DescribeLaunchTemplatesWithContext(context.Context, *ec2.DescribeLaunchTemplatesInput, ...request.Option) (*ec2.DescribeLaunchTemplatesOutput, error) {
-	if e.WantErr != nil {
-		return nil, e.WantErr
-	}
+func (e *EC2API) DescribeLaunchTemplatesWithContext(ctx context.Context, input *ec2.DescribeLaunchTemplatesInput, options ...request.Option) (*ec2.DescribeLaunchTemplatesOutput, error) {
 	if e.DescribeLaunchTemplatesOutput != nil {
 		return e.DescribeLaunchTemplatesOutput, nil
 	}
-	return &ec2.DescribeLaunchTemplatesOutput{LaunchTemplates: []*ec2.LaunchTemplate{{
-		LaunchTemplateName: aws.String("test-launch-template-name"),
-		LaunchTemplateId:   aws.String("test-launch-template-id"),
-	}}}, nil
+	output := &ec2.DescribeLaunchTemplatesOutput{}
+	for _, wanted := range input.LaunchTemplateNames {
+		for _, launchTemplate := range e.LaunchTemplates {
+			if launchTemplate.LaunchTemplateName == wanted {
+				output.LaunchTemplates = append(output.LaunchTemplates, launchTemplate)
+			}
+		}
+	}
+	if len(output.LaunchTemplates) == 0 {
+		return nil, awserr.New("InvalidLaunchTemplateName.NotFoundException", "not found", nil)
+	}
+	return output, nil
 }
 
 func (e *EC2API) DescribeSubnetsWithContext(context.Context, *ec2.DescribeSubnetsInput, ...request.Option) (*ec2.DescribeSubnetsOutput, error) {
-	if e.WantErr != nil {
-		return nil, e.WantErr
-	}
 	if e.DescribeSubnetsOutput != nil {
 		return e.DescribeSubnetsOutput, nil
 	}
@@ -117,19 +118,17 @@ func (e *EC2API) DescribeSubnetsWithContext(context.Context, *ec2.DescribeSubnet
 }
 
 func (e *EC2API) DescribeSecurityGroupsWithContext(context.Context, *ec2.DescribeSecurityGroupsInput, ...request.Option) (*ec2.DescribeSecurityGroupsOutput, error) {
-	if e.WantErr != nil {
-		return nil, e.WantErr
-	}
 	if e.DescribeSecurityGroupsOutput != nil {
 		return e.DescribeSecurityGroupsOutput, nil
 	}
-	return &ec2.DescribeSecurityGroupsOutput{SecurityGroups: []*ec2.SecurityGroup{{GroupId: aws.String("test-group")}}}, nil
+	return &ec2.DescribeSecurityGroupsOutput{SecurityGroups: []*ec2.SecurityGroup{
+		{GroupId: aws.String("test-security-group-1"), Tags: []*ec2.Tag{{Key: aws.String("Name"), Value: aws.String("test-security-group-1")}}},
+		{GroupId: aws.String("test-security-group-2"), Tags: []*ec2.Tag{{Key: aws.String("Name"), Value: aws.String("test-security-group-2")}}},
+		{GroupId: aws.String("test-security-group-3"), Tags: []*ec2.Tag{{Key: aws.String("Name"), Value: aws.String("test-security-group-3")}, {Key: aws.String("TestTag")}}},
+	}}, nil
 }
 
 func (e *EC2API) DescribeAvailabilityZonesWithContext(context.Context, *ec2.DescribeAvailabilityZonesInput, ...request.Option) (*ec2.DescribeAvailabilityZonesOutput, error) {
-	if e.WantErr != nil {
-		return nil, e.WantErr
-	}
 	if e.DescribeAvailabilityZonesOutput != nil {
 		return e.DescribeAvailabilityZonesOutput, nil
 	}
@@ -141,9 +140,6 @@ func (e *EC2API) DescribeAvailabilityZonesWithContext(context.Context, *ec2.Desc
 }
 
 func (e *EC2API) DescribeInstanceTypesPagesWithContext(ctx context.Context, input *ec2.DescribeInstanceTypesInput, fn func(*ec2.DescribeInstanceTypesOutput, bool) bool, opts ...request.Option) error {
-	if e.WantErr != nil {
-		return e.WantErr
-	}
 	if e.DescribeInstanceTypesOutput != nil {
 		fn(e.DescribeInstanceTypesOutput, false)
 		return nil
@@ -267,9 +263,6 @@ func (e *EC2API) DescribeInstanceTypesPagesWithContext(ctx context.Context, inpu
 }
 
 func (e *EC2API) DescribeInstanceTypeOfferingsPagesWithContext(ctx context.Context, input *ec2.DescribeInstanceTypeOfferingsInput, fn func(*ec2.DescribeInstanceTypeOfferingsOutput, bool) bool, opts ...request.Option) error {
-	if e.WantErr != nil {
-		return e.WantErr
-	}
 	if e.DescribeInstanceTypeOfferingsOutput != nil {
 		fn(e.DescribeInstanceTypeOfferingsOutput, false)
 		return nil

pkg/cloudprovider/aws/instance.go

@@ -87,8 +87,8 @@ func (p *InstanceProvider) Create(ctx context.Context,
 		},
 		LaunchTemplateConfigs: []*ec2.FleetLaunchTemplateConfigRequest{{
 			LaunchTemplateSpecification: &ec2.FleetLaunchTemplateSpecificationRequest{
-				LaunchTemplateId: launchTemplate.Id,
-				Version:          launchTemplate.Version,
+				LaunchTemplateId: aws.String(launchTemplate.Id),
+				Version:          aws.String(launchTemplate.Version),
 			},
 			Overrides: overrides,
 		}},