initial import

rpcx/auth/auth.go

@@ -0,0 +1,74 @@
+package auth
+
+import (
+	"context"
+	"time"
+
+	"zero/core/collection"
+	"zero/core/stores/redis"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+)
+
+const defaultExpiration = 5 * time.Minute
+
+type Authenticator struct {
+	store  *redis.Redis
+	key    string
+	cache  *collection.Cache
+	strict bool
+}
+
+func NewAuthenticator(store *redis.Redis, key string, strict bool) (*Authenticator, error) {
+	cache, err := collection.NewCache(defaultExpiration)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Authenticator{
+		store:  store,
+		key:    key,
+		cache:  cache,
+		strict: strict,
+	}, nil
+}
+
+func (a *Authenticator) Authenticate(ctx context.Context) error {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return status.Error(codes.Unauthenticated, missingMetadata)
+	}
+
+	apps, tokens := md[appKey], md[tokenKey]
+	if len(apps) == 0 || len(tokens) == 0 {
+		return status.Error(codes.Unauthenticated, missingMetadata)
+	}
+
+	app, token := apps[0], tokens[0]
+	if len(app) == 0 || len(token) == 0 {
+		return status.Error(codes.Unauthenticated, missingMetadata)
+	}
+
+	return a.validate(app, token)
+}
+
+func (a *Authenticator) validate(app, token string) error {
+	expect, err := a.cache.Take(app, func() (interface{}, error) {
+		return a.store.Hget(a.key, app)
+	})
+	if err != nil {
+		if a.strict {
+			return status.Error(codes.Internal, err.Error())
+		} else {
+			return nil
+		}
+	}
+
+	if token != expect {
+		return status.Error(codes.Unauthenticated, accessDenied)
+	}
+
+	return nil
+}

rpcx/auth/credential.go

@@ -0,0 +1,47 @@
+package auth
+
+import (
+	"context"
+
+	"google.golang.org/grpc/metadata"
+)
+
+type Credential struct {
+	App   string
+	Token string
+}
+
+func (c *Credential) GetRequestMetadata(context.Context, ...string) (map[string]string, error) {
+	return map[string]string{
+		appKey:   c.App,
+		tokenKey: c.Token,
+	}, nil
+}
+
+func (c *Credential) RequireTransportSecurity() bool {
+	return false
+}
+
+func ParseCredential(ctx context.Context) Credential {
+	var credential Credential
+
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return credential
+	}
+
+	apps, tokens := md[appKey], md[tokenKey]
+	if len(apps) == 0 || len(tokens) == 0 {
+		return credential
+	}
+
+	app, token := apps[0], tokens[0]
+	if len(app) == 0 || len(token) == 0 {
+		return credential
+	}
+
+	credential.App = app
+	credential.Token = token
+
+	return credential
+}

rpcx/auth/vars.go

@@ -0,0 +1,9 @@
+package auth
+
+const (
+	appKey   = "app"
+	tokenKey = "token"
+
+	accessDenied    = "access denied"
+	missingMetadata = "app/token required"
+)

rpcx/client.go

