Skip to content

Commit

Permalink
add a new function RouteRegister which supports *gin.RouterGroup… (#20)
Browse files Browse the repository at this point in the history
first parameter
  • Loading branch information
micanzhang authored Apr 23, 2020
1 parent 021a0ad commit 3f09087
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 2 deletions.
9 changes: 8 additions & 1 deletion pprof.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,16 @@ func getPrefix(prefixOptions ...string) string {
// the provided gin.Engine. prefixOptions is a optional. If not prefixOptions,
// the default path prefix is used, otherwise first prefixOptions will be path prefix.
func Register(r *gin.Engine, prefixOptions ...string) {
RouteRegister(&(r.RouterGroup), prefixOptions...)
}

// RouteRegister the standard HandlerFuncs from the net/http/pprof package with
// the provided gin.GrouterGroup. prefixOptions is a optional. If not prefixOptions,
// the default path prefix is used, otherwise first prefixOptions will be path prefix.
func RouteRegister(rg *gin.RouterGroup, prefixOptions ...string) {
prefix := getPrefix(prefixOptions...)

prefixRouter := r.Group(prefix)
prefixRouter := rg.Group(prefix)
{
prefixRouter.GET("/", pprofHandler(pprof.Index))
prefixRouter.GET("/cmdline", pprofHandler(pprof.Cmdline))
Expand Down
48 changes: 47 additions & 1 deletion pprof_test.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,12 @@
package pprof

import "testing"
import (
"net/http"
"net/http/httptest"
"testing"

"github.com/gin-gonic/gin"
)

func Test_getPrefix(t *testing.T) {
tests := []struct {
Expand All @@ -18,3 +24,43 @@ func Test_getPrefix(t *testing.T) {
}
}
}

func TestRegisterAndRouteRegister(t *testing.T) {
bearerToken := "Bearer token"
gin.SetMode(gin.ReleaseMode)
r := gin.New()
Register(r)
adminGroup := r.Group("/admin", func(c *gin.Context) {
if c.Request.Header.Get("Authorization") != bearerToken {
c.AbortWithStatus(http.StatusForbidden)
return
}
c.Next()
})
RouteRegister(adminGroup, "pprof")

req, _ := http.NewRequest(http.MethodGet, "/debug/pprof/", nil)
rw := httptest.NewRecorder()
r.ServeHTTP(rw, req)

if expected, got := http.StatusOK, rw.Code; expected != got {
t.Errorf("expected: %d, got: %d", expected, got)
}

req, _ = http.NewRequest(http.MethodGet, "/admin/pprof/", nil)
rw = httptest.NewRecorder()
r.ServeHTTP(rw, req)

if expected, got := http.StatusForbidden, rw.Code; expected != got {
t.Errorf("expected: %d, got: %d", expected, got)
}

req, _ = http.NewRequest(http.MethodGet, "/admin/pprof/", nil)
req.Header.Set("Authorization", bearerToken)
rw = httptest.NewRecorder()
r.ServeHTTP(rw, req)

if expected, got := http.StatusOK, rw.Code; expected != got {
t.Errorf("expected: %d, got: %d", expected, got)
}
}

0 comments on commit 3f09087

Please sign in to comment.