blackteay/open-im-server
1
0
Fork 0
mirror of https://github.com/openimsdk/open-im-server.git synced 2025-04-23 09:50:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

rpc custom header

Browse Source
This commit is contained in:
withchao 2023-03-22 10:11:18 +08:00
parent c0d72d91fb
commit 900858e4ea
4 changed files with 17 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
pkg/common/constant/constant.go
View File

@ -275,7 +275,7 @@ const OpUserID = "opUserID"
const ConnID = "connID" const ConnID = "connID"
const OpUserPlatform = "platform" const OpUserPlatform = "platform"
const Token = "token" const Token = "token"
const RpcMwCustom = "hCustom" const RpcCustomHeader = "customHeader" // rpc中间件自定义ctx参数
const CheckKey = "CheckKey" const CheckKey = "CheckKey"
const ( const (

4
pkg/common/mw/check.go
View File

@ -15,8 +15,8 @@ import (
) )
var ( var (
block cipher.Block
once sync.Once once sync.Once
block cipher.Block
) )
func init() { func init() {
@ -37,7 +37,7 @@ func initAesKey() {
func genReqKey(args []string) string { func genReqKey(args []string) string {
initAesKey() initAesKey()
plaintext := md5.Sum([]byte(strings.Join(args, ":"))) plaintext := md5.Sum([]byte(strings.Join(args, ":")))
var iv = make([]byte, aes.BlockSize, aes.BlockSize+md5.Size) iv := make([]byte, aes.BlockSize, aes.BlockSize+md5.Size)
if _, err := rand.Read(iv); err != nil { if _, err := rand.Read(iv); err != nil {
panic(err) panic(err)
} }

6
pkg/common/mw/rpc_client_interceptor.go
View File

@ -24,14 +24,18 @@ func rpcClientInterceptor(ctx context.Context, method string, req, resp interfac
} }
log.ZInfo(ctx, "rpc client req", "funcName", method, "req", rpcString(req)) log.ZInfo(ctx, "rpc client req", "funcName", method, "req", rpcString(req))
md := metadata.Pairs() md := metadata.Pairs()
if keys, _ := ctx.Value(constant.RpcMwCustom).([]string); len(keys) > 0 { if keys, _ := ctx.Value(constant.RpcCustomHeader).([]string); len(keys) > 0 {
for _, key := range keys { for _, key := range keys {
val, ok := ctx.Value(key).([]string) val, ok := ctx.Value(key).([]string)
if !ok { if !ok {
return errs.ErrInternalServer.Wrap(fmt.Sprintf("ctx missing key %s", key)) return errs.ErrInternalServer.Wrap(fmt.Sprintf("ctx missing key %s", key))
} }
if len(val) == 0 {
return errs.ErrInternalServer.Wrap(fmt.Sprintf("ctx key %s value is empty", key))
}
md.Set(key, val...) md.Set(key, val...)
} }
md.Set(constant.RpcCustomHeader, keys...)
} }
operationID, ok := ctx.Value(constant.OperationID).(string) operationID, ok := ctx.Value(constant.OperationID).(string)
if !ok { if !ok {

9
pkg/common/mw/rpc_server_interceptor.go
View File

@ -56,6 +56,15 @@ func rpcServerInterceptor(ctx context.Context, req interface{}, info *grpc.Unary
if !ok { if !ok {
return nil, status.New(codes.InvalidArgument, "missing metadata").Err() return nil, status.New(codes.InvalidArgument, "missing metadata").Err()
} }
if keys := md.Get(constant.RpcCustomHeader); len(keys) > 0 {
for _, key := range keys {
values := md.Get(key)
if len(values) == 0 {
return nil, status.New(codes.InvalidArgument, fmt.Sprintf("missing metadata key %s", key)).Err()
}
ctx = context.WithValue(ctx, key, values)
}
}
args := make([]string, 0, 4) args := make([]string, 0, 4)
if opts := md.Get(constant.OperationID); len(opts) != 1 || opts[0] == "" { if opts := md.Get(constant.OperationID); len(opts) != 1 || opts[0] == "" {
return nil, status.New(codes.InvalidArgument, "operationID error").Err() return nil, status.New(codes.InvalidArgument, "operationID error").Err()