diff --git a/pkg/common/constant/constant.go b/pkg/common/constant/constant.go index 2893f4ea7..d2188cbd9 100644 --- a/pkg/common/constant/constant.go +++ b/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 diff --git a/pkg/common/mw/check.go b/pkg/common/mw/check.go new file mode 100644 index 000000000..65c332da2 --- /dev/null +++ b/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 +} diff --git a/pkg/common/mw/check_test.go b/pkg/common/mw/check_test.go new file mode 100644 index 000000000..a70903810 --- /dev/null +++ b/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) + +} diff --git a/pkg/common/mw/rpc_client_interceptor.go b/pkg/common/mw/rpc_client_interceptor.go index 3ee56ddbd..d558965e0 100644 --- a/pkg/common/mw/rpc_client_interceptor.go +++ b/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)) diff --git a/pkg/common/mw/rpc_server_interceptor.go b/pkg/common/mw/rpc_server_interceptor.go index e6383b57c..f76fa5324 100644 --- a/pkg/common/mw/rpc_server_interceptor.go +++ b/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 {