Add handshake validators

Add handshake validators

agent/agent.go

@@ -52,7 +52,9 @@ var (
 	// hbd contains the heartbeat packet data
 	hbd []byte
 	// hrd contains the handshake response data
-	hrd  []byte
+	hrd []byte
+	// herd contains the handshake error response data
+	herd []byte
 	once sync.Once
 )
 
@@ -110,6 +112,7 @@ type (
 		Handle()
 		IPVersion() string
 		SendHandshakeResponse() error
+		SendHandshakeErrorResponse() error
 		SendRequest(ctx context.Context, serverID, route string, v interface{}) (*protos.Response, error)
 		AnswerWithError(ctx context.Context, mid uint, err error)
 	}
@@ -180,6 +183,7 @@ func newAgent(
 
 	once.Do(func() {
 		hbdEncode(heartbeatTime, packetEncoder, messageEncoder.IsCompressionEnabled(), serializerName)
+		herdEncode(heartbeatTime, packetEncoder, messageEncoder.IsCompressionEnabled(), serializerName)
 	})
 
 	a := &agentImpl{
@@ -475,6 +479,14 @@ func (a *agentImpl) onSessionClosed(s session.Session) {
 // SendHandshakeResponse sends a handshake response
 func (a *agentImpl) SendHandshakeResponse() error {
 	_, err := a.conn.Write(hrd)
+
+	return err
+}
+
+func (a *agentImpl) SendHandshakeErrorResponse() error {
+	a.SetStatus(constants.StatusClosed)
+	_, err := a.conn.Write(herd)
+
 	return err
 }
 
@@ -543,33 +555,63 @@ func hbdEncode(heartbeatTimeout time.Duration, packetEncoder codec.PacketEncoder
 			"serializer": serializerName,
 		},
 	}
-	data, err := gojson.Marshal(hData)
+
+	data, err := encodeAndCompress(hData, dataCompression)
 	if err != nil {
 		panic(err)
 	}
 
-	if dataCompression {
-		compressedData, err := compression.DeflateData(data)
-		if err != nil {
-			panic(err)
-		}
+	hrd, err = packetEncoder.Encode(packet.Handshake, data)
+	if err != nil {
+		panic(err)
+	}
 
-		if len(compressedData) < len(data) {
-			data = compressedData
-		}
+	hbd, err = packetEncoder.Encode(packet.Heartbeat, nil)
+	if err != nil {
+		panic(err)
 	}
+}
 
-	hrd, err = packetEncoder.Encode(packet.Handshake, data)
+func herdEncode(heartbeatTimeout time.Duration, packetEncoder codec.PacketEncoder, dataCompression bool, serializerName string) {
+	hErrData := map[string]interface{}{
+		"code": 400,
+		"sys": map[string]interface{}{
+			"heartbeat":  heartbeatTimeout.Seconds(),
+			"dict":       message.GetDictionary(),
+			"serializer": serializerName,
+		},
+	}
+
+	errData, err := encodeAndCompress(hErrData, dataCompression)
 	if err != nil {
 		panic(err)
 	}
 
-	hbd, err = packetEncoder.Encode(packet.Heartbeat, nil)
+	herd, err = packetEncoder.Encode(packet.Handshake, errData)
 	if err != nil {
 		panic(err)
 	}
 }
 
+func encodeAndCompress(data interface{}, dataCompression bool) ([]byte, error) {
+	encData, err := gojson.Marshal(data)
+	if err != nil {
+		return nil, err
+	}
+
+	if dataCompression {
+		compressedData, err := compression.DeflateData(encData)
+		if err != nil {
+			return nil, err
+		}
+
+		if len(compressedData) < len(encData) {
+			encData = compressedData
+		}
+	}
+	return encData, nil
+}
+
 func (a *agentImpl) reportChannelSize() {
 	chSendCapacity := a.messagesBufferSize - len(a.chSend)
 	if chSendCapacity == 0 {

agent/agent_test.go

@@ -84,7 +84,7 @@ func TestNewAgent(t *testing.T) {
 		func(typ packet.Type, d []byte) {
 			// cannot compare inside the expect because they are equivalent but not equal
 			assert.EqualValues(t, packet.Handshake, typ)
-		})
+		}).Times(2)
 	mockEncoder.EXPECT().Encode(gomock.Any(), gomock.Nil()).Do(
 		func(typ packet.Type, d []byte) {
 			assert.EqualValues(t, packet.Heartbeat, typ)