diff --git a/go.mod b/go.mod index 086e6462..84b49d9e 100644 --- a/go.mod +++ b/go.mod @@ -22,7 +22,7 @@ require ( go.uber.org/zap v1.13.0 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9 golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd - golang.org/x/tools v0.0.0-20191112195655-aa38f8e97acc // indirect + golang.org/x/tools v0.0.0-20191112195655-aa38f8e97acc google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013 google.golang.org/grpc v1.34.0 google.golang.org/protobuf v1.25.0 diff --git a/pkg/packets/packets.go b/pkg/packets/packets.go index c302a335..7dbd8c0d 100644 --- a/pkg/packets/packets.go +++ b/pkg/packets/packets.go @@ -178,21 +178,6 @@ func (r *Reader) ReadPacket() (Packet, error) { } fh := &FixHeader{PacketType: first >> 4, Flags: first & 15} //设置FixHeader length, err := EncodeRemainLength(r.bufr) - - var headerLen int - if length <= 127 { - headerLen = 2 - } else if length <= 16383 { - headerLen = 3 - } else if length <= 2097151 { - headerLen = 4 - } else if length <= 268435455 { - headerLen = 5 - } - if headerLen+length > 1234 { - return nil, codes.NewError(codes.RecvMaxExceeded) - } - if err != nil { return nil, err } @@ -296,6 +281,9 @@ func EncodeRemainLength(r io.ByteReader) (int, error) { return 0, err } vbi |= uint32(digit&127) << multiplier + if vbi > 268435455 { + return 0, codes.ErrMalformed + } if (digit & 128) == 0 { break } diff --git a/plugin/federation/hooks_test.go b/plugin/federation/hooks_test.go index c5abd11e..637c7430 100644 --- a/plugin/federation/hooks_test.go +++ b/plugin/federation/hooks_test.go @@ -302,6 +302,7 @@ func TestFederation_OnUnsubscribedWrapper(t *testing.T) { func TestFederation_OnSessionTerminatedWrapper(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() + a := assert.New(t) p, _ := New(testConfig) f := p.(*Federation) f.localSubStore.init(mem.NewStore()) @@ -338,19 +339,17 @@ func TestFederation_OnSessionTerminatedWrapper(t *testing.T) { }) onSessionTerminated(context.Background(), "client2", 0) - mockQueue.EXPECT().add(&Event{ - Event: &Event_Unsubscribe{ - Unsubscribe: &Unsubscribe{ - TopicName: "/topicB", - }, - }, - }) - mockQueue.EXPECT().add(&Event{ - Event: &Event_Unsubscribe{ - Unsubscribe: &Unsubscribe{ - TopicName: "/topicC", - }, - }, - }) + var b, c bool + mockQueue.EXPECT().add(gomock.Any()).Do(func(event *Event) { + if event.Event.(*Event_Unsubscribe).Unsubscribe.TopicName == "/topicB" { + b = true + } + if event.Event.(*Event_Unsubscribe).Unsubscribe.TopicName == "/topicC" { + c = true + } + }).Times(2) + onSessionTerminated(context.Background(), "client3", 0) + a.True(b) + a.True(c) }