Fix KICS container shutdown race condition and add OneAssist license support

- Create kicsshutdown package with thread-safe container name management
- Update signal handler to read container name from kicsshutdown instead of viper
- Prevents race conditions during SIGTERM cleanup
- Add support for OneAssist license in addition to Developer Assist
- Update GetUniqueID() to check both license types

Author: cx-atish-jadhav

cmd/main.go

@@ -9,6 +9,7 @@ import (
 	"syscall"
 
 	"github.com/checkmarx/ast-cli/internal/commands"
+	"github.com/checkmarx/ast-cli/internal/kicsshutdown"
 	"github.com/checkmarx/ast-cli/internal/logger"
 	"github.com/checkmarx/ast-cli/internal/params"
 	"github.com/checkmarx/ast-cli/internal/wrappers"
@@ -191,10 +192,6 @@ func exitListener() {
 }
 
 func signalHandler(signalChanel chan os.Signal) {
-	kicsRunArgs := []string{
-		killCommand,
-		viper.GetString(params.KicsContainerNameKey),
-	}
 	for {
 		s := <-signalChanel
 		switch s {
@@ -204,7 +201,12 @@ func signalHandler(signalChanel chan os.Signal) {
 				os.Exit(failureExitCode)
 			}
 			logger.PrintIfVerbose(string(out))
-			if strings.Contains(string(out), viper.GetString(params.KicsContainerNameKey)) {
+			kicsContainerName := kicsshutdown.GetKicsContainerName()
+			if kicsContainerName != "" && strings.Contains(string(out), kicsContainerName) {
+				kicsRunArgs := []string{
+					killCommand,
+					kicsContainerName,
+				}
 				out, err = exec.Command("docker", kicsRunArgs...).CombinedOutput()
 				logger.PrintIfVerbose(string(out))
 				if err != nil {

internal/commands/scan.go

@@ -35,6 +35,7 @@ import (
 	"github.com/pkg/errors"
 
 	"github.com/MakeNowJust/heredoc"
+	"github.com/checkmarx/ast-cli/internal/kicsshutdown"
 	commonParams "github.com/checkmarx/ast-cli/internal/params"
 	"github.com/checkmarx/ast-cli/internal/wrappers"
 	"github.com/mssola/user_agent"
@@ -279,7 +280,9 @@ func NewScanCommand(
 
 func scanRealtimeSubCommand() *cobra.Command {
 	kicsContainerID := uuid.New()
-	viper.Set(commonParams.KicsContainerNameKey, kicsContainerPrefixName+kicsContainerID.String())
+	kicsName := kicsContainerPrefixName + kicsContainerID.String()
+	viper.Set(commonParams.KicsContainerNameKey, kicsName)
+	kicsshutdown.SetKicsContainerName(kicsName)
 	realtimeScanCmd := &cobra.Command{
 		Use:   "kics-realtime",
 		Short: "Create and run kics scan",

internal/kicsshutdown/container_name.go

@@ -0,0 +1,25 @@
+// Package kicsshutdown holds the current KICS Docker container name for SIGTERM cleanup.
+// It is updated alongside viper wherever KicsContainerNameKey is set so the signal handler
+// can read the latest name without concurrent access to viper.
+package kicsshutdown
+
+import "sync"
+
+var (
+	mu   sync.RWMutex
+	name string
+)
+
+// SetKicsContainerName records the container name used for shutdown handling.
+func SetKicsContainerName(n string) {
+	mu.Lock()
+	defer mu.Unlock()
+	name = n
+}
+
+// GetKicsContainerName returns the last recorded container name.
+func GetKicsContainerName() string {
+	mu.RLock()
+	defer mu.RUnlock()
+	return name
+}

internal/services/realtimeengine/iacrealtime/container-manager.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/checkmarx/ast-cli/internal/commands/util"
+	"github.com/checkmarx/ast-cli/internal/kicsshutdown"
 	"github.com/checkmarx/ast-cli/internal/logger"
 	commonParams "github.com/checkmarx/ast-cli/internal/params"
 	"github.com/google/uuid"
@@ -32,6 +33,7 @@ func (dm *ContainerManager) GenerateContainerID() string {
 	containerID := uuid.New().String()
 	containerName := KicsContainerPrefix + containerID
 	viper.Set(commonParams.KicsContainerNameKey, containerName)
+	kicsshutdown.SetKicsContainerName(containerName)
 	return containerName
 }
 

internal/services/realtimeengine/iacrealtime/container-manager_test.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/checkmarx/ast-cli/internal/kicsshutdown"
 	commonParams "github.com/checkmarx/ast-cli/internal/params"
 	"github.com/google/uuid"
 	"github.com/spf13/viper"
@@ -46,6 +47,7 @@ func (m *MockContainerManager) GenerateContainerID() string {
 	containerName := KicsContainerPrefix + containerID
 	m.GeneratedContainerIDs = append(m.GeneratedContainerIDs, containerName)
 	viper.Set(commonParams.KicsContainerNameKey, containerName)
+	kicsshutdown.SetKicsContainerName(containerName)
 	return containerName
 }
 
@@ -102,6 +104,7 @@ func TestMockContainerManager_GenerateContainerID(t *testing.T) {
 
 	// Clear any existing value
 	viper.Set(commonParams.KicsContainerNameKey, "")
+	kicsshutdown.SetKicsContainerName("")
 
 	containerName := dm.GenerateContainerID()
 
@@ -125,6 +128,9 @@ func TestMockContainerManager_GenerateContainerID(t *testing.T) {
 	if viperValue != containerName {
 		t.Errorf("Viper should be set to '%s', got '%s'", containerName, viperValue)
 	}
+	if kicsshutdown.GetKicsContainerName() != containerName {
+		t.Errorf("kicsshutdown should be set to '%s', got '%s'", containerName, kicsshutdown.GetKicsContainerName())
+	}
 
 	// Test that mock recorded the generated ID
 	if len(dm.GeneratedContainerIDs) != 1 {
@@ -164,6 +170,7 @@ func TestMockContainerManager_RunKicsContainer(t *testing.T) {
 	// Set up test parameters
 	containerName := "test-container"
 	viper.Set(commonParams.KicsContainerNameKey, containerName)
+	kicsshutdown.SetKicsContainerName(containerName)
 
 	tests := []struct {
 		name      string
@@ -263,6 +270,9 @@ func TestMockContainerManager_Integration(t *testing.T) {
 	if viper.GetString(commonParams.KicsContainerNameKey) != containerName {
 		t.Error("Container name should be set in viper after generation")
 	}
+	if kicsshutdown.GetKicsContainerName() != containerName {
+		t.Error("Container name should be set in kicsshutdown after generation")
+	}
 
 	// Test running container
 	err := dm.RunKicsContainer("docker", "/tmp:/path")
@@ -379,6 +389,7 @@ func TestContainerManager_GenerateContainerID(t *testing.T) {
 
 	// Clear any existing value
 	viper.Set(commonParams.KicsContainerNameKey, "")
+	kicsshutdown.SetKicsContainerName("")
 
 	containerName := cm.GenerateContainerID()
 
@@ -402,6 +413,9 @@ func TestContainerManager_GenerateContainerID(t *testing.T) {
 	if viperValue != containerName {
 		t.Errorf("Viper should be set to '%s', got '%s'", containerName, viperValue)
 	}
+	if kicsshutdown.GetKicsContainerName() != containerName {
+		t.Errorf("kicsshutdown should be set to '%s', got '%s'", containerName, kicsshutdown.GetKicsContainerName())
+	}
 
 	// Test that subsequent calls generate different IDs
 	containerName2 := cm.GenerateContainerID()

internal/wrappers/jwt-helper.go

@@ -203,12 +203,19 @@ func GetUniqueID() string {
 	var uniqueID string
 	// Check License first
 	jwtWrapper := NewJwtWrapper()
-	isAllowed, err := jwtWrapper.IsAllowedEngine(commonParams.CheckmarxDevAssistType)
+	devAssistAllowed, err := jwtWrapper.IsAllowedEngine(commonParams.CheckmarxDevAssistType)
 	if err != nil {
 		logger.PrintIfVerbose("Failed to check engine allowance: " + err.Error())
 		return ""
 	}
-	if !isAllowed {
+
+	oneAssistAllowed, err := jwtWrapper.IsAllowedEngine(commonParams.CheckmarxOneAssistType)
+	if err != nil {
+		logger.PrintIfVerbose("Failed to check engine allowance: " + err.Error())
+		return ""
+	}
+
+	if !devAssistAllowed && !oneAssistAllowed {
 		return ""
 	}
 

internal/wrappers/jwt-helper_test.go

@@ -90,7 +90,7 @@ func TestGetUniqueID(t *testing.T) {
 		if result != "" {
 			assert.Equal(t, existingID, result)
 		} else {
-			t.Skip("Requires valid auth and 'Checkmarx Developer Assist' license")
+			t.Skip("Requires valid auth and 'Checkmarx Developer Assist' or 'Checkmarx OneAssist' license")
 		}
 	})
 
@@ -101,7 +101,7 @@ func TestGetUniqueID(t *testing.T) {
 		result := GetUniqueID()
 
 		if result == "" {
-			t.Skip("Requires valid auth and 'Checkmarx Developer Assist' license")
+			t.Skip("Requires valid auth and 'Checkmarx Developer Assist' or 'Checkmarx OneAssist' license")
 			return
 		}