From 435ba547436dc8b1b4684ff52bbadd2e528bb4be Mon Sep 17 00:00:00 2001
From: ionutboangiu <ionut.boangiu@itsyscom.com>
Date: Fri, 9 Aug 2024 20:24:43 +0300
Subject: [PATCH] Register prometheus and pprof endpoints using https mux

---
 cmd/cgr-engine/cgr-engine.go |  2 ++
 cores/server.go              | 36 +++++++++++++++++++++++++++++++++++-
 cores/server_it_test.go      | 10 ++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/cmd/cgr-engine/cgr-engine.go b/cmd/cgr-engine/cgr-engine.go
index fd4f5759d..f8af096b4 100644
--- a/cmd/cgr-engine/cgr-engine.go
+++ b/cmd/cgr-engine/cgr-engine.go
@@ -253,6 +253,8 @@ func startRPC(server *cores.Server, internalRaterChan,
 			cfg.TLSCfg().ServerName,
 			cfg.HTTPCfg().HTTPJsonRPCURL,
 			cfg.HTTPCfg().HTTPWSURL,
+			cfg.HTTPCfg().PrometheusURL,
+			cfg.HTTPCfg().PprofPath,
 			cfg.HTTPCfg().HTTPUseBasicAuth,
 			cfg.HTTPCfg().HTTPAuthUsers,
 			shdChan,
diff --git a/cores/server.go b/cores/server.go
index 50f56e340..5c4006676 100644
--- a/cores/server.go
+++ b/cores/server.go
@@ -465,7 +465,7 @@ func (s *Server) handleWebSocket(ws *websocket.Conn) {
 }
 
 func (s *Server) ServeHTTPTLS(addr, serverCrt, serverKey, caCert string, serverPolicy int,
-	serverName string, jsonRPCURL string, wsRPCURL string,
+	serverName, jsonRPCURL, wsRPCURL, promURL, pprofPath string,
 	useBasicAuth bool, userList map[string]string, shdChan *utils.SyncedChan) {
 	s.RLock()
 	enabled := s.rpcEnabled
@@ -496,6 +496,40 @@ func (s *Server) ServeHTTPTLS(addr, serverCrt, serverKey, caCert string, serverP
 			s.httpsMux.Handle(wsRPCURL, wsHandler)
 		}
 	}
+	if promURL != "" {
+		s.Lock()
+		s.httpEnabled = true
+		s.Unlock()
+		utils.Logger.Info(fmt.Sprintf("<HTTPS> prometheus metrics endpoint registered at %q", promURL))
+		promHandler := promhttp.Handler()
+		if useBasicAuth {
+			s.httpsMux.HandleFunc(promURL, use(promHandler.ServeHTTP, basicAuth(userList)))
+		} else {
+			s.httpsMux.Handle(promURL, promHandler)
+		}
+	}
+	if pprofPath != "" {
+		s.Lock()
+		s.httpEnabled = true
+		s.Unlock()
+		if !strings.HasSuffix(pprofPath, "/") {
+			pprofPath += "/"
+		}
+		utils.Logger.Info(fmt.Sprintf("<HTTPS> profiling endpoints registered at %q", pprofPath))
+		if useBasicAuth {
+			s.httpsMux.HandleFunc(pprofPath, use(pprof.Index, basicAuth(userList)))
+			s.httpsMux.HandleFunc(pprofPath+"cmdline", use(pprof.Cmdline, basicAuth(userList)))
+			s.httpsMux.HandleFunc(pprofPath+"profile", use(pprof.Profile, basicAuth(userList)))
+			s.httpsMux.HandleFunc(pprofPath+"symbol", use(pprof.Symbol, basicAuth(userList)))
+			s.httpsMux.HandleFunc(pprofPath+"trace", use(pprof.Trace, basicAuth(userList)))
+		} else {
+			s.httpsMux.HandleFunc(pprofPath, pprof.Index)
+			s.httpsMux.HandleFunc(pprofPath+"cmdline", pprof.Cmdline)
+			s.httpsMux.HandleFunc(pprofPath+"profile", pprof.Profile)
+			s.httpsMux.HandleFunc(pprofPath+"symbol", pprof.Symbol)
+			s.httpsMux.HandleFunc(pprofPath+"trace", pprof.Trace)
+		}
+	}
 	if !s.httpEnabled {
 		return
 	}
diff --git a/cores/server_it_test.go b/cores/server_it_test.go
index 933d10c2b..f9278b735 100644
--- a/cores/server_it_test.go
+++ b/cores/server_it_test.go
@@ -575,6 +575,8 @@ func testServeHTTPTLS(t *testing.T) {
 		cfg.TLSCfg().ServerName,
 		cfg.HTTPCfg().HTTPJsonRPCURL,
 		cfg.HTTPCfg().HTTPWSURL,
+		cfg.HTTPCfg().PrometheusURL,
+		cfg.HTTPCfg().PprofPath,
 		cfg.HTTPCfg().HTTPUseBasicAuth,
 		cfg.HTTPCfg().HTTPAuthUsers,
 		shdChan)
@@ -590,6 +592,8 @@ func testServeHTTPTLS(t *testing.T) {
 		cfg.TLSCfg().ServerName,
 		cfg.HTTPCfg().HTTPJsonRPCURL,
 		cfg.HTTPCfg().HTTPWSURL,
+		cfg.HTTPCfg().PrometheusURL,
+		cfg.HTTPCfg().PprofPath,
 		cfg.HTTPCfg().HTTPUseBasicAuth,
 		cfg.HTTPCfg().HTTPAuthUsers,
 		shdChan)
@@ -620,6 +624,8 @@ func testServeHTTPTLSWithBasicAuth(t *testing.T) {
 		cfg.TLSCfg().ServerName,
 		cfg.HTTPCfg().HTTPJsonRPCURL,
 		cfg.HTTPCfg().HTTPWSURL,
+		cfg.HTTPCfg().PrometheusURL,
+		cfg.HTTPCfg().PprofPath,
 		!cfg.HTTPCfg().HTTPUseBasicAuth,
 		cfg.HTTPCfg().HTTPAuthUsers,
 		shdChan)
@@ -649,6 +655,8 @@ func testServeHTTPTLSError(t *testing.T) {
 		cfg.TLSCfg().ServerName,
 		cfg.HTTPCfg().HTTPJsonRPCURL,
 		cfg.HTTPCfg().HTTPWSURL,
+		cfg.HTTPCfg().PrometheusURL,
+		cfg.HTTPCfg().PprofPath,
 		!cfg.HTTPCfg().HTTPUseBasicAuth,
 		cfg.HTTPCfg().HTTPAuthUsers,
 		shdChan)
@@ -678,6 +686,8 @@ func testServeHTTPTLSHttpNotEnabled(t *testing.T) {
 		cfg.TLSCfg().ServerName,
 		utils.EmptyString,
 		utils.EmptyString,
+		utils.EmptyString,
+		utils.EmptyString,
 		cfg.HTTPCfg().HTTPUseBasicAuth,
 		cfg.HTTPCfg().HTTPAuthUsers,
 		shdChan)