diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/remoting/activity/ClientManagerActivity.java b/proxy/src/main/java/org/apache/rocketmq/proxy/remoting/activity/ClientManagerActivity.java index 05d8e5fbe13..a4ee47c465e 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/remoting/activity/ClientManagerActivity.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/remoting/activity/ClientManagerActivity.java @@ -44,6 +44,9 @@ import java.util.Set; public class ClientManagerActivity extends AbstractRemotingActivity { + static final String EMPTY_HEARTBEAT_DATA_REMARK = "heartbeat data is empty"; + static final String MISSING_HEARTBEAT_DATA_SET_REMARK = "heartbeat producerDataSet and consumerDataSet are required"; + private final RemotingChannelManager remotingChannelManager; @@ -78,6 +81,13 @@ protected RemotingCommand processRequest0(ChannelHandlerContext ctx, RemotingCom protected RemotingCommand heartBeat(ChannelHandlerContext ctx, RemotingCommand request, ProxyContext context) { HeartbeatData heartbeatData = HeartbeatData.decode(request.getBody(), HeartbeatData.class); + if (heartbeatData == null) { + return RemotingCommand.buildErrorResponse(ResponseCode.INVALID_PARAMETER, EMPTY_HEARTBEAT_DATA_REMARK); + } + if (heartbeatData.getProducerDataSet() == null || heartbeatData.getConsumerDataSet() == null) { + return RemotingCommand.buildErrorResponse(ResponseCode.INVALID_PARAMETER, + MISSING_HEARTBEAT_DATA_SET_REMARK); + } String clientId = heartbeatData.getClientID(); for (ProducerData data : heartbeatData.getProducerDataSet()) { diff --git a/proxy/src/test/java/org/apache/rocketmq/proxy/remoting/activity/ClientManagerActivityTest.java b/proxy/src/test/java/org/apache/rocketmq/proxy/remoting/activity/ClientManagerActivityTest.java new file mode 100644 index 00000000000..2e0caf4f51c --- /dev/null +++ b/proxy/src/test/java/org/apache/rocketmq/proxy/remoting/activity/ClientManagerActivityTest.java @@ -0,0 +1,97 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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 org.apache.rocketmq.proxy.remoting.activity; + +import java.nio.charset.StandardCharsets; +import org.apache.rocketmq.proxy.common.ProxyContext; +import org.apache.rocketmq.proxy.config.InitConfigTest; +import org.apache.rocketmq.proxy.processor.MessagingProcessor; +import org.apache.rocketmq.proxy.remoting.channel.RemotingChannelManager; +import org.apache.rocketmq.remoting.protocol.RemotingCommand; +import org.apache.rocketmq.remoting.protocol.RequestCode; +import org.apache.rocketmq.remoting.protocol.ResponseCode; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.junit.MockitoJUnitRunner; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; + +@RunWith(MockitoJUnitRunner.class) +public class ClientManagerActivityTest extends InitConfigTest { + private ClientManagerActivity clientManagerActivity; + + @Mock + private MessagingProcessor messagingProcessor; + @Mock + private RemotingChannelManager remotingChannelManager; + + @Before + public void setup() { + this.clientManagerActivity = new ClientManagerActivity(null, messagingProcessor, remotingChannelManager); + } + + @Test + public void testHeartbeatShouldRejectEmptyBody() { + RemotingCommand request = RemotingCommand.createRequestCommand(RequestCode.HEART_BEAT, null); + + RemotingCommand response = clientManagerActivity.heartBeat(null, request, ProxyContext.create()); + + assertThat(response.getCode()).isEqualTo(ResponseCode.INVALID_PARAMETER); + assertThat(response.getRemark()).isEqualTo(ClientManagerActivity.EMPTY_HEARTBEAT_DATA_REMARK); + verifyNoClientRegistration(); + } + + @Test + public void testHeartbeatShouldRejectNullDataSets() { + RemotingCommand request = RemotingCommand.createRequestCommand(RequestCode.HEART_BEAT, null); + request.setBody("{\"clientID\":\"client-a\",\"producerDataSet\":null,\"consumerDataSet\":[]}" + .getBytes(StandardCharsets.UTF_8)); + + RemotingCommand response = clientManagerActivity.heartBeat(null, request, ProxyContext.create()); + + assertThat(response.getCode()).isEqualTo(ResponseCode.INVALID_PARAMETER); + assertThat(response.getRemark()).isEqualTo(ClientManagerActivity.MISSING_HEARTBEAT_DATA_SET_REMARK); + verifyNoClientRegistration(); + } + + @Test + public void testHeartbeatShouldRejectNullConsumerDataSet() { + RemotingCommand request = RemotingCommand.createRequestCommand(RequestCode.HEART_BEAT, null); + request.setBody("{\"clientID\":\"client-a\",\"producerDataSet\":[],\"consumerDataSet\":null}" + .getBytes(StandardCharsets.UTF_8)); + + RemotingCommand response = clientManagerActivity.heartBeat(null, request, ProxyContext.create()); + + assertThat(response.getCode()).isEqualTo(ResponseCode.INVALID_PARAMETER); + assertThat(response.getRemark()).isEqualTo(ClientManagerActivity.MISSING_HEARTBEAT_DATA_SET_REMARK); + verifyNoClientRegistration(); + } + + private void verifyNoClientRegistration() { + verify(messagingProcessor, never()).registerProducer(any(), any(), any()); + verify(messagingProcessor, never()).registerConsumer(any(), any(), any(), any(), any(), any(), any(), anyBoolean()); + verifyNoInteractions(remotingChannelManager); + } +}