Merge pull request #221 from daimaxiaxie/enhance-userdata

Enhance userdata to support mime

Author: jwcesign

charts/karpenter/crds/karpenter.k8s.alibabacloud_ecsnodeclasses.yaml

@@ -90,7 +90,7 @@ spec:
                       x-kubernetes-validations:
                       - message: '''alias'' is improperly formatted, must match the
                           format ''family'''
-                        rule: self.matches('^[a-zA-Z0-9]*$')
+                        rule: self.matches('^[a-zA-Z0-9]+@.+$')
                       - message: 'family is not supported, must be one of the following:
                           ''AlibabaCloudLinux3,ContainerOS'''
                         rule: self.find('^[^@]+') in ['AlibabaCloudLinux3', 'ContainerOS']

go.mod

@@ -9,9 +9,12 @@ require (
 	github.com/alibabacloud-go/tea v1.2.2
 	github.com/alibabacloud-go/tea-utils/v2 v2.0.7
 	github.com/alibabacloud-go/vpc-20160428/v6 v6.11.2
+	github.com/aws/karpenter-provider-aws v1.2.1
 	github.com/awslabs/operatorpkg v0.0.0-20250121140423-041752c305f4
 	github.com/cloudpilot-ai/priceserver v0.0.0-20241011010411-15ac0e19a857
 	github.com/mitchellh/hashstructure/v2 v2.0.2
+	github.com/onsi/ginkgo/v2 v2.22.2
+	github.com/onsi/gomega v1.36.2
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/samber/lo v1.49.1
 	github.com/stretchr/testify v1.10.0
@@ -27,10 +30,14 @@ require (
 )
 
 require (
+	github.com/awslabs/amazon-eks-ami/nodeadm v0.0.0-20240229193347-cfab22a10647 // indirect
 	github.com/fxamacker/cbor/v2 v2.7.0 // indirect
+	github.com/go-task/slim-sprig/v3 v3.0.0 // indirect
 	github.com/google/btree v1.1.3 // indirect
+	github.com/google/pprof v0.0.0-20241210010833-40e02aabc2ad // indirect
 	github.com/klauspost/compress v1.17.11 // indirect
 	github.com/x448/float16 v0.8.4 // indirect
+	golang.org/x/tools v0.29.0 // indirect
 	gopkg.in/evanphx/json-patch.v4 v4.12.0 // indirect
 )
 

go.sum

@@ -58,6 +58,10 @@ github.com/aliyun/credentials-go v1.4.3 h1:N3iHyvHRMyOwY1+0qBLSf3hb5JFiOujVSVuEp
 github.com/aliyun/credentials-go v1.4.3/go.mod h1:Jm6d+xIgwJVLVWT561vy67ZRP4lPTQxMbEYRuT2Ti1U=
 github.com/avast/retry-go v3.0.0+incompatible h1:4SOWQ7Qs+oroOTQOYnAHqelpCO0biHSxpiH9JdtuBj0=
 github.com/avast/retry-go v3.0.0+incompatible/go.mod h1:XtSnn+n/sHqQIpZ10K1qAevBhOOCWBLXXy3hyiqqBrY=
+github.com/aws/karpenter-provider-aws v1.2.1 h1:JPZ4EPM9iJ+lga6zDpsrHZuhojgKbrgWPqnwTFE3gNw=
+github.com/aws/karpenter-provider-aws v1.2.1/go.mod h1:ojjedpNzBJ0uKJ0Ge8/spDLbXWHIKGDbk8MUWh6pklE=
+github.com/awslabs/amazon-eks-ami/nodeadm v0.0.0-20240229193347-cfab22a10647 h1:8yRBVsjGmI7qQsPWtIrbWP+XfwHO9Wq7gdLVzjqiZFs=
+github.com/awslabs/amazon-eks-ami/nodeadm v0.0.0-20240229193347-cfab22a10647/go.mod h1:9NafTAUHL0FlMeL6Cu5PXnMZ1q/LnC9X2emLXHsVbM8=
 github.com/awslabs/operatorpkg v0.0.0-20250121140423-041752c305f4 h1:o5eK/Sh2Ndkv1xOw0Y3Wd9pIoFum8dvNY8MBYbFD9Rw=
 github.com/awslabs/operatorpkg v0.0.0-20250121140423-041752c305f4/go.mod h1:ajddvYQzaxpfjeVxLa02iedZnExetLxx45m3S86Dhj0=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=

pkg/apis/v1alpha1/ecsnodeclass.go

@@ -153,7 +153,7 @@ type ImageSelectorTerm struct {
 	// Valid families include: AlibabaCloudLinux3,ContainerOS
 	// Currently only supports version pinning to the latest image release, with that images version format (ex: "aliyun3@latest").
 	// Setting the version to latest will result in drift when a new Image is released. This is **not** recommended for production environments.
-	// +kubebuilder:validation:XValidation:message="'alias' is improperly formatted, must match the format 'family'",rule="self.matches('^[a-zA-Z0-9]*$')"
+	// +kubebuilder:validation:XValidation:message="'alias' is improperly formatted, must match the format 'family'",rule="self.matches('^[a-zA-Z0-9]+@.+$')"
 	// +kubebuilder:validation:XValidation:message="family is not supported, must be one of the following: 'AlibabaCloudLinux3,ContainerOS'",rule="self.find('^[^@]+') in ['AlibabaCloudLinux3', 'ContainerOS']"
 	// +kubebuilder:validation:MaxLength=30
 	// +optional

pkg/providers/cluster/ackmanaged.go

@@ -143,8 +143,40 @@ func (a *ACKManaged) UserData(ctx context.Context,
 	taints []corev1.Taint,
 	kubeletCfg *v1alpha1.KubeletConfiguration,
 	userData *string) (string, error) {
+
+	attach, err := a.getClusterAttachScripts(ctx)
+	if err != nil {
+		return "", err
+	}
+	ackScript := a.ackBootstrap(attach, labels, taints, kubeletCfg)
+	cloudInit := NewCloudInit()
+
+	if err := cloudInit.Merge(&ackScript); err != nil {
+		return "", err
+	}
+	if err := cloudInit.Merge(userData); err != nil {
+		return "", err
+	}
+
+	return cloudInit.Script()
+}
+
+func (a *ACKManaged) FeatureFlags() FeatureFlags {
+	if cni, err := a.GetClusterCNI(context.TODO()); err == nil && cni == ClusterCNITypeFlannel {
+		return FeatureFlags{
+			PodsPerCoreEnabled:           false,
+			SupportsENILimitedPodDensity: false,
+		}
+	}
+	return FeatureFlags{
+		PodsPerCoreEnabled:           true,
+		SupportsENILimitedPodDensity: true,
+	}
+}
+
+func (a *ACKManaged) getClusterAttachScripts(ctx context.Context) (string, error) {
 	if cachedScript, ok := a.cache.Get(a.clusterID); ok {
-		return a.resolveUserData(cachedScript.(string), labels, taints, kubeletCfg, userData), nil
+		return cachedScript.(string), nil
 	}
 
 	reqPara := &ackclient.DescribeClusterAttachScriptsRequest{
@@ -174,20 +206,7 @@ func (a *ACKManaged) UserData(ctx context.Context,
 	}
 
 	a.cache.SetDefault(a.clusterID, respStr)
-	return a.resolveUserData(respStr, labels, taints, kubeletCfg, userData), nil
-}
-
-func (a *ACKManaged) FeatureFlags() FeatureFlags {
-	if cni, err := a.GetClusterCNI(context.TODO()); err == nil && cni == ClusterCNITypeFlannel {
-		return FeatureFlags{
-			PodsPerCoreEnabled:           false,
-			SupportsENILimitedPodDensity: false,
-		}
-	}
-	return FeatureFlags{
-		PodsPerCoreEnabled:           true,
-		SupportsENILimitedPodDensity: true,
-	}
+	return respStr, nil
 }
 
 // We need to manually retrieve the runtime configuration of the nodepool, with the default node pool prioritized.
@@ -239,21 +258,13 @@ func (a *ACKManaged) getClusterAttachRuntimeConfiguration(ctx context.Context) (
 		tea.StringValue(targetNodepool.KubernetesConfig.RuntimeVersion), nil
 }
 
-func (a *ACKManaged) resolveUserData(respStr string, labels map[string]string, taints []corev1.Taint,
-	kubeletCfg *v1alpha1.KubeletConfiguration, userData *string) string {
-	preUserData, postUserData := parseCustomUserData(userData)
+func (a *ACKManaged) ackBootstrap(respStr string, labels map[string]string, taints []corev1.Taint,
+	kubeletCfg *v1alpha1.KubeletConfiguration) string {
 
 	var script bytes.Buffer
 	// Add bash script header
 	script.WriteString("#!/bin/bash\n\n")
 
-	// Insert preUserData if available
-	if preUserData != "" {
-		// Pre-userData: scripts to be executed before node registration
-		script.WriteString("echo \"Executing preUserData...\"\n")
-		script.WriteString(preUserData + "\n\n")
-	}
-
 	// Clean up the input string
 	script.WriteString(respStr + " ")
 	// Add labels
@@ -264,15 +275,7 @@ func (a *ACKManaged) resolveUserData(respStr string, labels map[string]string, t
 	// Add taints
 	script.WriteString(fmt.Sprintf("--taints %s\n\n", a.formatTaints(taints)))
 
-	// Insert postUserData if available
-	if postUserData != "" {
-		// Post-userData: scripts to be executed after node registration
-		script.WriteString("echo \"Executing postUserData...\"\n")
-		script.WriteString(postUserData + "\n")
-	}
-
-	// Encode to base64
-	return base64.StdEncoding.EncodeToString(script.Bytes())
+	return script.String()
 }
 
 func (a *ACKManaged) formatLabels(labels map[string]string) string {
@@ -310,16 +313,3 @@ func convertNodeClassKubeletConfigToACKNodeConfig(kubeletCfg *v1alpha1.KubeletCo
 	}
 	return base64.StdEncoding.EncodeToString(data)
 }
-
-const userDataSeparator = "#===USERDATA_SEPARATOR==="
-
-// By default, the UserData is executed after the node registration is completed.
-// If a user requires tasks to be executed both before and after node registration,
-// they must split the userdata into preUserData and postUserData using a SEPARATOR.
-func parseCustomUserData(userData *string) (string, string) {
-	parts := strings.Split(tea.StringValue(userData), userDataSeparator)
-	if len(parts) == 2 {
-		return strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1])
-	}
-	return "", tea.StringValue(userData)
-}

pkg/providers/cluster/cloudinit.go

@@ -0,0 +1,90 @@
+/*
+Copyright 2024 The CloudPilot AI Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package cluster
+
+import (
+	"mime"
+	"strings"
+
+	awsmime "github.com/aws/karpenter-provider-aws/pkg/providers/amifamily/bootstrap/mime"
+	"github.com/samber/lo"
+)
+
+const (
+	contentTypeStage                string = `stage`
+	contentTypeShellScriptMediaType string = `text/x-shellscript`
+)
+
+type CloudInit struct {
+	entries []awsmime.Entry
+}
+
+func NewCloudInit() *CloudInit {
+	return &CloudInit{
+		entries: make([]awsmime.Entry, 0),
+	}
+}
+
+func (c *CloudInit) Script() (string, error) {
+	c.sort()
+	mimeArchive := awsmime.Archive(c.entries)
+	userData, err := mimeArchive.Serialize()
+	if err != nil {
+		return "", err
+	}
+	return userData, nil
+}
+
+func (c *CloudInit) Merge(userdata *string) error {
+	userData := lo.FromPtr(userdata)
+	if userData == "" {
+		return nil
+	}
+	if strings.HasPrefix(strings.TrimSpace(userData), "MIME-Version:") ||
+		strings.HasPrefix(strings.TrimSpace(userData), "Content-Type:") {
+		archive, err := awsmime.NewArchive(userData)
+		if err != nil {
+			return err
+		}
+		c.entries = append(c.entries, archive...)
+		return nil
+	}
+	// Fallback to YAML or shall script if UserData is not in MIME format. Determine the content type for the
+	// generated MIME header depending on the type of the custom UserData.
+	c.entries = append(c.entries, awsmime.Entry{
+		ContentType: awsmime.ContentTypeShellScript,
+		Content:     userData,
+	})
+	return nil
+}
+
+func (c *CloudInit) sort() {
+	var pre, non []awsmime.Entry
+	for _, entry := range c.entries {
+		mediaType, params, err := mime.ParseMediaType(string(entry.ContentType))
+		if err != nil {
+			non = append(non, entry)
+			continue
+		}
+		if stage, ok := params[contentTypeStage]; ok && mediaType == contentTypeShellScriptMediaType && stage == "pre" {
+			pre = append(pre, entry)
+		} else {
+			non = append(non, entry)
+		}
+	}
+	c.entries = append(pre, non...)
+}