add mechanism to handle exceptions on packet level

Author: JavaSaBr

src/main/java/com/ss/mqtt/broker/exception/ConnectionRejectException.java

@@ -1,11 +1,19 @@
 package com.ss.mqtt.broker.exception;
 
 import com.ss.mqtt.broker.model.ConnectAckReasonCode;
-import lombok.RequiredArgsConstructor;
+import lombok.Getter;
 import org.jetbrains.annotations.NotNull;
 
-@RequiredArgsConstructor
-public class ConnectionRejectException extends RuntimeException {
+public class ConnectionRejectException extends MqttException {
 
-    private final @NotNull ConnectAckReasonCode reasonCode;
+    private final @Getter @NotNull ConnectAckReasonCode reasonCode;
+
+    public ConnectionRejectException(@NotNull ConnectAckReasonCode reasonCode) {
+        this.reasonCode = reasonCode;
+    }
+
+    public ConnectionRejectException(@NotNull Throwable cause, @NotNull ConnectAckReasonCode reasonCode) {
+        super(cause);
+        this.reasonCode = reasonCode;
+    }
 }

src/main/java/com/ss/mqtt/broker/exception/MqttException.java

@@ -0,0 +1,17 @@
+package com.ss.mqtt.broker.exception;
+
+import org.jetbrains.annotations.NotNull;
+
+public class MqttException extends RuntimeException {
+
+    public MqttException() {
+    }
+
+    public MqttException(@NotNull String message) {
+        super(message);
+    }
+
+    public MqttException(@NotNull Throwable cause) {
+        super(cause);
+    }
+}

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

@@ -1,5 +1,6 @@
 package com.ss.mqtt.broker.network;
 
+import com.ss.mqtt.broker.exception.ConnectionRejectException;
 import com.ss.mqtt.broker.model.ConnectAckReasonCode;
 import com.ss.mqtt.broker.model.MqttPropertyConstants;
 import com.ss.mqtt.broker.network.packet.PacketType;
@@ -55,6 +56,19 @@ public void reject(@NotNull ConnectAckReasonCode returnCode) {
     private void onConnected(@NotNull ConnectInPacket connect) {
         connection.setMqttVersion(connect.getMqttVersion());
 
+        var exception = connect.getException();
+
+        if (exception instanceof ConnectionRejectException) {
+
+            var reasonCode = ((ConnectionRejectException) exception).getReasonCode();
+
+            connection
+                .sendWithFeedback(getPacketOutFactory().newConnectAck(this, reasonCode))
+                .thenAccept(sent -> connection.close());
+
+            return;
+        }
+
         sessionExpiryInterval = connect.getSessionExpiryInterval();
         receiveMax = connect.getReceiveMax();
         maximumPacketSize = connect.getMaximumPacketSize();

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

@@ -2,7 +2,6 @@
 
 import com.ss.mqtt.broker.config.MqttConnectionConfig;
 import com.ss.mqtt.broker.model.MqttVersion;
-import com.ss.mqtt.broker.model.QoS;
 import com.ss.mqtt.broker.network.packet.MqttPacketReader;
 import com.ss.mqtt.broker.network.packet.MqttPacketWriter;
 import com.ss.mqtt.broker.network.packet.in.MqttReadablePacket;
@@ -76,7 +75,9 @@ public boolean isSupported(@NotNull MqttVersion mqttVersion) {
             channel,
             bufferAllocator,
             this::updateLastActivity,
-            this::nextPacketToWrite
+            this::nextPacketToWrite,
+            this::onWrittenPacket,
+            this::onSentPacket
         );
     }
 

src/main/java/com/ss/mqtt/broker/network/packet/MqttPacketWriter.java

@@ -4,11 +4,14 @@
 import com.ss.mqtt.broker.network.packet.out.MqttWritablePacket;
 import com.ss.mqtt.broker.util.MqttDataUtils;
 import com.ss.rlib.network.BufferAllocator;
+import com.ss.rlib.network.packet.WritablePacket;
 import com.ss.rlib.network.packet.impl.AbstractPacketWriter;
 import org.jetbrains.annotations.NotNull;
 
 import java.nio.ByteBuffer;
 import java.nio.channels.AsynchronousSocketChannel;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 public class MqttPacketWriter extends AbstractPacketWriter<MqttWritablePacket, MqttConnection> {
@@ -18,13 +21,23 @@ public MqttPacketWriter(
         @NotNull AsynchronousSocketChannel channel,
         @NotNull BufferAllocator bufferAllocator,
         @NotNull Runnable updateActivityFunction,
-        @NotNull Supplier<MqttWritablePacket> nextWritePacketSupplier
+        @NotNull Supplier<@NotNull WritablePacket> nextWritePacketSupplier,
+        @NotNull Consumer<@NotNull WritablePacket> writtenPacketHandler,
+        @NotNull BiConsumer<@NotNull WritablePacket, Boolean> sentPacketHandler
     ) {
-        super(connection, channel, bufferAllocator, updateActivityFunction, nextWritePacketSupplier);
+        super(
+            connection,
+            channel,
+            bufferAllocator,
+            updateActivityFunction,
+            nextWritePacketSupplier,
+            writtenPacketHandler,
+            sentPacketHandler
+        );
     }
 
     @Override
-    protected int getTotalSize(@NotNull MqttWritablePacket packet, int expectedLength) {
+    protected int getTotalSize(@NotNull WritablePacket packet, int expectedLength) {
         return 1 + MqttDataUtils.sizeOfMbi(expectedLength) + expectedLength;
     }
 

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

@@ -219,7 +219,7 @@ public class ConnectInPacket extends MqttReadablePacket {
     public ConnectInPacket(byte info) {
         super(info);
         this.userProperties = Array.empty();
-        this.mqttVersion = MqttVersion.MQTT_5;
+        this.mqttVersion = MqttVersion.MQTT_3_1_1;
         this.clientId = StringUtils.EMPTY;
         this.willTopic = StringUtils.EMPTY;
         this.username = StringUtils.EMPTY;
@@ -270,7 +270,7 @@ protected void readVariableHeader(@NotNull MqttConnection connection, @NotNull B
                  the reserved flag is not 0 it is a Malformed Packet. Refer to section 4.13 for information about handling
                  errors.
                 */
-                throw new IllegalStateException("non-zero reserved flag");
+                throw new ConnectionRejectException(ConnectAckReasonCode.MALFORMED_PACKET);
             }
         }
 

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

@@ -1,5 +1,8 @@
 package com.ss.mqtt.broker.network.packet.in;
 
+import com.ss.mqtt.broker.exception.ConnectionRejectException;
+import com.ss.mqtt.broker.exception.MqttException;
+import com.ss.mqtt.broker.model.ConnectAckReasonCode;
 import com.ss.mqtt.broker.model.MqttVersion;
 import com.ss.mqtt.broker.model.PacketProperty;
 import com.ss.mqtt.broker.model.StringPair;
@@ -11,6 +14,7 @@
 import com.ss.rlib.network.packet.impl.AbstractReadablePacket;
 import lombok.Getter;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
@@ -24,6 +28,11 @@ public abstract class MqttReadablePacket extends AbstractReadablePacket<MqttConn
      */
     protected @Getter @NotNull Array<StringPair> userProperties;
 
+    /**
+     * The happened exception during parsing this packet.
+     */
+    protected @Getter @Nullable Exception exception;
+
     protected MqttReadablePacket(byte info) {
         this.userProperties = Array.empty();
     }
@@ -41,6 +50,17 @@ protected void readImpl(@NotNull MqttConnection connection, @NotNull ByteBuffer
         readPayload(connection, buffer);
     }
 
+    @Override
+    protected void handleException(@NotNull ByteBuffer buffer, @NotNull Exception exception) {
+        super.handleException(buffer, exception);
+
+        if (!(exception instanceof MqttException)) {
+            exception = new ConnectionRejectException(exception, ConnectAckReasonCode.PROTOCOL_ERROR);
+        }
+
+        this.exception = exception;
+    }
+
     protected boolean isPropertiesSupported(@NotNull MqttConnection connection, @NotNull ByteBuffer buffer) {
         return connection.isSupported(MqttVersion.MQTT_5);
     }