[broker-5] add validating UTF-8 string during packet reading

Author: JavaSaBr

src/main/java/com/ss/mqtt/broker/config/MqttBrokerConfig.java

@@ -13,7 +13,6 @@
 import com.ss.mqtt.broker.service.SubscriptionService;
 import com.ss.mqtt.broker.service.impl.*;
 import com.ss.rlib.network.*;
-import com.ss.mqtt.broker.service.impl.SimplePublishingService;
 import com.ss.rlib.network.impl.DefaultBufferAllocator;
 import com.ss.rlib.network.server.ServerNetwork;
 import lombok.RequiredArgsConstructor;
@@ -81,28 +80,26 @@ private interface ChannelFactory extends
     }
 
     @Bean
-    @NotNull ServerNetwork<MqttConnection> deviceNetwork(
+    @NotNull ServerNetwork<@NotNull MqttConnection> deviceNetwork(
         @NotNull ServerNetworkConfig networkConfig,
         @NotNull BufferAllocator bufferAllocator,
         @NotNull MqttConnectionConfig deviceConnectionConfig,
-        @NotNull SubscriptionService subscriptionService,
         PacketInHandler @NotNull [] devicePacketHandlers
     ) {
         return NetworkFactory.newServerNetwork(
             networkConfig,
             deviceConnectionFactory(
                 bufferAllocator,
                 deviceConnectionConfig,
-                subscriptionService,
                 devicePacketHandlers
             )
         );
     }
 
     @Bean
     @NotNull InetSocketAddress deviceNetworkAddress(
-        @NotNull ServerNetwork<MqttConnection> deviceNetwork,
-        @NotNull Consumer<MqttConnection> mqttConnectionConsumer
+        @NotNull ServerNetwork<@NotNull MqttConnection> deviceNetwork,
+        @NotNull Consumer<@NotNull MqttConnection> mqttConnectionConsumer
     ) {
 
         var address = new InetSocketAddress("localhost", 1883);
@@ -124,7 +121,7 @@ private interface ChannelFactory extends
     }
 
     @Bean
-    @NotNull Consumer<MqttConnection> mqttConnectionConsumer() {
+    @NotNull Consumer<@NotNull MqttConnection> mqttConnectionConsumer() {
         return mqttConnection -> {
             log.info("Accepted connection: {}", mqttConnection);
             var client = (UnsafeMqttClient) mqttConnection.getClient();
@@ -167,7 +164,6 @@ private interface ChannelFactory extends
     private @NotNull ChannelFactory deviceConnectionFactory(
         @NotNull BufferAllocator bufferAllocator,
         @NotNull MqttConnectionConfig connectionConfig,
-        @NotNull SubscriptionService subscriptionService,
         PacketInHandler @NotNull [] packetHandlers
     ) {
         return (network, channel) -> new MqttConnection(
@@ -176,7 +172,6 @@ private interface ChannelFactory extends
             bufferAllocator,
             100,
             packetHandlers,
-            subscriptionService,
             connectionConfig,
             DeviceMqttClient::new
         );

src/main/java/com/ss/mqtt/broker/network/MqttConnection.java

@@ -9,14 +9,14 @@
 import com.ss.mqtt.broker.network.packet.in.MqttReadablePacket;
 import com.ss.mqtt.broker.network.packet.in.handler.PacketInHandler;
 import com.ss.mqtt.broker.network.packet.out.MqttWritablePacket;
-import com.ss.mqtt.broker.service.SubscriptionService;
 import com.ss.rlib.network.BufferAllocator;
 import com.ss.rlib.network.Connection;
 import com.ss.rlib.network.Network;
 import com.ss.rlib.network.NetworkCryptor;
 import com.ss.rlib.network.impl.AbstractConnection;
 import com.ss.rlib.network.packet.PacketReader;
 import com.ss.rlib.network.packet.PacketWriter;
+import lombok.AccessLevel;
 import lombok.Getter;
 import lombok.Setter;
 import lombok.extern.log4j.Log4j2;
@@ -26,27 +26,27 @@
 import java.util.function.Function;
 
 @Log4j2
-@Getter
 public class MqttConnection extends AbstractConnection<MqttReadablePacket, MqttWritablePacket> {
 
+    @Getter(AccessLevel.PROTECTED)
     private final @NotNull PacketReader packetReader;
+
+    @Getter(AccessLevel.PROTECTED)
     private final @NotNull PacketWriter packetWriter;
 
-    private final PacketInHandler @NotNull [] packetHandlers;
-    private final @NotNull SubscriptionService subscriptionService;
+    private final @Getter PacketInHandler @NotNull [] packetHandlers;
 
     private final @Getter @NotNull MqttClient client;
     private final @Getter @NotNull MqttConnectionConfig config;
 
-    private volatile @Setter @NotNull MqttVersion mqttVersion;
+    private volatile @Getter @Setter @NotNull MqttVersion mqttVersion;
 
     public MqttConnection(
         @NotNull Network<? extends Connection<MqttReadablePacket, MqttWritablePacket>> network,
         @NotNull AsynchronousSocketChannel channel,
         @NotNull BufferAllocator bufferAllocator,
         int maxPacketsByRead,
         PacketInHandler @NotNull [] packetHandlers,
-        @NotNull SubscriptionService subscriptionService,
         @NotNull MqttConnectionConfig config,
         @NotNull Function<MqttConnection, UnsafeMqttClient> clientFactory
     ) {
@@ -56,7 +56,6 @@ public MqttConnection(
         this.mqttVersion = MqttVersion.MQTT_3_1_1;
         this.packetReader = createPacketReader();
         this.packetWriter = createPacketWriter();
-        this.subscriptionService = subscriptionService;
         this.client = clientFactory.apply(this);
     }
 

src/main/java/com/ss/mqtt/broker/network/client/impl/AbstractMqttClient.java

@@ -15,19 +15,20 @@
 import lombok.extern.log4j.Log4j2;
 import org.jetbrains.annotations.NotNull;
 
+@Getter
 @Log4j2
 @EqualsAndHashCode(of = "clientId")
 public abstract class AbstractMqttClient implements UnsafeMqttClient {
 
-    protected final @Getter@NotNull MqttConnection connection;
+    protected final @NotNull MqttConnection connection;
 
-    private volatile @Setter @Getter @NotNull String clientId;
+    private volatile @Setter @NotNull String clientId;
 
-    private volatile @Getter long sessionExpiryInterval = MqttPropertyConstants.SESSION_EXPIRY_INTERVAL_DEFAULT;
-    private volatile @Getter int receiveMax = MqttPropertyConstants.RECEIVE_MAXIMUM_DEFAULT;
-    private volatile @Getter int maximumPacketSize = MqttPropertyConstants.MAXIMUM_PACKET_SIZE_DEFAULT;
-    private volatile @Getter int topicAliasMaximum = MqttPropertyConstants.TOPIC_ALIAS_MAXIMUM_DEFAULT;
-    private volatile @Getter int keepAlive = MqttPropertyConstants.SERVER_KEEP_ALIVE_MAX;
+    private volatile long sessionExpiryInterval = MqttPropertyConstants.SESSION_EXPIRY_INTERVAL_DEFAULT;
+    private volatile int receiveMax = MqttPropertyConstants.RECEIVE_MAXIMUM_DEFAULT;
+    private volatile int maximumPacketSize = MqttPropertyConstants.MAXIMUM_PACKET_SIZE_DEFAULT;
+    private volatile int topicAliasMaximum = MqttPropertyConstants.TOPIC_ALIAS_MAXIMUM_DEFAULT;
+    private volatile int keepAlive = MqttPropertyConstants.SERVER_KEEP_ALIVE_MAX;
 
     public AbstractMqttClient(@NotNull MqttConnection connection) {
         this.connection = connection;

src/main/java/com/ss/mqtt/broker/network/packet/in/MqttReadablePacket.java

@@ -13,16 +13,38 @@
 import com.ss.rlib.common.util.array.ArrayFactory;
 import com.ss.rlib.network.packet.impl.AbstractReadablePacket;
 import lombok.Getter;
+import lombok.RequiredArgsConstructor;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.nio.ByteBuffer;
+import java.nio.CharBuffer;
+import java.nio.charset.CharsetDecoder;
+import java.nio.charset.CodingErrorAction;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.Set;
 
 public abstract class MqttReadablePacket extends AbstractReadablePacket<MqttConnection> {
 
+    @Getter
+    @RequiredArgsConstructor
+    private static class Utf8Decoder {
+        private final CharsetDecoder decoder;
+        private final ByteBuffer inBuffer;
+        private final CharBuffer outBuffer;
+    }
+
+    private static final ThreadLocal<Utf8Decoder> LOCAL_DECODER = ThreadLocal.withInitial(() -> {
+
+        var decoder = StandardCharsets.UTF_8.newDecoder()
+            .onMalformedInput(CodingErrorAction.REPORT)
+            .onUnmappableCharacter(CodingErrorAction.REPORT);
+
+        return new Utf8Decoder(decoder, ByteBuffer.allocate(10240), CharBuffer.allocate(10240));
+    });
+
+
     /**
      * The list of user properties.
      */
@@ -161,9 +183,30 @@ protected long readUnsignedInt(@NotNull ByteBuffer buffer) {
 
     @Override
     protected @NotNull String readString(@NotNull ByteBuffer buffer) {
-        var stringData = new byte[readShort(buffer) & 0xFFFF];
-        buffer.get(stringData);
-        return new String(stringData, StandardCharsets.UTF_8);
+
+        var utf8Decoder = LOCAL_DECODER.get();
+        var inBuffer = utf8Decoder.getInBuffer();
+
+        var stringLength = readShort(buffer) & 0xFFFF;
+
+        if (stringLength > inBuffer.capacity()) {
+            throw new ConnectionRejectException(ConnectAckReasonCode.MALFORMED_PACKET);
+        }
+
+        var decoder = utf8Decoder.getDecoder();
+        var outBuffer = utf8Decoder.getOutBuffer();
+
+        buffer.get(inBuffer.clear().array(), 0, stringLength);
+
+        decoder.reset();
+
+        var result = decoder.decode(inBuffer.position(stringLength).flip(), outBuffer.clear(), true);
+
+        if (result.isError()) {
+            throw new ConnectionRejectException(ConnectAckReasonCode.MALFORMED_PACKET);
+        }
+
+        return new String(inBuffer.array(), 0, stringLength, StandardCharsets.UTF_8);
     }
 
     protected @NotNull byte[] readBytes(@NotNull ByteBuffer buffer) {

src/main/java/com/ss/mqtt/broker/network/packet/in/handler/ConnectInPacketHandler.java

@@ -9,7 +9,6 @@
 import com.ss.rlib.common.util.StringUtils;
 import lombok.RequiredArgsConstructor;
 import org.jetbrains.annotations.NotNull;
-import reactor.core.publisher.Mono;
 
 @RequiredArgsConstructor
 public class ConnectInPacketHandler extends AbstractPacketHandler<UnsafeMqttClient, ConnectInPacket> {
@@ -38,12 +37,6 @@ protected void handleImpl(@NotNull UnsafeMqttClient client, @NotNull ConnectInPa
                 .switchIfEmpty(fromRunnable(() -> client.reject(ConnectAckReasonCode.CLIENT_IDENTIFIER_NOT_VALID)))
                 .subscribe();
         } else {
-
-            if (!clientIdRegistry.validate(requestedClientId)) {
-                client.reject(ConnectAckReasonCode.CLIENT_IDENTIFIER_NOT_VALID);
-                return;
-            }
-
             clientIdRegistry.register(requestedClientId)
                 .subscribe(result -> {
                     if(!result) {

src/main/java/com/ss/mqtt/broker/network/packet/in/handler/PublishInPacketHandler.java

@@ -9,7 +9,7 @@
 @RequiredArgsConstructor
 public class PublishInPacketHandler extends AbstractPacketHandler<UnsafeMqttClient, PublishInPacket> {
 
-    private final PublishingService publishingService;
+    private final @NotNull PublishingService publishingService;
 
     @Override
     protected void handleImpl(@NotNull UnsafeMqttClient client, @NotNull PublishInPacket packet) {

src/main/java/com/ss/mqtt/broker/network/packet/in/handler/SubscribeInPacketHandler.java

@@ -9,7 +9,7 @@
 @RequiredArgsConstructor
 public class SubscribeInPacketHandler extends AbstractPacketHandler<UnsafeMqttClient, SubscribeInPacket> {
 
-    private final SubscriptionService subscriptionService;
+    private final @NotNull SubscriptionService subscriptionService;
 
     @Override
     protected void handleImpl(@NotNull UnsafeMqttClient client, @NotNull SubscribeInPacket packet) {

src/main/java/com/ss/mqtt/broker/network/packet/in/handler/UnsubscribeInPacketHandler.java

@@ -9,7 +9,7 @@
 @RequiredArgsConstructor
 public class UnsubscribeInPacketHandler extends AbstractPacketHandler<UnsafeMqttClient, UnsubscribeInPacket> {
 
-    private final SubscriptionService subscriptionService;
+    private final @NotNull SubscriptionService subscriptionService;
 
     @Override
     protected void handleImpl(@NotNull UnsafeMqttClient client, @NotNull UnsubscribeInPacket packet) {

src/main/java/com/ss/mqtt/broker/service/impl/SimpleClientIdRegistry.java

@@ -32,23 +32,22 @@ public SimpleClientIdRegistry(@NotNull String availableChars, int maxClientIdLen
     @Override
     public @NotNull Mono<Boolean> register(@NotNull String clientId) {
 
-        var value = clientIdRegistry.getInReadLock(clientId, ObjectDictionary::get);
+        if (!validate(clientId)) {
+            return Mono.just(Boolean.FALSE);
+        }
 
+        var value = clientIdRegistry.getInReadLock(clientId, ObjectDictionary::get);
         if (value != null) {
             return Mono.just(Boolean.FALSE);
         }
 
         var stamp = clientIdRegistry.writeLock();
         try {
-
             value = clientIdRegistry.get(clientId);
-
             if (value != null) {
                 return Mono.just(Boolean.FALSE);
             }
-
             clientIdRegistry.put(clientId, CLIENT_ID_VALUE);
-
         } finally {
             clientIdRegistry.writeUnlock(stamp);
         }

src/main/java/com/ss/mqtt/broker/service/impl/SimpleSubscriptions.java

@@ -48,14 +48,11 @@ public class SimpleSubscriptions implements Subscriptions {
         @NotNull String topicFilter,
         @NotNull MqttClient mqttClient
     ) {
-
         var subscribers = subscriptions.getOrDefault(topicFilter, Array.empty());
-
         if (subscribers.removeIf(subscriber -> mqttClient.equals(subscriber.getMqttClient()))) {
             return UnsubscribeAckReasonCode.SUCCESS;
+        } else {
+            return UnsubscribeAckReasonCode.NO_SUBSCRIPTION_EXISTED;
         }
-
-        return UnsubscribeAckReasonCode.NO_SUBSCRIPTION_EXISTED;
     }
-
 }