feat(config,main): add option to configure pprof

Co-authored-by: PrimalPimmy <Prashant20.pm@gmail.com>
Signed-off-by: PrimalPimmy <Prashant20.pm@gmail.com>
Signed-off-by: daemon1024 <barun1024@gmail.com>

Author: daemon1024

KubeArmor/config/config.go

@@ -31,10 +31,12 @@ type KubearmorConfig struct {
 	Visibility     string // Container visibility to use
 	HostVisibility string // Host visibility to use
 
-	Policy     bool // Enable/Disable policy enforcement
-	HostPolicy bool // Enable/Disable host policy enforcement
-	KVMAgent   bool // Enable/Disable KVM Agent
-	K8sEnv     bool // Is k8s env ?
+	Policy      bool   // Enable/Disable policy enforcement
+	EnablePPROF bool   // Enable pprof to be used
+	PPROFPort   string // pprof Port to use
+	HostPolicy  bool   // Enable/Disable host policy enforcement
+	KVMAgent    bool   // Enable/Disable KVM Agent
+	K8sEnv      bool   // Is k8s env ?
 
 	Debug bool // Enable/Disable KubeArmor debug mode
 
@@ -70,6 +72,8 @@ const (
 	PIDFilePath                          string = "/opt/kubearmor/kubearmor.pid"
 	ConfigCluster                        string = "cluster"
 	ConfigHost                           string = "host"
+	ConfigEnablePPROF                    string = "pprofEnable"
+	ConfigPPROFPort                      string = "pprofPort"
 	ConfigGRPC                           string = "gRPC"
 	ConfigTLSCertPath                    string = "tlsCertPath"
 	ConfigTLSCertProvider                string = "tlsCertProvider"
@@ -122,6 +126,8 @@ func readCmdLineParams() {
 	hostVisStr := flag.String(ConfigHostVisibility, "default", "Host Visibility to use [process,file,network,capabilities,none] (default \"none\" for k8s, \"process,file,network,capabilities\" for VM)")
 
 	policyB := flag.Bool(ConfigKubearmorPolicy, true, "enabling KubeArmorPolicy")
+	pprofEnableB := flag.Bool(ConfigEnablePPROF, false, "enables pprof to be used")
+	pprofPortStr := flag.String(ConfigPPROFPort, "8080", "PPROF port number")
 	hostPolicyB := flag.Bool(ConfigKubearmorHostPolicy, false, "enabling KubeArmorHostPolicy")
 	kvmAgentB := flag.Bool(ConfigKubearmorVM, false, "enabling KubeArmorVM")
 	k8sEnvB := flag.Bool(ConfigK8sEnv, true, "is k8s env?")
@@ -180,6 +186,8 @@ func readCmdLineParams() {
 	viper.SetDefault(ConfigVisibility, *visStr)
 	viper.SetDefault(ConfigHostVisibility, *hostVisStr)
 
+	viper.SetDefault(ConfigEnablePPROF, *pprofEnableB)
+	viper.SetDefault(ConfigPPROFPort, *pprofPortStr)
 	viper.SetDefault(ConfigKubearmorPolicy, *policyB)
 	viper.SetDefault(ConfigKubearmorHostPolicy, *hostPolicyB)
 	viper.SetDefault(ConfigKubearmorVM, *kvmAgentB)
@@ -263,6 +271,8 @@ func LoadConfig() error {
 	GlobalCfg.Visibility = viper.GetString(ConfigVisibility)
 	GlobalCfg.HostVisibility = viper.GetString(ConfigHostVisibility)
 
+	GlobalCfg.EnablePPROF = viper.GetBool(ConfigEnablePPROF)
+	GlobalCfg.PPROFPort = viper.GetString(ConfigPPROFPort)
 	GlobalCfg.Policy = viper.GetBool(ConfigKubearmorPolicy)
 	GlobalCfg.HostPolicy = viper.GetBool(ConfigKubearmorHostPolicy)
 	GlobalCfg.KVMAgent = viper.GetBool(ConfigKubearmorVM)

KubeArmor/main.go

@@ -5,8 +5,11 @@
 package main
 
 import (
+	"net/http"
+	"net/http/pprof"
 	"os"
 	"path/filepath"
+	"time"
 
 	cfg "github.com/kubearmor/KubeArmor/KubeArmor/config"
 	"github.com/kubearmor/KubeArmor/KubeArmor/core"
@@ -76,5 +79,27 @@ func main() {
 		return
 	}
 
+	if cfg.GlobalCfg.EnablePPROF {
+		go func() {
+			mux := http.NewServeMux()
+
+			mux.HandleFunc("/debug/pprof/", pprof.Index)
+			mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+			mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
+			mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+
+			server := &http.Server{
+				Addr:              "0.0.0.0:" + cfg.GlobalCfg.PPROFPort,
+				Handler:           mux,
+				ReadHeaderTimeout: 10 * time.Second,
+			}
+
+			kg.Printf("pprof server is running on http://0.0.0.0:" + cfg.GlobalCfg.PPROFPort)
+			if err := server.ListenAndServe(); err != nil {
+				kg.Errf("Could not expose a pprof server due to %s", err.Error())
+			}
+		}()
+	}
+
 	core.KubeArmor()
 }