diff --git a/store/src/main/java/org/apache/rocketmq/store/ha/HAConnectionState.java b/store/src/main/java/org/apache/rocketmq/store/ha/HAConnectionState.java index 4f0c5ca9095..fd246db5b4e 100644 --- a/store/src/main/java/org/apache/rocketmq/store/ha/HAConnectionState.java +++ b/store/src/main/java/org/apache/rocketmq/store/ha/HAConnectionState.java @@ -37,5 +37,10 @@ public enum HAConnectionState { /** * Connection shutdown. */ - SHUTDOWN, + SHUTDOWN; + + public static HAConnectionState fromOrdinal(int ordinal) { + HAConnectionState[] states = values(); + return ordinal >= 0 && ordinal < states.length ? states[ordinal] : null; + } } diff --git a/store/src/main/java/org/apache/rocketmq/store/ha/autoswitch/AutoSwitchHAClient.java b/store/src/main/java/org/apache/rocketmq/store/ha/autoswitch/AutoSwitchHAClient.java index 3dd14f4e35c..eb3e601d149 100644 --- a/store/src/main/java/org/apache/rocketmq/store/ha/autoswitch/AutoSwitchHAClient.java +++ b/store/src/main/java/org/apache/rocketmq/store/ha/autoswitch/AutoSwitchHAClient.java @@ -504,7 +504,7 @@ protected boolean processReadResult(ByteBuffer byteBufferRead) { AutoSwitchHAClient.this.processPosition += headerSize + bodySize; AutoSwitchHAClient.this.waitForRunning(1); LOGGER.error("State not matched, masterState:{}, slaveState:{}, bodySize:{}, offset:{}, masterEpoch:{}, masterEpochStartOffset:{}, confirmOffset:{}", - HAConnectionState.values()[masterState], AutoSwitchHAClient.this.currentState, bodySize, masterOffset, masterEpoch, masterEpochStartOffset, confirmOffset); + HAConnectionState.fromOrdinal(masterState), AutoSwitchHAClient.this.currentState, bodySize, masterOffset, masterEpoch, masterEpochStartOffset, confirmOffset); return false; } diff --git a/store/src/main/java/org/apache/rocketmq/store/ha/autoswitch/AutoSwitchHAConnection.java b/store/src/main/java/org/apache/rocketmq/store/ha/autoswitch/AutoSwitchHAConnection.java index cc55937aebb..e67ce46b0c6 100644 --- a/store/src/main/java/org/apache/rocketmq/store/ha/autoswitch/AutoSwitchHAConnection.java +++ b/store/src/main/java/org/apache/rocketmq/store/ha/autoswitch/AutoSwitchHAConnection.java @@ -310,7 +310,12 @@ protected boolean processReadResult(ByteBuffer byteBufferRead) { int diff = byteBufferRead.position() - ReadSocketService.this.processPosition; if (diff >= AutoSwitchHAClient.MIN_HEADER_SIZE) { int readPosition = ReadSocketService.this.processPosition; - HAConnectionState slaveState = HAConnectionState.values()[byteBufferRead.getInt(readPosition)]; + int slaveStateOrdinal = byteBufferRead.getInt(readPosition); + HAConnectionState slaveState = HAConnectionState.fromOrdinal(slaveStateOrdinal); + if (slaveState == null) { + LOGGER.error("Received illegal slave state ordinal {}", slaveStateOrdinal); + return false; + } switch (slaveState) { case HANDSHAKE: diff --git a/store/src/test/java/org/apache/rocketmq/store/ha/HAConnectionStateTest.java b/store/src/test/java/org/apache/rocketmq/store/ha/HAConnectionStateTest.java new file mode 100644 index 00000000000..7e98c78abbe --- /dev/null +++ b/store/src/test/java/org/apache/rocketmq/store/ha/HAConnectionStateTest.java @@ -0,0 +1,32 @@ +/* + * 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.store.ha; + +import org.junit.Assert; +import org.junit.Test; + +public class HAConnectionStateTest { + + @Test + public void testFromOrdinal() { + Assert.assertEquals(HAConnectionState.READY, HAConnectionState.fromOrdinal(0)); + Assert.assertEquals(HAConnectionState.SHUTDOWN, + HAConnectionState.fromOrdinal(HAConnectionState.values().length - 1)); + Assert.assertNull(HAConnectionState.fromOrdinal(-1)); + Assert.assertNull(HAConnectionState.fromOrdinal(HAConnectionState.values().length)); + } +} diff --git a/store/src/test/java/org/apache/rocketmq/store/ha/autoswitch/AutoSwitchHAStateValidationTest.java b/store/src/test/java/org/apache/rocketmq/store/ha/autoswitch/AutoSwitchHAStateValidationTest.java new file mode 100644 index 00000000000..73365f412d4 --- /dev/null +++ b/store/src/test/java/org/apache/rocketmq/store/ha/autoswitch/AutoSwitchHAStateValidationTest.java @@ -0,0 +1,98 @@ +/* + * 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.store.ha.autoswitch; + +import java.lang.reflect.Field; +import java.net.InetSocketAddress; +import java.nio.ByteBuffer; +import java.nio.channels.ServerSocketChannel; +import java.nio.channels.SocketChannel; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.rocketmq.common.BrokerConfig; +import org.apache.rocketmq.store.DefaultMessageStore; +import org.apache.rocketmq.store.config.MessageStoreConfig; +import org.junit.Test; + +import static org.junit.Assert.assertFalse; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class AutoSwitchHAStateValidationTest { + + @Test + public void testServerReaderRejectsUnknownSlaveStateOrdinal() throws Exception { + AutoSwitchHAService haService = mock(AutoSwitchHAService.class); + DefaultMessageStore messageStore = mock(DefaultMessageStore.class); + when(haService.getConnectionCount()).thenReturn(new AtomicInteger()); + when(haService.getDefaultMessageStore()).thenReturn(messageStore); + when(messageStore.getBrokerConfig()).thenReturn(new BrokerConfig()); + when(messageStore.getMessageStoreConfig()).thenReturn(new MessageStoreConfig()); + + try (ServerSocketChannel serverSocket = ServerSocketChannel.open(); + SocketChannel slaveSocket = SocketChannel.open()) { + serverSocket.bind(new InetSocketAddress("127.0.0.1", 0)); + slaveSocket.connect(serverSocket.getLocalAddress()); + + try (SocketChannel masterSocket = serverSocket.accept()) { + AutoSwitchHAConnection connection = new AutoSwitchHAConnection( + haService, masterSocket, mock(EpochFileCache.class)); + try { + AutoSwitchHAConnection.ReadSocketService readSocketService = + getReadSocketService(connection); + AutoSwitchHAConnection.ReadSocketService.HAServerReader reader = + readSocketService.new HAServerReader(); + ByteBuffer frame = ByteBuffer.allocate(AutoSwitchHAClient.MIN_HEADER_SIZE); + frame.putInt(Integer.MAX_VALUE); + frame.position(AutoSwitchHAClient.MIN_HEADER_SIZE); + + assertFalse(reader.processReadResult(frame)); + } finally { + connection.shutdown(); + } + } + } + } + + @Test + public void testClientReaderRejectsUnknownMasterStateOrdinal() throws Exception { + AutoSwitchHAService haService = mock(AutoSwitchHAService.class); + DefaultMessageStore messageStore = mock(DefaultMessageStore.class); + when(haService.getDefaultMessageStore()).thenReturn(messageStore); + when(messageStore.getBrokerConfig()).thenReturn(new BrokerConfig()); + when(messageStore.getMessageStoreConfig()).thenReturn(new MessageStoreConfig()); + AutoSwitchHAClient client = new AutoSwitchHAClient( + haService, messageStore, mock(EpochFileCache.class), 1L); + try { + ByteBuffer frame = ByteBuffer.allocate(AutoSwitchHAConnection.HANDSHAKE_HEADER_SIZE); + frame.putInt(Integer.MAX_VALUE); + frame.putInt(0); + frame.putLong(0); + frame.putInt(0); + + assertFalse(client.new HAClientReader().processReadResult(frame)); + } finally { + client.shutdown(); + } + } + + private AutoSwitchHAConnection.ReadSocketService getReadSocketService( + AutoSwitchHAConnection connection) throws ReflectiveOperationException { + Field field = AutoSwitchHAConnection.class.getDeclaredField("readSocketService"); + field.setAccessible(true); + return (AutoSwitchHAConnection.ReadSocketService) field.get(connection); + } +}