added ability to filter ExpvarHandler output with regexp

2 files changed, 64 insertions, 6 deletions

diff --git a/fasthttputil/expvar.go b/fasthttputil/expvar.go
index b2679a2..a942f4c 100644
--- a/fasthttputil/expvar.go
+++ b/fasthttputil/expvar.go
@@ -3,26 +3,58 @@ package fasthttputil
 import (
 	"expvar"
 	"fmt"
+	"regexp"
 
 	"github.com/valyala/fasthttp"
 )
 
+var (
+	expvarHandlerCalls = expvar.NewInt("expvarHandlerCalls")
+	expvarRegexpErrors = expvar.NewInt("expvarRegexpErrors")
+)
+
 // ExpvarHandler dumps json representation of expvars to http response.
 //
+// Expvars may be filtered by regexp provided via 'r' query argument.
+//
 // See https://golang.org/pkg/expvar/ for details.
 func ExpvarHandler(ctx *fasthttp.RequestCtx) {
+	expvarHandlerCalls.Add(1)
+
 	ctx.Response.Reset()
 
+	r, err := getExpvarRegexp(ctx)
+	if err != nil {
+		expvarRegexpErrors.Add(1)
+		fmt.Fprintf(ctx, "Error when obtaining expvar regexp: %s", err)
+		ctx.SetStatusCode(fasthttp.StatusBadRequest)
+		return
+	}
+
 	fmt.Fprintf(ctx, "{\n")
 	first := true
 	expvar.Do(func(kv expvar.KeyValue) {
 		if !first {
 			fmt.Fprintf(ctx, ",\n")
 		}
-		first = false
-		fmt.Fprintf(ctx, "%q: %s", kv.Key, kv.Value)
+		if r.MatchString(kv.Key) {
+			first = false
+			fmt.Fprintf(ctx, "\t%q: %s", kv.Key, kv.Value)
+		}
 	})
 	fmt.Fprintf(ctx, "\n}\n")
 
 	ctx.SetContentType("application/json; charset=utf-8")
 }
+
+func getExpvarRegexp(ctx *fasthttp.RequestCtx) (*regexp.Regexp, error) {
+	r := string(ctx.QueryArgs().Peek("r"))
+	if len(r) == 0 {
+		r = "."
+	}
+	rr, err := regexp.Compile(r)
+	if err != nil {
+		return nil, fmt.Errorf("cannot parse r=%q: %s", r, err)
+	}
+	return rr, nil
+}
diff --git a/fasthttputil/expvar_test.go b/fasthttputil/expvar_test.go
index 4c409a0..fea2b25 100644
--- a/fasthttputil/expvar_test.go
+++ b/fasthttputil/expvar_test.go
@@ -3,18 +3,21 @@ package fasthttputil
 import (
 	"encoding/json"
 	"expvar"
+	"strings"
 	"testing"
 
 	"github.com/valyala/fasthttp"
 )
 
-func TestExpvarHandler(t *testing.T) {
+func TestExpvarHandlerBasic(t *testing.T) {
 	expvar.Publish("customVar", expvar.Func(func() interface{} {
 		return "foobar"
 	}))
 
 	var ctx fasthttp.RequestCtx
 
+	expvarHandlerCalls.Set(0)
+
 	ExpvarHandler(&ctx)
 
 	body := ctx.Response.Body()
@@ -31,11 +34,34 @@ func TestExpvarHandler(t *testing.T) {
 		t.Fatalf("cannot locate memstats expvar")
 	}
 
-	v, ok := m["customVar"]
+	v := m["customVar"]
+	sv, ok := v.(string)
 	if !ok {
-		t.Fatalf("cannot locate customVar")
+		t.Fatalf("unexpected custom var type %T. Expecting string", v)
 	}
-	if v != "foobar" {
+	if sv != "foobar" {
 		t.Fatalf("unexpected custom var value: %q. Expecting %q", v, "foobar")
 	}
+
+	v = m["expvarHandlerCalls"]
+	fv, ok := v.(float64)
+	if !ok {
+		t.Fatalf("unexpected expvarHandlerCalls type %T. Expecting float64", v)
+	}
+	if int(fv) != 1 {
+		t.Fatalf("unexpected value for expvarHandlerCalls: %v. Expecting %v", fv, 1)
+	}
+}
+
+func TestExpvarHandlerRegexp(t *testing.T) {
+	var ctx fasthttp.RequestCtx
+	ctx.QueryArgs().Set("r", "cmd")
+	ExpvarHandler(&ctx)
+	body := string(ctx.Response.Body())
+	if !strings.Contains(body, `"cmdline"`) {
+		t.Fatalf("missing 'cmdline' expvar")
+	}
+	if strings.Contains(body, `"memstats"`) {
+		t.Fatalf("unexpected memstats expvar found")
+	}
 }