Merge pull request #19 from fabi200123/add-sec-group-option

[WIP] Add extra-specs option to add SecurityGroups to the runner
cloudbase · Nov 18, 2024 · fe1b339 · fe1b339
2 parents e56d555 + bbf03cc
commit fe1b339
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 37 deletions.
diff --git a/internal/client/aws.go b/internal/client/aws.go
@@ -244,13 +244,14 @@ func (a *AwsCli) CreateRunningInstance(ctx context.Context, spec *spec.RunnerSpe
 	}
 
 	resp, err := a.client.RunInstances(ctx, &ec2.RunInstancesInput{
-		ImageId:      aws.String(spec.BootstrapParams.Image),
-		InstanceType: types.InstanceType(spec.BootstrapParams.Flavor),
-		MaxCount:     aws.Int32(1),
-		MinCount:     aws.Int32(1),
-		SubnetId:     aws.String(spec.SubnetID),
-		UserData:     aws.String(udata),
-		KeyName:      spec.SSHKeyName,
+		ImageId:          aws.String(spec.BootstrapParams.Image),
+		InstanceType:     types.InstanceType(spec.BootstrapParams.Flavor),
+		MaxCount:         aws.Int32(1),
+		MinCount:         aws.Int32(1),
+		SubnetId:         aws.String(spec.SubnetID),
+		SecurityGroupIds: spec.SecurityGroupIds,
+		UserData:         aws.String(udata),
+		KeyName:          spec.SSHKeyName,
 		TagSpecifications: []types.TagSpecification{
 			{
 				ResourceType: types.ResourceTypeInstance,

diff --git a/internal/spec/spec.go b/internal/spec/spec.go
@@ -73,11 +73,12 @@ func newExtraSpecsFromBootstrapData(data params.BootstrapInstance) (*extraSpecs,
 }
 
 type extraSpecs struct {
-	SubnetID        *string  `json:"subnet_id,omitempty" jsonschema:"pattern=^subnet-[0-9a-fA-F]{17}$"`
-	SSHKeyName      *string  `json:"ssh_key_name,omitempty" jsonschema:"description=The name of the Key Pair to use for the instance."`
-	DisableUpdates  *bool    `json:"disable_updates,omitempty" jsonschema:"description=Disable automatic updates on the VM."`
-	EnableBootDebug *bool    `json:"enable_boot_debug,omitempty" jsonschema:"description=Enable boot debug on the VM"`
-	ExtraPackages   []string `json:"extra_packages,omitempty" jsonschema:"description=Extra packages to install on the VM"`
+	SubnetID         *string  `json:"subnet_id,omitempty" jsonschema:"pattern=^subnet-[0-9a-fA-F]{17}$"`
+	SSHKeyName       *string  `json:"ssh_key_name,omitempty" jsonschema:"description=The name of the Key Pair to use for the instance."`
+	SecurityGroupIds []string `json:"security_group_ids,omitempty" jsonschema:"description=The security groups IDs to associate with the instance. Default: Amazon EC2 uses the default security group."`
+	DisableUpdates   *bool    `json:"disable_updates,omitempty" jsonschema:"description=Disable automatic updates on the VM."`
+	EnableBootDebug  *bool    `json:"enable_boot_debug,omitempty" jsonschema:"description=Enable boot debug on the VM"`
+	ExtraPackages    []string `json:"extra_packages,omitempty" jsonschema:"description=Extra packages to install on the VM"`
 	// The Cloudconfig struct from common package
 	cloudconfig.CloudConfigSpec
 }
@@ -112,15 +113,16 @@ func GetRunnerSpecFromBootstrapParams(cfg *config.Config, data params.BootstrapI
 }
 
 type RunnerSpec struct {
-	Region          string
-	DisableUpdates  bool
-	ExtraPackages   []string
-	EnableBootDebug bool
-	Tools           params.RunnerApplicationDownload
-	BootstrapParams params.BootstrapInstance
-	SubnetID        string
-	SSHKeyName      *string
-	ControllerID    string
+	Region           string
+	DisableUpdates   bool
+	ExtraPackages    []string
+	EnableBootDebug  bool
+	Tools            params.RunnerApplicationDownload
+	BootstrapParams  params.BootstrapInstance
+	SecurityGroupIds []string
+	SubnetID         string
+	SSHKeyName       *string
+	ControllerID     string
 }
 
 func (r *RunnerSpec) Validate() error {
@@ -142,6 +144,10 @@ func (r *RunnerSpec) MergeExtraSpecs(extraSpecs *extraSpecs) {
 		r.SSHKeyName = extraSpecs.SSHKeyName
 	}
 
+	if len(extraSpecs.SecurityGroupIds) > 0 {
+		r.SecurityGroupIds = extraSpecs.SecurityGroupIds
+	}
+
 	if extraSpecs.DisableUpdates != nil {
 		r.DisableUpdates = *extraSpecs.DisableUpdates
 	}

diff --git a/internal/spec/spec_test.go b/internal/spec/spec_test.go
@@ -36,14 +36,15 @@ func TestExtraSpecsFromBootstrapData(t *testing.T) {
 		{
 			name: "valid bootstrap data",
 			input: params.BootstrapInstance{
-				ExtraSpecs: json.RawMessage(`{"subnet_id": "subnet-0a0a0a0a0a0a0a0a0", "ssh_key_name": "ssh_key_name", "disable_updates": true, "enable_boot_debug": true, "extra_packages": ["package1", "package2"], "runner_install_template": "IyEvYmluL2Jhc2gKZWNobyBJbnN0YWxsaW5nIHJ1bm5lci4uLg==", "pre_install_scripts": {"setup.sh": "IyEvYmluL2Jhc2gKZWNobyBTZXR1cCBzY3JpcHQuLi4="}, "extra_context": {"key": "value"}}`),
+				ExtraSpecs: json.RawMessage(`{"subnet_id": "subnet-0a0a0a0a0a0a0a0a0", "ssh_key_name": "ssh_key_name", "security_group_ids": ["sg-018c35963edfb1cce", "sg-018c35963edfb1cee"], "disable_updates": true, "enable_boot_debug": true, "extra_packages": ["package1", "package2"], "runner_install_template": "IyEvYmluL2Jhc2gKZWNobyBJbnN0YWxsaW5nIHJ1bm5lci4uLg==", "pre_install_scripts": {"setup.sh": "IyEvYmluL2Jhc2gKZWNobyBTZXR1cCBzY3JpcHQuLi4="}, "extra_context": {"key": "value"}}`),
 			},
 			expectedOutput: &extraSpecs{
-				SubnetID:        aws.String("subnet-0a0a0a0a0a0a0a0a0"),
-				SSHKeyName:      aws.String("ssh_key_name"),
-				DisableUpdates:  aws.Bool(true),
-				EnableBootDebug: aws.Bool(true),
-				ExtraPackages:   []string{"package1", "package2"},
+				SubnetID:         aws.String("subnet-0a0a0a0a0a0a0a0a0"),
+				SSHKeyName:       aws.String("ssh_key_name"),
+				SecurityGroupIds: []string{"sg-018c35963edfb1cce", "sg-018c35963edfb1cee"},
+				DisableUpdates:   aws.Bool(true),
+				EnableBootDebug:  aws.Bool(true),
+				ExtraPackages:    []string{"package1", "package2"},
 				CloudConfigSpec: cloudconfig.CloudConfigSpec{
 					RunnerInstallTemplate: []byte("#!/bin/bash\necho Installing runner..."),
 					PreInstallScripts: map[string][]byte{
@@ -74,6 +75,15 @@ func TestExtraSpecsFromBootstrapData(t *testing.T) {
 			},
 			errString: "",
 		},
+		{
+			name: "specs just with security_group_ids",
+			input: params.BootstrapInstance{
+				ExtraSpecs: json.RawMessage(`{"security_group_ids": ["sg-018c35963edfb1cce", "sg-018c35963edfb1cee"]}`),
+			},
+			expectedOutput: &extraSpecs{
+				SecurityGroupIds: []string{"sg-018c35963edfb1cce", "sg-018c35963edfb1cee"},
+			},
+		},
 		{
 			name: "specs just with disable_updates",
 			input: params.BootstrapInstance{
@@ -174,6 +184,14 @@ func TestExtraSpecsFromBootstrapData(t *testing.T) {
 			expectedOutput: nil,
 			errString:      "ssh_key_name: Invalid type. Expected: string, given: integer",
 		},
+		{
+			name: "invalid type for security_group_ids",
+			input: params.BootstrapInstance{
+				ExtraSpecs: json.RawMessage(`{"security_group_ids": "sg-018c35963edfb1cce"}`),
+			},
+			expectedOutput: nil,
+			errString:      "security_group_ids: Invalid type. Expected: array, given: string",
+		},
 		{
 			name: "invalid type for disable_updates",
 			input: params.BootstrapInstance{
@@ -267,7 +285,7 @@ func TestGetRunnerSpecFromBootstrapParams(t *testing.T) {
 
 	data := params.BootstrapInstance{
 		Name:       "mock-name",
-		ExtraSpecs: json.RawMessage(`{"subnet_id": "subnet-0a0a0a0a0a0a0a0a0", "ssh_key_name": "ssh_key_name", "disable_updates": true, "enable_boot_debug": true, "extra_packages": ["package1", "package2"], "runner_install_template": "IyEvYmluL2Jhc2gKZWNobyBJbnN0YWxsaW5nIHJ1bm5lci4uLg==", "pre_install_scripts": {"setup.sh": "IyEvYmluL2Jhc2gKZWNobyBTZXR1cCBzY3JpcHQuLi4="}, "extra_context": {"key": "value"}}`),
+		ExtraSpecs: json.RawMessage(`{"subnet_id": "subnet-0a0a0a0a0a0a0a0a0", "ssh_key_name": "ssh_key_name", "security_group_ids": ["sg-018c35963edfb1cce", "sg-018c35963edfb1cee"], "disable_updates": true, "enable_boot_debug": true, "extra_packages": ["package1", "package2"], "runner_install_template": "IyEvYmluL2Jhc2gKZWNobyBJbnN0YWxsaW5nIHJ1bm5lci4uLg==", "pre_install_scripts": {"setup.sh": "IyEvYmluL2Jhc2gKZWNobyBTZXR1cCBzY3JpcHQuLi4="}, "extra_context": {"key": "value"}}`),
 	}
 
 	config := &config.Config{
@@ -283,15 +301,16 @@ func TestGetRunnerSpecFromBootstrapParams(t *testing.T) {
 		Region:   "region",
 	}
 	expectedRunnerSpec := &RunnerSpec{
-		Region:          "region",
-		DisableUpdates:  true,
-		ExtraPackages:   []string{"package1", "package2"},
-		EnableBootDebug: true,
-		SubnetID:        "subnet-0a0a0a0a0a0a0a0a0",
-		Tools:           Mocktools,
-		ControllerID:    "controller_id",
-		BootstrapParams: data,
-		SSHKeyName:      aws.String("ssh_key_name"),
+		Region:           "region",
+		DisableUpdates:   true,
+		ExtraPackages:    []string{"package1", "package2"},
+		EnableBootDebug:  true,
+		SubnetID:         "subnet-0a0a0a0a0a0a0a0a0",
+		Tools:            Mocktools,
+		ControllerID:     "controller_id",
+		BootstrapParams:  data,
+		SSHKeyName:       aws.String("ssh_key_name"),
+		SecurityGroupIds: []string{"sg-018c35963edfb1cce", "sg-018c35963edfb1cee"},
 	}
 
 	runnerSpec, err := GetRunnerSpecFromBootstrapParams(config, data, "controller_id")