verify req key

pkg/common/constant/constant.go

@@ -275,6 +275,8 @@ const OpUserID = "opUserID"
 const ConnID = "connID"
 const OpUserPlatform = "platform"
 const Token = "token"
+const RpcMwCustom = "hCustom"
+const CheckKey = "CheckKey"
 
 const (
 	UnreliableNotification    = 1

pkg/common/mw/check.go

@@ -0,0 +1,65 @@
+package mw
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/md5"
+	"encoding/base64"
+	"errors"
+	"fmt"
+	"github.com/OpenIMSDK/Open-IM-Server/pkg/common/config"
+	"math/rand"
+	"strings"
+	"sync"
+	"time"
+)
+
+var (
+	block cipher.Block
+	once  sync.Once
+)
+
+func init() {
+	rand.Seed(time.Now().UnixNano())
+}
+
+func initAesKey() {
+	once.Do(func() {
+		key := md5.Sum([]byte("openim:" + config.Config.Secret))
+		var err error
+		block, err = aes.NewCipher(key[:])
+		if err != nil {
+			panic(err)
+		}
+	})
+}
+
+func genReqKey(args []string) string {
+	initAesKey()
+	plaintext := md5.Sum([]byte(strings.Join(args, ":")))
+	var iv = make([]byte, aes.BlockSize, aes.BlockSize+md5.Size)
+	if _, err := rand.Read(iv); err != nil {
+		panic(err)
+	}
+	ciphertext := make([]byte, md5.Size)
+	cipher.NewCBCEncrypter(block, iv).CryptBlocks(ciphertext, plaintext[:])
+	return base64.StdEncoding.EncodeToString(append(iv, ciphertext...))
+}
+
+func verifyReqKey(args []string, key string) error {
+	initAesKey()
+	k, err := base64.StdEncoding.DecodeString(key)
+	if err != nil {
+		return fmt.Errorf("invalid key %v", err)
+	}
+	if len(k) != aes.BlockSize+md5.Size {
+		return errors.New("invalid key")
+	}
+	plaintext := make([]byte, md5.Size)
+	cipher.NewCBCDecrypter(block, k[:aes.BlockSize]).CryptBlocks(plaintext, k[aes.BlockSize:])
+	sum := md5.Sum([]byte(strings.Join(args, ":")))
+	if string(plaintext) != string(sum[:]) {
+		return errors.New("mismatch key")
+	}
+	return nil
+}

pkg/common/mw/check_test.go

@@ -0,0 +1,28 @@
+package mw
+
+import (
+	"fmt"
+	"github.com/OpenIMSDK/Open-IM-Server/pkg/common/config"
+	"testing"
+)
+
+func TestCheck(t *testing.T) {
+	config.Config.Secret = "123456"
+
+	args := []string{"1", "2", "3"}
+
+	key := genReqKey(args)
+	fmt.Println("key:", key)
+	err := verifyReqKey(args, key)
+
+	fmt.Println(err)
+
+	args = []string{"4", "5", "6"}
+
+	key = genReqKey(args)
+	fmt.Println("key:", key)
+	err = verifyReqKey(args, key)
+
+	fmt.Println(err)
+
+}

pkg/common/mw/rpc_client_interceptor.go

@@ -3,6 +3,7 @@ package mw
 import (
 	"context"
 	"errors"
+	"fmt"
 	"github.com/OpenIMSDK/Open-IM-Server/pkg/common/constant"
 	"github.com/OpenIMSDK/Open-IM-Server/pkg/common/log"
 	"github.com/OpenIMSDK/Open-IM-Server/pkg/errs"
@@ -22,20 +23,34 @@ func rpcClientInterceptor(ctx context.Context, method string, req, resp interfac
 		return errs.ErrInternalServer.Wrap("call rpc request context is nil")
 	}
 	log.ZInfo(ctx, "rpc client req", "funcName", method, "req", rpcString(req))
+	md := metadata.Pairs()
+	if keys, _ := ctx.Value(constant.RpcMwCustom).([]string); len(keys) > 0 {
+		for _, key := range keys {
+			val, ok := ctx.Value(key).([]string)
+			if !ok {
+				return errs.ErrInternalServer.Wrap(fmt.Sprintf("ctx missing key %s", key))
+			}
+			md.Set(key, val...)
+		}
+	}
 	operationID, ok := ctx.Value(constant.OperationID).(string)
 	if !ok {
 		log.ZWarn(ctx, "ctx missing operationID", errors.New("ctx missing operationID"), "funcName", method)
 		return errs.ErrArgs.Wrap("ctx missing operationID")
 	}
-	md := metadata.Pairs(constant.OperationID, operationID)
+	md.Set(constant.OperationID, operationID)
+	args := make([]string, 0, 4)
+	args = append(args, constant.OperationID, operationID)
 	opUserID, ok := ctx.Value(constant.OpUserID).(string)
 	if ok {
-		md.Append(constant.OpUserID, opUserID)
+		md.Set(constant.OpUserID, opUserID)
+		args = append(args, constant.OpUserID, opUserID)
 	}
 	opUserIDPlatformID, ok := ctx.Value(constant.OpUserPlatform).(string)
 	if ok {
-		md.Append(constant.OpUserPlatform, opUserIDPlatformID)
+		md.Set(constant.OpUserPlatform, opUserIDPlatformID)
 	}
+	md.Set(constant.CheckKey, genReqKey(args))
 	err = invoker(metadata.NewOutgoingContext(ctx, md), method, req, resp, cc, opts...)
 	if err == nil {
 		log.ZInfo(ctx, "rpc client resp", "funcName", method, "resp", rpcString(resp))

pkg/common/mw/rpc_server_interceptor.go

@@ -56,17 +56,27 @@ func rpcServerInterceptor(ctx context.Context, req interface{}, info *grpc.Unary
 	if !ok {
 		return nil, status.New(codes.InvalidArgument, "missing metadata").Err()
 	}
+	args := make([]string, 0, 4)
 	if opts := md.Get(constant.OperationID); len(opts) != 1 || opts[0] == "" {
 		return nil, status.New(codes.InvalidArgument, "operationID error").Err()
 	} else {
+		args = append(args, constant.OperationID, opts[0])
 		ctx = context.WithValue(ctx, constant.OperationID, opts[0])
 	}
 	if opts := md.Get(constant.OpUserID); len(opts) == 1 {
+		args = append(args, constant.OpUserID, opts[0])
 		ctx = context.WithValue(ctx, constant.OpUserID, opts[0])
 	}
 	if opts := md.Get(constant.OpUserPlatform); len(opts) == 1 {
 		ctx = context.WithValue(ctx, constant.OpUserPlatform, opts[0])
 	}
+	if opts := md.Get(constant.CheckKey); len(opts) != 1 || opts[0] == "" {
+		return nil, status.New(codes.InvalidArgument, "check key empty").Err()
+	} else {
+		if err := verifyReqKey(args, opts[0]); err != nil {
+			return nil, status.New(codes.InvalidArgument, err.Error()).Err()
+		}
+	}
 	log.ZInfo(ctx, "rpc server req", "funcName", funcName, "req", rpcString(req))
 	resp, err = handler(ctx, req)
 	if err == nil {