Repository: keikoproj/instance-manager (go)

Issue:
Add UserData support for EKS
**Is this a BUG REPORT or FEATURE REQUEST?**:
Feature Request

**What happened**:

Can't find where to place the UserData to inject into the Autoscale Group through Cloudformation.

**What you expected to happen**:

Be able to inject UserData to CloudFormation

- Kubernetes version:

1.15 , 1.16
```


Pull request:
User data support
Fixes #104 

TBD
- [x] BDD passing 
- [x] Documentation
- [x] Unit tests

Patch:
diff --git a/README.md b/README.md
index bbd858d..e3d5967 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,5 @@
 # instance-manager
+
 [![Build Status](https://travis-ci.org/keikoproj/instance-manager.svg?branch=master)](https://travis-ci.org/keikoproj/instance-manager)
 [![codecov](https://codecov.io/gh/keikoproj/instance-manager/branch/master/graph/badge.svg)](https://codecov.io/gh/keikoproj/instance-manager)
 [![Go Report Card](https://goreportcard.com/badge/github.com/keikoproj/instance-manager)](https://goreportcard.com/report/github.com/keikoproj/instance-manager)
@@ -12,15 +13,14 @@
 
 - [Installation](#installation)
 - [Usage Example](#usage-example)
-    * [Currently supported provisioners](#currently-supported-provisioners)
-    * [EKS sample spec](#EKS-sample-spec)
-    * [Submit and Verify](#submit-and-verify)
-    * [Features](#features)
-    * [Alpha-2 Version](#alpha-2-version)
+  - [Currently supported provisioners](#currently-supported-provisioners)
+  - [EKS sample spec](#EKS-sample-spec)
+  - [Submit and Verify](#submit-and-verify)
+  - [Features](#features)
+  - [Alpha-2 Version](#alpha-2-version)
 - [Contributing](#contributing)
 - [Developer Guide](#developer-guide)
 
-
 Worker nodes in Kubernetes clusters work best if provisioned and managed using a logical grouping. Kops introduced the term “InstanceGroup” for this logical grouping. In AWS, an InstanceGroup maps to an AutoScalingGroup.
 
 Given a particular cluster, there should be a way to create, read, upgrade and delete worker nodes from within the cluster itself. This enables use-cases where worker nodes can be created in response to Kubernetes events, InstanceGroups can be automatically assigned to namespaces for multi-tenancy, etc.
@@ -45,7 +45,6 @@ _For installation instructions and more examples of usage, please refer to the [
 
 To create an instance group, submit an InstanceGroup custom resource in your cluster, and the controller will provision and bootstrap it to your cluster, and allow you to modify it from within the cluster.
 
-
 ### EKS sample spec
 
 ```yaml
@@ -90,7 +89,6 @@ spec:
         value: some-value
 ```
 
-
 ### Submit and Verify
 
 ```bash
@@ -242,6 +240,7 @@ spec:
 ```
 
 #### EKS Fargate
+
 The purpose of the fargate provisioner is to enable the management of Fargate profiles.
 
 By associating EKS clusters with a Fargate Profile, pods can be identified for execution through profile selectors. If a to-be-scheduled pod matches any of the selectors in the Fargate Profile, then that pod is scheduled on Fargate.
@@ -286,8 +285,8 @@ spec:
     tags:
       key1: "value1"
       key2: "value2"
-
 ```