@@ -0,0 +1,69 @@
+package rpcx
+
+import (
+	"log"
+	"time"
+
+	"zero/core/discov"
+	"zero/core/rpc"
+	"zero/rpcx/auth"
+
+	"google.golang.org/grpc"
+)
+
+type RpcClient struct {
+	client rpc.Client
+}
+
+func MustNewClient(c RpcClientConf, options ...rpc.ClientOption) *RpcClient {
+	cli, err := NewClient(c, options...)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return cli
+}
+
+func NewClient(c RpcClientConf, options ...rpc.ClientOption) (*RpcClient, error) {
+	var opts []rpc.ClientOption
+	if c.HasCredential() {
+		opts = append(opts, rpc.WithDialOption(grpc.WithPerRPCCredentials(&auth.Credential{
+			App:   c.App,
+			Token: c.Token,
+		})))
+	}
+	if c.Timeout > 0 {
+		opts = append(opts, rpc.WithTimeout(time.Duration(c.Timeout)*time.Millisecond))
+	}
+	opts = append(opts, options...)
+
+	var client rpc.Client
+	var err error
+	if len(c.Server) > 0 {
+		client, err = rpc.NewDirectClient(c.Server, opts...)
+	} else if err = c.Etcd.Validate(); err == nil {
+		client, err = rpc.NewRoundRobinRpcClient(c.Etcd.Hosts, c.Etcd.Key, opts...)
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	return &RpcClient{
+		client: client,
+	}, nil
+}
+
+func NewClientNoAuth(c discov.EtcdConf) (*RpcClient, error) {
+	client, err := rpc.NewRoundRobinRpcClient(c.Hosts, c.Key)
+	if err != nil {
+		return nil, err
+	}
+
+	return &RpcClient{
+		client: client,
+	}, nil
+}
+
+func (rc *RpcClient) Next() (*grpc.ClientConn, bool) {
+	return rc.client.Next()
+}

rpcx/config.go

@@ -0,0 +1,67 @@
+package rpcx
+
+import (
+	"zero/core/discov"
+	"zero/core/service"
+	"zero/core/stores/redis"
+)
+
+type (
+	RpcServerConf struct {
+		service.ServiceConf
+		ListenOn      string
+		Etcd          discov.EtcdConf    `json:",optional"`
+		Auth          bool               `json:",optional"`
+		Redis         redis.RedisKeyConf `json:",optional"`
+		StrictControl bool               `json:",optional"`
+		// pending forever is not allowed
+		// never set it to 0, if zero, the underlying will set to 2s automatically
+		Timeout      int64 `json:",default=2000"`
+		CpuThreshold int64 `json:",default=900,range=[0:1000]"`
+	}
+
+	RpcClientConf struct {
+		Etcd    discov.EtcdConf `json:",optional"`
+		Server  string          `json:",optional=!Etcd"`
+		App     string          `json:",optional"`
+		Token   string          `json:",optional"`
+		Timeout int64           `json:",optional"`
+	}
+)
+
+func NewDirectClientConf(server, app, token string) RpcClientConf {
+	return RpcClientConf{
+		Server: server,
+		App:    app,
+		Token:  token,
+	}
+}
+
+func NewEtcdClientConf(hosts []string, key, app, token string) RpcClientConf {
+	return RpcClientConf{
+		Etcd: discov.EtcdConf{
+			Hosts: hosts,
+			Key:   key,
+		},
+		App:   app,
+		Token: token,
+	}
+}
+
+func (sc RpcServerConf) HasEtcd() bool {
+	return len(sc.Etcd.Hosts) > 0 && len(sc.Etcd.Key) > 0
+}
+
+func (sc RpcServerConf) Validate() error {
+	if sc.Auth {
+		if err := sc.Redis.Validate(); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (cc RpcClientConf) HasCredential() bool {
+	return len(cc.App) > 0 && len(cc.Token) > 0
+}

rpcx/etc/config.json

@@ -0,0 +1,20 @@
+{
+    "Log": {
+        "Access": "logs/access.log",
+        "Error": "logs/error.log",
+        "Stat": "logs/stat.log"
+    },
+    "MetricsUrl": "http://localhost:2222/add",
+    "ListenOn": "localhost:3456",
+    "Etcd": {
+        "Hosts": [
+            "localhost:2379"
+        ],
+        "Key": "rpcx"
+    },
+    "Redis": {
+        "Host": "localhost:6379",
+        "Type": "node",
+        "Key": "apps"
+    }
+}

rpcx/interceptors/authinterceptor.go

@@ -0,0 +1,31 @@
+package interceptors
+
+import (
+	"context"
+
+	"zero/rpcx/auth"
+
+	"google.golang.org/grpc"
+)
+
+func StreamAuthorizeInterceptor(authenticator *auth.Authenticator) grpc.StreamServerInterceptor {
+	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo,
+		handler grpc.StreamHandler) error {
+		if err := authenticator.Authenticate(stream.Context()); err != nil {
+			return err
+		}
+
+		return handler(srv, stream)
+	}
+}
+
+func UnaryAuthorizeInterceptor(authenticator *auth.Authenticator) grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler) (interface{}, error) {
+		if err := authenticator.Authenticate(ctx); err != nil {
+			return nil, err
+		}
+
+		return handler(ctx, req)
+	}
+}

rpcx/proxy.go

