From cd4fdcd9bbabdac340b1ca150bdce01694eba516 Mon Sep 17 00:00:00 2001
From: Ankur Kothiwal <ankur.kothiwal99@gmail.com>
Date: Thu, 6 Apr 2023 14:10:06 +0530
Subject: [PATCH] fix logic for missing kubearmor-relay deployment

Signed-off-by: Ankur Kothiwal <ankur.kothiwal99@gmail.com>
---
 src/cluster/k8sClientHandler.go  |  3 ++-
 src/plugin/kubearmor.go          | 14 ++++++++++----
 src/systempolicy/systemPolicy.go | 11 +++++++++++
 src/types/constants.go           |  3 +++
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/src/cluster/k8sClientHandler.go b/src/cluster/k8sClientHandler.go
index 84adcb55..ac71b7a4 100644
--- a/src/cluster/k8sClientHandler.go
+++ b/src/cluster/k8sClientHandler.go
@@ -491,7 +491,8 @@ func GetKubearmorRelayURL() string {
 		log.Error().Msg(err.Error())
 		return ""
 	}
-	if pods == nil {
+	if pods == nil || len(pods.Items) == 0 {
+		log.Error().Msgf("Unable to find kubearmor-relay")
 		return ""
 	}
 	namespace = pods.Items[0].Namespace
diff --git a/src/plugin/kubearmor.go b/src/plugin/kubearmor.go
index 28b4e89b..3bb76f73 100644
--- a/src/plugin/kubearmor.go
+++ b/src/plugin/kubearmor.go
@@ -19,6 +19,7 @@ import (
 	"github.com/accuknox/auto-policy-discovery/src/types"
 	pb "github.com/kubearmor/KubeArmor/protobuf"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
 )
 
 // Global Variable
@@ -313,13 +314,18 @@ func ConvertKubeArmorLogToKnoxSystemLog(relayLog *pb.Alert) (types.KnoxSystemLog
 func ConnectKubeArmorRelay(cfg types.ConfigKubeArmorRelay) *grpc.ClientConn {
 	addr := net.JoinHostPort(cfg.KubeArmorRelayURL, cfg.KubeArmorRelayPort)
 
-	conn, err := grpc.Dial(addr, grpc.WithInsecure())
+	// Check for kubearmor-relay with 30s timeout
+	ctx, cf1 := context.WithTimeout(context.Background(), time.Second*30)
+	defer cf1()
+
+	// Blocking grpc Dial: in case of a bad connection, fails with timeout
+	conn, err := grpc.DialContext(ctx, addr, grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithBlock())
 	if err != nil {
-		log.Error().Msg("err connecting kubearmor relay. " + err.Error())
+		log.Error().Msg("Error connecting kubearmor relay: " + err.Error())
 		return nil
 	}
 
-	log.Info().Msg("connected to kubearmor relay " + addr)
+	log.Info().Msg("Connected to kubearmor relay " + addr)
 	return conn
 }
 
@@ -327,7 +333,7 @@ func GetSystemAlertsFromKubeArmorRelay(trigger int) []*pb.Alert {
 	results := []*pb.Alert{}
 	KubeArmorRelayLogsMutex.Lock()
 	if len(KubeArmorRelayLogs) == 0 {
-		log.Info().Msgf("KubeArmor Relay traffic flow not exist")
+		log.Info().Msgf("KubeArmor Relay traffic flow does not exist")
 		KubeArmorRelayLogsMutex.Unlock()
 		return results
 	}
diff --git a/src/systempolicy/systemPolicy.go b/src/systempolicy/systemPolicy.go
index 4d811583..cafa3efb 100644
--- a/src/systempolicy/systemPolicy.go
+++ b/src/systempolicy/systemPolicy.go
@@ -1464,6 +1464,17 @@ func StartSystemLogRcvr() {
 	for {
 		if cfg.GetCfgSystemLogFrom() == "kubearmor" {
 			url := cluster.GetKubearmorRelayURL()
+			if url == "" {
+				log.Error().Msg("kubearmor-relay url not found, retrying...")
+				for i := 0; i < types.Maxtries; i++ {
+					time.Sleep(10 * time.Second)
+					url = cluster.GetKubearmorRelayURL()
+					if url != "" {
+						break
+					}
+				}
+				return
+			}
 			plugin.StartKubeArmorRelay(SystemStopChan, types.ConfigKubeArmorRelay{
 				KubeArmorRelayURL:  url,
 				KubeArmorRelayPort: cfg.CurrentCfg.ConfigKubeArmorRelay.KubeArmorRelayPort,
diff --git a/src/types/constants.go b/src/types/constants.go
index 907c23cf..c56b2f7f 100644
--- a/src/types/constants.go
+++ b/src/types/constants.go
@@ -44,4 +44,7 @@ const (
 	// K8sNetworkPolicy
 	K8sNwPolicyAPIVersion = "networking.k8s.io/v1"
 	K8sNwPolicyKind       = "NetworkPolicy"
+
+	// max no. of tries to connect to kubearmor-relay
+	Maxtries = 6
 )