55 "fmt"
66 "io"
77 "net/http"
8- "regexp "
8+ "slices "
99 "time"
1010
1111 "github.com/gin-gonic/gin"
@@ -21,7 +21,7 @@ import (
2121// @Failure 500 {object} map[string]interface{}
2222// @Router /pprof/evaluator_upload/{param} [get]
2323func GetEvaluatorUploadPprof (c * gin.Context ) {
24- pprofHandler (c , utils . CoreCfg . EvaluatorUploadPrivateAddress )
24+ pprofHandler (c , "evaluator-upload" )
2525}
2626
2727// @Summary Get profile info
@@ -34,7 +34,7 @@ func GetEvaluatorUploadPprof(c *gin.Context) {
3434// @Failure 500 {object} map[string]interface{}
3535// @Router /pprof/evaluator_recalc/{param} [get]
3636func GetEvaluatorRecalcPprof (c * gin.Context ) {
37- pprofHandler (c , utils . CoreCfg . EvaluatorRecalcPrivateAddress )
37+ pprofHandler (c , "evaluator-recalc" )
3838}
3939
4040// @Summary Get profile info
@@ -47,7 +47,7 @@ func GetEvaluatorRecalcPprof(c *gin.Context) {
4747// @Failure 500 {object} map[string]interface{}
4848// @Router /pprof/listener/{param} [get]
4949func GetListenerPprof (c * gin.Context ) {
50- pprofHandler (c , utils . CoreCfg . ListenerPrivateAddress )
50+ pprofHandler (c , "listener" )
5151}
5252
5353// @Summary Get profile info
@@ -60,20 +60,19 @@ func GetListenerPprof(c *gin.Context) {
6060// @Failure 500 {object} map[string]interface{}
6161// @Router /pprof/manager/{param} [get]
6262func GetManagerPprof (c * gin.Context ) {
63- pprofHandler (c , utils . CoreCfg . ManagerPrivateAddress )
63+ pprofHandler (c , "manager" )
6464}
6565
66- var paramRegexp = regexp . MustCompile ( "^( heap| profile| block| mutex| trace)$" )
66+ var allowedParams = [] string { " heap" , " profile" , " block" , " mutex" , " trace" }
6767
68- func pprofHandler (c * gin.Context , address string ) {
68+ func pprofHandler (c * gin.Context , serviceName string ) {
6969 query := c .Request .URL .RawQuery
7070 param := c .Param ("param" )
71- match := paramRegexp .FindStringSubmatch (param )
72- if len (match ) < 1 {
71+ if ! slices .Contains (allowedParams , param ) {
7372 c .Status (http .StatusBadRequest )
7473 return
7574 }
76- data , err := getPprof (address , match [ 0 ] , query )
75+ data , err := getPprof (serviceName , param , query )
7776 if err != nil {
7877 c .JSON (http .StatusInternalServerError , gin.H {"err" : err .Error ()})
7978 return
@@ -82,19 +81,32 @@ func pprofHandler(c *gin.Context, address string) {
8281 c .Data (http .StatusOK , "application/octet-stream" , data )
8382}
8483
85- func getPprof (address , param , query string ) ([]byte , error ) {
84+ func getPprof (serviceName , param , query string ) ([]byte , error ) {
8685 client := & http.Client {
8786 Timeout : time .Second * 60 ,
8887 }
8988 if len (query ) > 0 {
9089 param = param + "?" + query
9190 }
91+ var address string
92+ switch serviceName {
93+ case "manager" :
94+ address = utils .CoreCfg .ManagerPrivateAddress
95+ case "listener" :
96+ address = utils .CoreCfg .ListenerPrivateAddress
97+ case "evaluator-upload" :
98+ address = utils .CoreCfg .EvaluatorUploadPrivateAddress
99+ case "evaluator-recalc" :
100+ address = utils .CoreCfg .EvaluatorRecalcPrivateAddress
101+ default :
102+ return nil , fmt .Errorf ("invalid service name: %s" , serviceName )
103+ }
92104 urlPath := fmt .Sprintf ("%s/debug/pprof/%s" , address , param )
93- req , err := http .NewRequest (http .MethodGet , urlPath , nil )
105+ req , err := http .NewRequest (http .MethodGet , urlPath , nil ) // #nosec G704
94106 if err != nil {
95107 return nil , err
96108 }
97- res , err := client .Do (req )
109+ res , err := client .Do (req ) // #nosec G704
98110 if err != nil {
99111 return nil , err
100112 }
0 commit comments