From 72f6ecee3eeaa722cf0cc245a6186705f2bfb715 Mon Sep 17 00:00:00 2001
From: stefin9898 <112696367+stefin9898@users.noreply.github.com>
Date: Thu, 18 May 2023 09:12:46 +0530
Subject: [PATCH] Added routes in http handle for pprof (#728)

---
 src/main.go | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/src/main.go b/src/main.go
index 6efbad3b..8dcc803e 100644
--- a/src/main.go
+++ b/src/main.go
@@ -3,12 +3,11 @@ package main
 import (
 	"math/rand"
 	"net"
+	"net/http"
+	"net/http/pprof"
 	"os"
 	"time"
 
-	"net/http"
-	_ "net/http/pprof"
-
 	"github.com/accuknox/auto-policy-discovery/src/cluster"
 	"github.com/accuknox/auto-policy-discovery/src/config"
 	"github.com/accuknox/auto-policy-discovery/src/libs"
@@ -38,14 +37,16 @@ func init() {
 	log = logger.GetInstance()
 
 	//Get pprof flag
-	pprof := viper.GetBool("pprof")
-	if pprof {
+	if viper.GetBool("pprof") {
 		// Server for pprof
 		go func() {
 			log.Info().Msgf("pprof enabled on :6060\n")
-			r := http.NewServeMux()
-			r.Handle("/debug/", http.DefaultServeMux)
-			http.ListenAndServe(":6060", r)
+			http.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
+			http.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
+			http.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
+			http.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
+			http.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
+			http.ListenAndServe(":6060", nil)
 		}()
 	}