From d124f2243e5583e1f41b444b262411a1da8ccfdc Mon Sep 17 00:00:00 2001
From: Grzegorz Nosek <grzegorz.nosek@sysdig.com>
Date: Fri, 27 Oct 2023 17:29:21 +0200
Subject: [PATCH] feat: Configurable task capabilities

Signed-off-by: Grzegorz Nosek <grzegorz.nosek@sysdig.com>
---
 pkg/hocon/build.go                            |  5 ++
 pkg/kilt/types.go                             |  1 +
 .../cloudformation/cfnpatcher/cfn_test.go     |  3 ++
 runtimes/cloudformation/cfnpatcher/patcher.go | 47 ++++++++++++-------
 4 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/pkg/hocon/build.go b/pkg/hocon/build.go
index 648c777..a88c33f 100644
--- a/pkg/hocon/build.go
+++ b/pkg/hocon/build.go
@@ -21,6 +21,11 @@ func extractBuild(config *configuration.Config) (*kilt.Build, error) {
 		b.Command = make([]string, 0)
 	}
 
+	b.Capabilities = config.GetStringList("build.capabilities")
+	if b.Capabilities == nil {
+		b.Capabilities = make([]string, 0)
+	}
+
 	b.EnvironmentVariables = extractToStringMap(config, "build.environment_variables")
 
 	if config.IsArray("build.mount") {
diff --git a/pkg/kilt/types.go b/pkg/kilt/types.go
index 1d78952..93f60a8 100644
--- a/pkg/kilt/types.go
+++ b/pkg/kilt/types.go
@@ -29,6 +29,7 @@ type Build struct {
 	EntryPoint           []string
 	Command              []string
 	EnvironmentVariables map[string]string
+	Capabilities         []string
 
 	Resources []BuildResource
 }
diff --git a/runtimes/cloudformation/cfnpatcher/cfn_test.go b/runtimes/cloudformation/cfnpatcher/cfn_test.go
index 38a4e0d..ebae41e 100644
--- a/runtimes/cloudformation/cfnpatcher/cfn_test.go
+++ b/runtimes/cloudformation/cfnpatcher/cfn_test.go
@@ -70,6 +70,7 @@ build {
 			entry_point: ["/kilt/wait"]
 		}
 	]
+	capabilities: ["SYS_PTRACE"]
 }
 `
 
@@ -88,6 +89,7 @@ build {
 			entry_point: ["/kilt/wait"]
 		}
 	]
+	capabilities: ["SYS_PTRACE"]
 }
 `
 
@@ -106,6 +108,7 @@ build {
 			}
 		}
 	]
+	capabilities: ["SYS_PTRACE"]
 }
 `
 
diff --git a/runtimes/cloudformation/cfnpatcher/patcher.go b/runtimes/cloudformation/cfnpatcher/patcher.go
index 4b2de2f..6747691 100644
--- a/runtimes/cloudformation/cfnpatcher/patcher.go
+++ b/runtimes/cloudformation/cfnpatcher/patcher.go
@@ -206,27 +206,40 @@ func applyContainerDefinitionPatch(ctx context.Context, container *gabs.Containe
 
 	}
 
-	// We need to add SYS_PTRACE capability to the container
-	if !container.Exists("LinuxParameters") {
-		emptyMap := make(map[string]interface{})
-		_, err = container.Set(emptyMap, "LinuxParameters")
-		if err != nil {
-			return fmt.Errorf("could not add LinuxParameters: %w", err)
+	if len(patch.Capabilities) > 0 {
+		capabilities := make([]interface{}, len(patch.Capabilities))
+		for i, v := range patch.Capabilities {
+			capabilities[i] = v
+		}
+		// We need to add capabilities to the container
+		if !container.Exists("LinuxParameters") {
+			emptyMap := make(map[string]interface{})
+			_, err = container.Set(emptyMap, "LinuxParameters")
+			if err != nil {
+				return fmt.Errorf("could not add LinuxParameters: %w", err)
+			}
 		}
-	}
 
-	if !container.Exists("LinuxParameters", "Capabilities") {
-		emptyMap := make(map[string]interface{})
-		_, err = container.Set(emptyMap, "LinuxParameters", "Capabilities")
-		if err != nil {
-			return fmt.Errorf("could not add LinuxParameters.Capabilities: %w", err)
+		if !container.Exists("LinuxParameters", "Capabilities") {
+			emptyMap := make(map[string]interface{})
+			_, err = container.Set(emptyMap, "LinuxParameters", "Capabilities")
+			if err != nil {
+				return fmt.Errorf("could not add LinuxParameters.Capabilities: %w", err)
+			}
 		}
-	}
 
-	// fargate only supports SYS_PTRACE
-	_, err = container.Set([]string{"SYS_PTRACE"}, "LinuxParameters", "Capabilities", "Add")
-	if err != nil {
-		return fmt.Errorf("could not add LinuxParamaters.Capabilities.Add: %w", err)
+		if !container.Exists("LinuxParameters", "Capabilities", "Add") {
+			emptyList := make([]interface{}, 0)
+			_, err = container.Set(emptyList, "LinuxParameters", "Capabilities", "Add")
+			if err != nil {
+				return fmt.Errorf("could not add LinuxParameters.Capabilities.Add: %w", err)
+			}
+		}
+
+		err := container.ArrayConcat(capabilities, "LinuxParameters", "Capabilities", "Add")
+		if err != nil {
+			return fmt.Errorf("could not append to LinuxParameters.Capabilities.Add: %w", err)
+		}
 	}
 
 	return nil