@@ -0,0 +1,66 @@
+package rpcx
+
+import (
+	"context"
+	"sync"
+
+	"zero/core/rpc"
+	"zero/core/syncx"
+	"zero/rpcx/auth"
+
+	"google.golang.org/grpc"
+)
+
+type RpcProxy struct {
+	backend     string
+	clients     map[string]*RpcClient
+	options     []rpc.ClientOption
+	sharedCalls syncx.SharedCalls
+	lock        sync.Mutex
+}
+
+func NewRpcProxy(backend string, opts ...rpc.ClientOption) *RpcProxy {
+	return &RpcProxy{
+		backend:     backend,
+		clients:     make(map[string]*RpcClient),
+		options:     opts,
+		sharedCalls: syncx.NewSharedCalls(),
+	}
+}
+
+func (p *RpcProxy) TakeConn(ctx context.Context) (*grpc.ClientConn, error) {
+	cred := auth.ParseCredential(ctx)
+	key := cred.App + "/" + cred.Token
+	val, err := p.sharedCalls.Do(key, func() (interface{}, error) {
+		p.lock.Lock()
+		client, ok := p.clients[key]
+		p.lock.Unlock()
+		if ok {
+			return client, nil
+		}
+
+		client, err := NewClient(RpcClientConf{
+			Server: p.backend,
+			App:    cred.App,
+			Token:  cred.Token,
+		}, p.options...)
+		if err != nil {
+			return nil, err
+		}
+
+		p.lock.Lock()
+		p.clients[key] = client
+		p.lock.Unlock()
+		return client, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	conn, ok := val.(*RpcClient).Next()
+	if !ok {
+		return nil, grpc.ErrServerStopped
+	}
+
+	return conn, nil
+}

rpcx/server.go

@@ -0,0 +1,126 @@
+package rpcx
+
+import (
+	"log"
+	"os"
+	"strings"
+	"time"
+
+	"zero/core/load"
+	"zero/core/logx"
+	"zero/core/netx"
+	"zero/core/rpc"
+	"zero/core/rpc/serverinterceptors"
+	"zero/core/stat"
+	"zero/rpcx/auth"
+	"zero/rpcx/interceptors"
+)
+
+const envPodIp = "POD_IP"
+
+type RpcServer struct {
+	server   rpc.Server
+	register rpc.RegisterFn
+}
+
+func MustNewServer(c RpcServerConf, register rpc.RegisterFn) *RpcServer {
+	server, err := NewServer(c, register)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return server
+}
+
+func NewServer(c RpcServerConf, register rpc.RegisterFn) (*RpcServer, error) {
+	var err error
+	if err = c.Validate(); err != nil {
+		return nil, err
+	}
+
+	var server rpc.Server
+	metrics := stat.NewMetrics(c.ListenOn)
+	if c.HasEtcd() {
+		listenOn := figureOutListenOn(c.ListenOn)
+		server, err = rpc.NewRpcPubServer(c.Etcd.Hosts, c.Etcd.Key, listenOn, rpc.WithMetrics(metrics))
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		server = rpc.NewRpcServer(c.ListenOn, rpc.WithMetrics(metrics))
+	}
+
+	server.SetName(c.Name)
+	if err = setupInterceptors(server, c, metrics); err != nil {
+		return nil, err
+	}
+
+	rpcServer := &RpcServer{
+		server:   server,
+		register: register,
+	}
+	if err = c.SetUp(); err != nil {
+		return nil, err
+	}
+
+	return rpcServer, nil
+}
+
+func (rs *RpcServer) Start() {
+	if err := rs.server.Start(rs.register); err != nil {
+		logx.Error(err)
+		panic(err)
+	}
+}
+
+func (rs *RpcServer) Stop() {
+	logx.Close()
+}
+
+func figureOutListenOn(listenOn string) string {
+	fields := strings.Split(listenOn, ":")
+	if len(fields) == 0 {
+		return listenOn
+	}
+
+	host := fields[0]
+	if len(host) > 0 && host != "0.0.0.0" {
+		return listenOn
+	}
+
+	ip := os.Getenv(envPodIp)
+	if len(ip) == 0 {
+		ip = netx.InternalIp()
+	}
+	if len(ip) == 0 {
+		return listenOn
+	} else {
+		return strings.Join(append([]string{ip}, fields[1:]...), ":")
+	}
+}
+
+func setupInterceptors(server rpc.Server, c RpcServerConf, metrics *stat.Metrics) error {
+	if c.CpuThreshold > 0 {
+		shedder := load.NewAdaptiveShedder(load.WithCpuThreshold(c.CpuThreshold))
+		server.AddUnaryInterceptors(serverinterceptors.UnarySheddingInterceptor(shedder, metrics))
+	}
+
+	if c.Timeout > 0 {
+		server.AddUnaryInterceptors(serverinterceptors.UnaryTimeoutInterceptor(
+			time.Duration(c.Timeout) * time.Millisecond))
+	}
+
+	server.AddUnaryInterceptors(serverinterceptors.UnaryTracingInterceptor(c.Name))
+
+	if c.Auth {
+		authenticator, err := auth.NewAuthenticator(c.Redis.NewRedis(), c.Redis.Key, c.StrictControl)
+		if err != nil {
+			return err
+		}
+
+		server.AddStreamInterceptors(interceptors.StreamAuthorizeInterceptor(authenticator))
+		server.AddUnaryInterceptors(interceptors.UnaryAuthorizeInterceptor(authenticator))
+	}
+
+	return nil
+}