-
Notifications
You must be signed in to change notification settings - Fork 786
/
util.go
117 lines (107 loc) · 3.86 KB
/
util.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
/*
* Copyright 2021 CloudWeGo Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package codec
import (
"errors"
"fmt"
"github.com/cloudwego/kitex/pkg/remote"
"github.com/cloudwego/kitex/pkg/rpcinfo"
)
const (
// FrontMask is used in protocol sniffing.
FrontMask = 0x0000ffff
)
// SetOrCheckMethodName is used to set method name to invocation.
func SetOrCheckMethodName(methodName string, message remote.Message) error {
ri := message.RPCInfo()
ink := ri.Invocation()
callMethodName := ink.MethodName()
if methodName == "" {
return fmt.Errorf("method name that receive is empty")
}
if callMethodName == methodName {
return nil
}
// the server's callMethodName may not be empty if RPCInfo is based on connection multiplexing
// for the server side callMethodName ! = methodName is normal
if message.RPCRole() == remote.Client {
return fmt.Errorf("wrong method name, expect=%s, actual=%s", callMethodName, methodName)
}
svcInfo, err := message.SpecifyServiceInfo(ink.ServiceName(), methodName)
if err != nil {
return err
}
if ink, ok := ink.(rpcinfo.InvocationSetter); ok {
ink.SetMethodName(methodName)
ink.SetPackageName(svcInfo.GetPackageName())
ink.SetServiceName(svcInfo.ServiceName)
} else {
return errors.New("the interface Invocation doesn't implement InvocationSetter")
}
// unknown method doesn't set methodName for RPCInfo.To(), or lead inconsistent with old version
rpcinfo.AsMutableEndpointInfo(ri.To()).SetMethod(methodName)
return nil
}
// SetOrCheckSeqID is used to check the sequence ID.
func SetOrCheckSeqID(seqID int32, message remote.Message) error {
switch message.MessageType() {
case remote.Call, remote.Oneway:
if ink, ok := message.RPCInfo().Invocation().(rpcinfo.InvocationSetter); ok {
ink.SetSeqID(seqID)
} else {
return errors.New("the interface Invocation doesn't implement InvocationSetter")
}
case remote.Reply:
expectSeqID := message.RPCInfo().Invocation().SeqID()
if expectSeqID != seqID {
methodName := message.RPCInfo().Invocation().MethodName()
return remote.NewTransErrorWithMsg(remote.BadSequenceID, fmt.Sprintf("method[%s] out of order sequence response, expect[%d], receive[%d]", methodName, expectSeqID, seqID))
}
case remote.Exception:
// don't check, proxy may build Exception with seqID = 0
// thrift 0.13 check seqID for Exception but thrift 0.9.2 doesn't check
}
return nil
}
// UpdateMsgType updates msg type.
func UpdateMsgType(msgType uint32, message remote.Message) error {
rpcRole := message.RPCRole()
mt := remote.MessageType(msgType)
if mt == message.MessageType() {
return nil
}
if rpcRole == remote.Server {
if mt != remote.Call && mt != remote.Oneway && mt != remote.Stream {
return remote.NewTransErrorWithMsg(remote.InvalidMessageTypeException, fmt.Sprintf("server side, invalid message type %d", mt))
}
} else {
if mt != remote.Reply && mt != remote.Exception && mt != remote.Stream {
return remote.NewTransErrorWithMsg(remote.InvalidMessageTypeException, fmt.Sprintf("client side, invalid message type %d", mt))
}
}
message.SetMessageType(mt)
return nil
}
// NewDataIfNeeded is used to create the data if not exist.
func NewDataIfNeeded(method string, message remote.Message) error {
if message.Data() != nil {
return nil
}
if message.NewData(method) {
return nil
}
return remote.NewTransErrorWithMsg(remote.InternalError, "message data for codec is nil")
}