fix: Fix max packet size validation

DrmagicE · Apr 18, 2021 · d986194 · d986194
1 parent 2ed97bd
commit d986194
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 30 deletions.
diff --git 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
@@ -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
@@ -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)
 }