+
 Read more about the [Fargate Profile](https://docs.aws.amazon.com/eks/latest/userguide/fargate-profile.html).
 
 Note that the eks-fargate provisioner does not accept a Fargate profile name.  Instead, the provisioner creates a unique profile name based upon the cluster name, instance group name and namespace.
@@ -307,8 +306,8 @@ Properties:
   ManagedPolicyArns:
   - "arn:aws:iam::aws:policy/AmazonEKSFargatePodExecutionRolePolicy"
   Path: /
-
 ```
+
 Most likely an execution role with access to addtional AWS resources will be required.  In this case, the above IAM role can be used as the basis to create a new, custom role with the IAM policies specific to your pods. Create your new role and your pod specific policies and use the new role's ARN as the *podExecutionRoleArn* parameter value in eks-fargate spec.
 
 Here is an example of a role with an additional policy for S3 access.
diff --git a/api/v1alpha1/instancegroup_types.go b/api/v1alpha1/instancegroup_types.go
index e530c4b..fecf664 100644
--- a/api/v1alpha1/instancegroup_types.go
+++ b/api/v1alpha1/instancegroup_types.go
@@ -49,6 +49,10 @@ const (
 	ReconcileReady ReconcileState = "Ready"
 	ReconcileErr   ReconcileState = "Error"
 
+	// Userdata bootstrap stages
+	PreBootstrapStage  = "PreBootstrap"
+	PostBootstrapStage = "PostBootstrap"
+
 	LifecycleStateNormal      = "normal"
 	LifecycleStateSpot        = "spot"
 	CRDStrategyName           = "crd"
@@ -174,12 +178,18 @@ type EKSConfiguration struct {
 	Tags                        []map[string]string `json:"tags,omitempty"`
 	Labels                      map[string]string   `json:"labels,omitempty"`
 	Taints                      []corev1.Taint      `json:"taints,omitempty"`
+	UserData                    []UserDataStage     `json:"userData,omitempty"`
 	ExistingRoleName            string              `json:"roleName,omitempty"`
 	ExistingInstanceProfileName string              `json:"instanceProfileName,omitempty"`
 	ManagedPolicies             []string            `json:"managedPolicies,omitempty"`
 	MetricsCollection           []string            `json:"metricsCollection,omitempty"`
 }
 
+type UserDataStage struct {
+	Stage string `json:"stage,omitempty"`
+	Data  string `json:"data,omitempty"`
+}
+
 type NodeVolume struct {
 	Name string `json:"name,omitempty"`
 	Type string `json:"type,omitempty"`
@@ -396,6 +406,9 @@ func (c *EKSConfiguration) SetTaints(taints []corev1.Taint) {
 func (c *EKSConfiguration) GetManagedPolicies() []string {
 	return c.ManagedPolicies
 }
+func (c *EKSConfiguration) GetUserData() []UserDataStage {
+	return c.UserData
+}
 func (c *EKSConfiguration) SetManagedPolicies(policies []string) {
 	c.ManagedPolicies = policies
 }
diff --git a/api/v1alpha1/zz_generated.deepcopy.go b/api/v1alpha1/zz_generated.deepcopy.go
index 11ff037..7cd7008 100644
--- a/api/v1alpha1/zz_generated.deepcopy.go
+++ b/api/v1alpha1/zz_generated.deepcopy.go
@@ -21,7 +21,7 @@ package v1alpha1
 
 import (
 	"k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/runtime"
+	runtime "k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/util/intstr"
 )
 
@@ -115,6 +115,11 @@ func (in *EKSConfiguration) DeepCopyInto(out *EKSConfiguration) {
 			(*in)[i].DeepCopyInto(&(*out)[i])
 		}
 	}
+	if in.UserData != nil {
+		in, out := &in.UserData, &out.UserData
+		*out = make([]UserDataStage, len(*in))
+		copy(*out, *in)
+	}
 	if in.ManagedPolicies != nil {
 		in, out := &in.ManagedPolicies, &out.ManagedPolicies
 		*out = make([]string, len(*in))
@@ -443,3 +448,18 @@ func (in *RollingUpdateStrategy) DeepCopy() *RollingUpdateStrategy {
 	in.DeepCopyInto(out)
 	return out
 }
+
+// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
+func (in *UserDataStage) DeepCopyInto(out *UserDataStage) {
+	*out = *in
+}
+
+// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new UserDataStage.
+func (in *UserDataStage) DeepCopy() *UserDataStage {
+	if in == nil {
+		return nil
+	}
+	out := new(UserDataStage)
+	in.DeepCopyInto(out)
+	return out
+}
diff --git a/config/crd/bases/instancemgr.keikoproj.io_instancegroups.yaml b/config/crd/bases/instancemgr.keikoproj.io_instancegroups.yaml
index c30ff5f..0e6fa04 100644
--- a/config/crd/bases/instancemgr.keikoproj.io_instancegroups.yaml
+++ b/config/crd/bases/instancemgr.keikoproj.io_instancegroups.yaml
@@ -149,6 +149,15 @@ spec:
                         - key
                         type: object
                       type: array
+                    userData:
+                      items:
+                        properties:
+                          data:
+                            type: string
+                          stage:
+                            type: string
+                        type: object
+                      type: array
                     volumes:
                       items:
                         properties:
diff --git a/controllers/providers/aws/aws.go b/controllers/providers/aws/aws.go
index 011ca8e..457c4fd 100644
--- a/controllers/providers/aws/aws.go
+++ b/controllers/providers/aws/aws.go
@@ -16,10 +16,11 @@ limitations under the License.
 package aws
 
 import (
+	"bytes"
 	"encoding/base64"
-	"fmt"
 	"os"
 	"strings"
+	"text/template"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -101,10 +102,29 @@ const (
 	IAMPolicyPrefix = "arn:aws:iam::aws:policy"
 )
 
-func DefaultEksUserDataFmt() string {
-	return `#!/bin/bash
+type EKSUserData struct {
+	ClusterName   string
+	Arguments     string
+	PreBootstrap  string
+	PostBootstrap string
+}
+
+var UserDataPayload = `#!/bin/bash
+{{ .PreBootstrap }}
 set -o xtrace
-/etc/eks/bootstrap.sh %s %s`
+/etc/eks/bootstrap.sh {{ .ClusterName }} {{ .Arguments }}
+set +o xtrace
+{{ .PostBootstrap }}`
+
+func RenderUserData(data EKSUserData) []byte {
+	out := &bytes.Buffer{}
+	tmpl := template.New("userData")
+	var err error
+	if tmpl, err = tmpl.Parse(UserDataPayload); err != nil {
+		log.Error(err, "failed to parse userData template")
+	}
+	tmpl.Execute(out, data)
+	return out.Bytes()
 }
 
 func (w *AwsWorker) RoleExist(name string) (*iam.Role, bool) {
@@ -234,9 +254,14 @@ func (w *AwsWorker) SetResumeProcesses(name string, processesToResume []string)
 	return nil
 }
 
-func (w *AwsWorker) GetBasicUserData(clusterName, bootstrapArgs string) string {
-	userData := fmt.Sprintf(DefaultEksUserDataFmt(), clusterName, bootstrapArgs)
-	return base64.StdEncoding.EncodeToString([]byte(userData))
+func (w *AwsWorker) GetBasicUserData(clusterName, args, preScript, postScript string) string {
+	userData := RenderUserData(EKSUserData{
+		ClusterName:   clusterName,
+		Arguments:     args,
+		PreBootstrap:  preScript,
+		PostBootstrap: postScript,
+	})
+	return base64.StdEncoding.EncodeToString(userData)
 }
 
 func (w *AwsWorker) NewTag(key, val, resource string) *autoscaling.Tag {
diff --git a/controllers/provisioners/eks/helpers.go b/controllers/provisioners/eks/helpers.go
index cafe914..f6733f8 100644
--- a/controllers/provisioners/eks/helpers.go
+++ b/controllers/provisioners/eks/helpers.go
@@ -37,14 +37,15 @@ import (
 
 func (ctx *EksInstanceGroupContext) GetLaunchConfigurationInput(name string) *autoscaling.CreateLaunchConfigurationInput {
 	var (
-		instanceGroup   = ctx.GetInstanceGroup()
-		configuration   = instanceGroup.GetEKSConfiguration()
-		clusterName     = configuration.GetClusterName()
-		state           = ctx.GetDiscoveredState()
-		instanceProfile = state.GetInstanceProfile()
-		devices         = ctx.GetBlockDeviceList()
-		args            = ctx.GetBootstrapArgs()
-		userData        = ctx.AwsWorker.GetBasicUserData(clusterName, args)
+		instanceGroup         = ctx.GetInstanceGroup()
+		configuration         = instanceGroup.GetEKSConfiguration()
+		clusterName           = configuration.GetClusterName()
+		state                 = ctx.GetDiscoveredState()
+		instanceProfile       = state.GetInstanceProfile()
+		devices               = ctx.GetBlockDeviceList()
+		args                  = ctx.GetBootstrapArgs()
+		preScript, postScript = ctx.GetUserDataStages()
+		userData              = ctx.AwsWorker.GetBasicUserData(clusterName, args, preScript, postScript)
 	)
 
 	input := &autoscaling.CreateLaunchConfigurationInput{
@@ -65,6 +66,27 @@ func (ctx *EksInstanceGroupContext) GetLaunchConfigurationInput(name string) *au
 	return input
 }
 
+func (ctx *EksInstanceGroupContext) GetUserDataStages() (string, string) {
+	var (
+		instanceGroup = ctx.GetInstanceGroup()
+		configuration = instanceGroup.GetEKSConfiguration()
+		userData      = configuration.GetUserData()
+	)
+
+	var preScript, postScript string
+	for _, stage := range userData {
+		switch {
+		case strings.EqualFold(stage.Stage, v1alpha1.PreBootstrapStage):
+			preScript = stage.Data
+		case strings.EqualFold(stage.Stage, v1alpha1.PostBootstrapStage):
+			postScript = stage.Data
+		default:
+			ctx.Log.Info("invalid userdata stage will not be rendered", "stage", stage.Stage, "data", stage.Data)
+		}
+	}
+	return preScript, postScript
+}
+
 func (ctx *EksInstanceGroupContext) GetAddedTags(asgName string) []*autoscaling.Tag {
 	var (
 		tags          []*autoscaling.Tag
diff --git a/docs/EKS.md b/docs/EKS.md
index 0494de1..276e07b 100644
--- a/docs/EKS.md
+++ b/docs/EKS.md
@@ -82,10 +82,29 @@ spec:
       # GroupTotalCapacity
       # All (will enable all above metrics)
       metricsCollection: <[]string> : must be a list of metric names to enable collection for
+
+      # customize UserData passed into launch configuration
+      userData: <[]UserDataStage> : must be a list of UserDataStage
+```
+
+### UserDataStage
+
+UserDataStage represents a custom userData script
+
+```yaml
+spec:
+  provisioner: eks
+  eks:
+    configuration:
+      userData:
+      - stage: <string> : represents the stage of the script, allowed values are PreBootstrap, PostBootstrap
+        data: <string> : represents the script payload to inject
 ```
 
 ### NodeVolume
 
+NodeVolume represents a custom EBS volume
+
 ```yaml
 spec:
   provisioner: eks
diff --git a/docs/README.md b/docs/README.md
new file mode 100644
index 0000000..5d2fb57
--- /dev/null
+++ b/docs/README.md
@@ -0,0 +1,10 @@
+# Documentation
+
+- [Installation](install)
+- [FAQ](faq)
+- [EKS Provisioner API Reference](eks)
+
+<!-- Markdown link -->
+[install]: https://github.com/keikoproj/instance-manager/blob/master/docs/INSTALL.md
+[faq]: https://github.com/keikoproj/instance-manager/blob/master/docs/FAQ.md
+[eks]: https://github.com/keikoproj/instance-manager/blob/master/docs/EKS.md


Test patch:
diff --git a/controllers/provisioners/eks/helpers_test.go b/controllers/provisioners/eks/helpers_test.go
index b81a2ac..8af29e0 100644
--- a/controllers/provisioners/eks/helpers_test.go
+++ b/controllers/provisioners/eks/helpers_test.go
@@ -185,6 +185,49 @@ func TestGetLabelList(t *testing.T) {
 	}
 }
 
+func TestGetUserDataStages(t *testing.T) {
+	var (
+		g             = gomega.NewGomegaWithT(t)
+		k             = MockKubernetesClientSet()
+		ig            = MockInstanceGroup()
+		configuration = ig.GetEKSConfiguration()
+		asgMock       = NewAutoScalingMocker()
+		iamMock       = NewIamMocker()
+		eksMock       = NewEksMocker()
+	)
+
+	w := MockAwsWorker(asgMock, iamMock, eksMock)
+	ctx := MockContext(ig, k, w)
+
+	tests := []struct {
+		preBootstrapScript  string
+		postBootstrapScript string
+	}{
+		{},
+		{preBootstrapScript: "", postBootstrapScript: ""},
+		{preBootstrapScript: "prebootstrap", postBootstrapScript: "postbootstrap"},
+	}
+
+	for i, tc := range tests {
+		t.Logf("Test #%v - %+v", i, tc)
+		configuration.UserData = append(configuration.UserData, v1alpha1.UserDataStage{
+			Stage: v1alpha1.PreBootstrapStage,
+			Data:  tc.preBootstrapScript,
+		})
+		configuration.UserData = append(configuration.UserData, v1alpha1.UserDataStage{
+			Stage: v1alpha1.PostBootstrapStage,
+			Data:  tc.postBootstrapScript,
+		})
+		configuration.UserData = append(configuration.UserData, v1alpha1.UserDataStage{
+			Stage: "invalid-stage",
+			Data:  "test",
+		})
+		preScript, postScript := ctx.GetUserDataStages()
+		g.Expect(preScript).To(gomega.Equal(tc.preBootstrapScript))
+		g.Expect(postScript).To(gomega.Equal(tc.postBootstrapScript))
+	}
+}
+
 func TestIsRetryable(t *testing.T) {
 	var (
 		g  = gomega.NewGomegaWithT(t)