From 17a330ddfd5c83b6c88bee4464ce3dc10dfbd357 Mon Sep 17 00:00:00 2001
From: Prateek <prateeknandle@gmail.com>
Date: Thu, 11 Jul 2024 18:20:17 +0530
Subject: [PATCH] fix(core):timeout when host & cluster security policies crds
 are not found

Signed-off-by: Prateek <prateeknandle@gmail.com>
---
 KubeArmor/core/kubeArmor.go  | 15 ++++++-------
 KubeArmor/core/kubeUpdate.go | 41 +++++++++++++++++++++++++++++-------
 2 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/KubeArmor/core/kubeArmor.go b/KubeArmor/core/kubeArmor.go
index 261d86a1d2..9b1ba26e0e 100644
--- a/KubeArmor/core/kubeArmor.go
+++ b/KubeArmor/core/kubeArmor.go
@@ -708,8 +708,8 @@ func KubeArmor() {
 
 	// == //
 
+	timeout, err := time.ParseDuration(cfg.GlobalCfg.InitTimeout)
 	if dm.K8sEnabled && cfg.GlobalCfg.Policy {
-		timeout, err := time.ParseDuration(cfg.GlobalCfg.InitTimeout)
 		if err != nil {
 			dm.Logger.Warnf("Not a valid InitTimeout duration: %q, defaulting to '60s'", cfg.GlobalCfg.InitTimeout)
 			timeout = 60 * time.Second
@@ -726,14 +726,12 @@ func KubeArmor() {
 		dm.Logger.Print("Started to monitor security policies")
 
 		// watch cluster security policies
-		clusterSecurityPoliciesSynced := dm.WatchClusterSecurityPolicies()
+		clusterSecurityPoliciesSynced := dm.WatchClusterSecurityPolicies(timeout)
 		if clusterSecurityPoliciesSynced == nil {
-			// destroy the daemon
-			dm.DestroyKubeArmorDaemon()
-
-			return
+			dm.Logger.Warn("error while monitoring cluster security policies, informer cache not synced")
+		} else {
+			dm.Logger.Print("Started to monitor cluster security policies")
 		}
-		dm.Logger.Print("Started to monitor cluster security policies")
 
 		// watch default posture
 		defaultPostureSynced := dm.WatchDefaultPosture()
@@ -776,8 +774,7 @@ func KubeArmor() {
 
 	if dm.K8sEnabled && cfg.GlobalCfg.HostPolicy {
 		// watch host security policies
-		go dm.WatchHostSecurityPolicies()
-		dm.Logger.Print("Started to monitor host security policies")
+		go dm.WatchHostSecurityPolicies(timeout)
 	}
 
 	if !dm.K8sEnabled && (enableContainerPolicy || cfg.GlobalCfg.HostPolicy) {
diff --git a/KubeArmor/core/kubeUpdate.go b/KubeArmor/core/kubeUpdate.go
index 509ffd39e1..f3d345d1a3 100644
--- a/KubeArmor/core/kubeUpdate.go
+++ b/KubeArmor/core/kubeUpdate.go
@@ -1676,13 +1676,22 @@ func (dm *KubeArmorDaemon) WatchSecurityPolicies() cache.InformerSynced {
 }
 
 // WatchClusterSecurityPolicies Function
-func (dm *KubeArmorDaemon) WatchClusterSecurityPolicies() cache.InformerSynced {
-	for {
-		if !K8s.CheckCustomResourceDefinition("kubearmorclusterpolicies") {
-			time.Sleep(time.Second * 1)
-			continue
-		} else {
-			break
+func (dm *KubeArmorDaemon) WatchClusterSecurityPolicies(timeout time.Duration) cache.InformerSynced {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	crdFound := false
+	for !crdFound {
+		select {
+		case <-ctx.Done():
+			dm.Logger.Warn("timeout while monitoring cluster security policies, kubearmorclusterpolicies CRD not found")
+			return nil
+		default:
+			if K8s.CheckCustomResourceDefinition("kubearmorclusterpolicies") {
+				crdFound = true
+			} else {
+				time.Sleep(time.Second * 1)
+			}
 		}
 	}
 
@@ -2260,8 +2269,24 @@ func (dm *KubeArmorDaemon) ParseAndUpdateHostSecurityPolicy(event tp.K8sKubeArmo
 }
 
 // WatchHostSecurityPolicies Function
-func (dm *KubeArmorDaemon) WatchHostSecurityPolicies() {
+func (dm *KubeArmorDaemon) WatchHostSecurityPolicies(timeout time.Duration) {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
 	for {
+		select {
+		case <-ctx.Done():
+			dm.Logger.Warn("timeout while monitoring host security policies, kubearmorhostpolicies CRD not found")
+			return
+		default:
+			if !K8s.CheckCustomResourceDefinition("kubearmorhostpolicies") {
+				time.Sleep(time.Second * 1)
+				continue
+			}
+		}
+
+		dm.Logger.Print("Started to monitor host security policies")
+
 		if !K8s.CheckCustomResourceDefinition("kubearmorhostpolicies") {
 			time.Sleep(time.Second * 1)
 			continue