fix: detect the Toolkit properly

kubeshop · Jul 25, 2024 · 421317f · 421317f
1 parent bc0eb69
commit 421317f
Show file tree

Hide file tree

Showing 8 changed files with 28 additions and 15 deletions.
diff --git a/cmd/testworkflow-init/main.go b/cmd/testworkflow-init/main.go
@@ -235,6 +235,12 @@ func main() {
 				continue
 			}
 
+			// Configure the environment
+			orchestration.Setup.UseCurrentEnv()
+			if !action.Execute.Toolkit {
+				_ = os.Unsetenv("TK_REF")
+			}
+
 			// List all the parents
 			leaf := []*data.StepData{step}
 			for i := range step.Parents {

diff --git a/cmd/testworkflow-init/orchestration/setup.go b/cmd/testworkflow-init/orchestration/setup.go
@@ -156,11 +156,6 @@ func (c *setup) UseEnv(group string) {
 		_ = os.Setenv("PATH", fmt.Sprintf("%s:%s", os.Getenv("PATH"), data.InternalBinPath))
 	}
 
-	// TODO: Delete unnecessary variables for non-toolkit operations
-	//if !toolkit {
-	//	_ = os.Unsetenv("TK_REF")
-	//}
-
 	// Compute dynamic environment variables
 	addonMachine := expressions.CombinedMachines(data.RefSuccessMachine, data.AliasMachine, data.StateMachine, libs.NewFsMachine(os.DirFS("/"), cwd))
 	localEnvMachine := expressions.NewMachine().

diff --git a/pkg/testworkflows/testworkflowprocessor/action/actiontypes/lite/action.go b/pkg/testworkflows/testworkflowprocessor/action/actiontypes/lite/action.go
@@ -19,6 +19,7 @@ type ActionDeclare struct {
 type ActionExecute struct {
 	Ref      string `json:"r"`
 	Negative bool   `json:"n,omitempty"`
+	Toolkit  bool   `json:"t,omitempty"`
 }
 
 type ActionPause struct {

diff --git a/pkg/testworkflows/testworkflowprocessor/action/containerize.go b/pkg/testworkflows/testworkflowprocessor/action/containerize.go
@@ -14,7 +14,6 @@ import (
 	stage2 "github.com/kubeshop/testkube/pkg/testworkflows/testworkflowprocessor/stage"
 )
 
-// TODO: Disallow bypassing
 func CreateContainer(groupId int, defaultContainer stage2.Container, actions []actiontypes.Action) (cr corev1.Container, actionsCleanup []actiontypes.Action, err error) {
 	actions = slices.Clone(actions)
 	actionsCleanup = actions

diff --git a/pkg/testworkflows/testworkflowprocessor/action/process.go b/pkg/testworkflows/testworkflowprocessor/action/process.go
@@ -78,6 +78,7 @@ func process(currentStatus string, parents []string, stage stage2.Stage, machine
 			Execute: &lite.ActionExecute{
 				Ref:      exec.Ref(),
 				Negative: exec.Negative(),
+				Toolkit:  exec.IsToolkit(),
 			},
 		})
 	}

diff --git a/pkg/testworkflows/testworkflowprocessor/stage/container.go b/pkg/testworkflows/testworkflowprocessor/stage/container.go
@@ -16,8 +16,9 @@ import (
 )
 
 type container struct {
-	parent *container
-	Cr     testworkflowsv1.ContainerConfig `expr:"include"`
+	parent  *container
+	toolkit bool
+	Cr      testworkflowsv1.ContainerConfig `expr:"include"`
 }
 
 type ContainerComposition interface {
@@ -47,6 +48,7 @@ type ContainerAccessors interface {
 
 	HasVolumeAt(path string) bool
 	ToContainerConfig() testworkflowsv1.ContainerConfig
+	IsToolkit() bool
 }
 
 type ContainerMutations[T any] interface {
@@ -410,6 +412,10 @@ func (c *container) ApplyImageData(image *imageinspector.Info, resolvedImageName
 	return nil
 }
 
+func (c *container) IsToolkit() bool {
+	return c.toolkit || slices.Contains(c.Env(), BypassToolkitCheck)
+}
+
 func (c *container) EnableToolkit(ref string) Container {
 	return c.
 		AppendEnv(corev1.EnvVar{

diff --git a/pkg/testworkflows/testworkflowprocessor/stage/containerstage.go b/pkg/testworkflows/testworkflowprocessor/stage/containerstage.go
@@ -2,11 +2,18 @@ package stage
 
 import (
 	"github.com/pkg/errors"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/util/rand"
 
 	"github.com/kubeshop/testkube/pkg/expressions"
 	"github.com/kubeshop/testkube/pkg/imageinspector"
 )
 
+var BypassToolkitCheck = corev1.EnvVar{
+	Name:  "TK_TC_SECURITY",
+	Value: rand.String(20),
+}
+
 type containerStage struct {
 	stageMetadata
 	stageLifecycle
@@ -16,6 +23,7 @@ type containerStage struct {
 type ContainerStage interface {
 	Stage
 	Container() Container
+	IsToolkit() bool
 }
 
 func NewContainerStage(ref string, container Container) ContainerStage {
@@ -72,3 +80,7 @@ func (s *containerStage) Container() Container {
 func (s *containerStage) HasPause() bool {
 	return s.paused
 }
+
+func (s *containerStage) IsToolkit() bool {
+	return s.container.IsToolkit()
+}
diff --git a/pkg/testworkflows/testworkflowprocessor/utils.go b/pkg/testworkflows/testworkflowprocessor/utils.go
@@ -2,18 +2,11 @@ package testworkflowprocessor
 
 import (
 	batchv1 "k8s.io/api/batch/v1"
-	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/apimachinery/pkg/util/rand"
 
 	"github.com/kubeshop/testkube/pkg/testworkflows/testworkflowprocessor/constants"
 )
 
-var BypassToolkitCheck = corev1.EnvVar{
-	Name:  "TK_TC_SECURITY",
-	Value: rand.String(20),
-}
-
 func AnnotateControlledBy(obj metav1.Object, rootId, id string) {
 	labels := obj.GetLabels()
 	if labels == nil {