Create & Init Project...

library/net/rpc/interceptor/BUILD.bazel

@@ -0,0 +1,47 @@
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_library",
+    "go_test",
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["interceptor.go"],
+    importpath = "go-common/library/net/rpc/interceptor",
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//library/ecode:go_default_library",
+        "//library/log:go_default_library",
+        "//library/net/rpc/context:go_default_library",
+        "//library/stat:go_default_library",
+        "//vendor/golang.org/x/time/rate:go_default_library",
+    ],
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = ["interceptor_test.go"],
+    embed = [":go_default_library"],
+    rundir = ".",
+    tags = ["automanaged"],
+    deps = [
+        "//library/ecode:go_default_library",
+        "//library/net/rpc/context:go_default_library",
+        "//vendor/golang.org/x/time/rate:go_default_library",
+    ],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [":package-srcs"],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)

library/net/rpc/interceptor/interceptor.go

@@ -0,0 +1,86 @@
+package interceptor
+
+import (
+	"fmt"
+	"net"
+	"strconv"
+	"strings"
+	"time"
+
+	"go-common/library/ecode"
+	"go-common/library/log"
+	"go-common/library/net/rpc/context"
+	"go-common/library/stat"
+
+	"golang.org/x/time/rate"
+)
+
+var stats = stat.RPCServer
+
+const (
+	_innerService = "inner"
+)
+
+// Interceptor is rpc interceptor
+type Interceptor struct {
+	// for limit rate
+	rateLimits map[string]*rate.Limiter
+	// for auth
+	token string
+}
+
+// NewInterceptor new a interceptor
+func NewInterceptor(token string) *Interceptor {
+	in := &Interceptor{token: token}
+	in.rateLimits = make(map[string]*rate.Limiter)
+	return in
+}
+
+// Rate check the call is limit or not
+func (i *Interceptor) Rate(c context.Context) error {
+	limit, ok := i.rateLimits[c.ServiceMethod()]
+	if ok && !limit.Allow() {
+		return ecode.Degrade
+	}
+	return nil
+}
+
+// Stat add stat info to ops
+func (i *Interceptor) Stat(c context.Context, args interface{}, err error) {
+	const noUser = "no_user"
+	var (
+		user   = c.User()
+		method = c.ServiceMethod()
+		tmsub  = time.Since(c.Now())
+	)
+	if user == "" {
+		user = noUser
+	}
+	stats.Timing(user, int64(tmsub/time.Millisecond), method)
+	stats.Incr(user, method, strconv.Itoa((ecode.Cause(err).Code())))
+	if err != nil {
+		log.Errorv(c,
+			log.KV("args", fmt.Sprintf("%v", args)),
+			log.KV("method", method),
+			log.KV("duration", tmsub),
+			log.KV("error", fmt.Sprintf("%+v", err)),
+		)
+	} else {
+		if !strings.HasPrefix(method, _innerService) || bool(log.V(2)) {
+			log.Infov(c,
+				log.KV("args", fmt.Sprintf("%v", args)),
+				log.KV("method", method),
+				log.KV("duration", tmsub),
+				log.KV("error", fmt.Sprintf("%+v", err)),
+			)
+		}
+	}
+}
+
+// Auth check token has auth
+func (i *Interceptor) Auth(c context.Context, addr net.Addr, token string) error {
+	if i.token != token {
+		return ecode.RPCNoAuth
+	}
+	return nil
+}

library/net/rpc/interceptor/interceptor_test.go

@@ -0,0 +1,54 @@
+package interceptor
+
+import (
+	ctx "context"
+	"errors"
+	"sync"
+	"testing"
+
+	"go-common/library/ecode"
+	"go-common/library/net/rpc/context"
+
+	"golang.org/x/time/rate"
+)
+
+var (
+	once sync.Once
+	i    *Interceptor
+	c    context.Context
+)
+
+func interceptor() {
+	i = NewInterceptor("test token")
+	c = context.NewContext(ctx.TODO(), "testMethod", "test user", 0)
+}
+
+func TestRate(t *testing.T) {
+	once.Do(interceptor)
+	if err := i.Rate(c); err != nil {
+		t.Errorf("TestRate error(%v)", err)
+		t.FailNow()
+	}
+	i.rateLimits["testMethod"] = rate.NewLimiter(1, 0)
+	if err := i.Rate(c); err != ecode.Degrade {
+		t.Errorf("TestRate error(%v)", err)
+		t.FailNow()
+	}
+}
+
+func TestStat(t *testing.T) {
+	once.Do(interceptor)
+	i.Stat(c, nil, errors.New("test error"))
+}
+
+func TestAuth(t *testing.T) {
+	once.Do(interceptor)
+	if err := i.Auth(c, nil, "test token"); err != nil {
+		t.Errorf("TestAuth error(%v)", err)
+		t.FailNow()
+	}
+	if err := i.Auth(c, nil, "token"); err != ecode.RPCNoAuth {
+		t.Errorf("TestAuth error(%v)", err)
+		t.FailNow()
+	}
+}