From 92842e0f0e8c1a1a800bf1eadae1a71d6164cb04 Mon Sep 17 00:00:00 2001
From: Andrii Shvaika <ashvayka@thingsboard.io>
Date: Thu, 3 Feb 2022 17:59:55 +0200
Subject: [PATCH] Improved IP Filter to delete next handlers from pipeline

---
 .../transport/mqtt/MqttTransportHandler.java  | 13 ++++++++--
 .../transport/mqtt/limits/IpFilter.java       |  2 ++
 .../transport/mqtt/limits/ProxyIpFilter.java  | 24 ++++++++++++-------
 3 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/MqttTransportHandler.java b/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/MqttTransportHandler.java
index 75ec21a8df..a07982d718 100644
--- a/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/MqttTransportHandler.java
+++ b/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/MqttTransportHandler.java
@@ -188,7 +188,16 @@ public class MqttTransportHandler extends ChannelInboundHandlerAdapter implement
     }
 
     InetSocketAddress getAddress(ChannelHandlerContext ctx) {
-        return ctx.channel().attr(MqttTransportService.ADDRESS).get();
+        var address = ctx.channel().attr(MqttTransportService.ADDRESS).get();
+        if (address == null) {
+            log.trace("[{}] Received empty address.", ctx.channel().id());
+            InetSocketAddress remoteAddress = (InetSocketAddress) ctx.channel().remoteAddress();
+            log.trace("[{}] Going to use address: {}", ctx.channel().id(), remoteAddress);
+            return remoteAddress;
+        } else {
+            log.trace("[{}] Received address: {}", ctx.channel().id(), address);
+        }
+        return address;
     }
 
     void processMqttMsg(ChannelHandlerContext ctx, MqttMessage msg) {
@@ -757,7 +766,7 @@ public class MqttTransportHandler extends ChannelInboundHandlerAdapter implement
     }
 
     void processConnect(ChannelHandlerContext ctx, MqttConnectMessage msg) {
-        log.debug("[{}] Processing connect msg for client: {}!", sessionId, msg.payload().clientIdentifier());
+        log.debug("[{}][{}] Processing connect msg for client: {}!", address, sessionId, msg.payload().clientIdentifier());
         String userName = msg.payload().userName();
         String clientId = msg.payload().clientIdentifier();
         if (DataConstants.PROVISION.equals(userName) || DataConstants.PROVISION.equals(clientId)) {
diff --git a/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/limits/IpFilter.java b/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/limits/IpFilter.java
index 05da48b0dd..479f38fd8c 100644
--- a/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/limits/IpFilter.java
+++ b/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/limits/IpFilter.java
@@ -36,7 +36,9 @@ public class IpFilter extends AbstractRemoteAddressFilter<InetSocketAddress> {
 
     @Override
     protected boolean accept(ChannelHandlerContext ctx, InetSocketAddress remoteAddress) throws Exception {
+        log.trace("[{}] Received msg: {}", ctx.channel().id(), remoteAddress);
         if(context.checkAddress(remoteAddress)){
+            log.trace("[{}] Setting address: {}", ctx.channel().id(), remoteAddress);
             ctx.channel().attr(MqttTransportService.ADDRESS).set(remoteAddress);
             return true;
         } else {
diff --git a/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/limits/ProxyIpFilter.java b/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/limits/ProxyIpFilter.java
index 87e92bb5b3..c7a70c6d44 100644
--- a/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/limits/ProxyIpFilter.java
+++ b/common/transport/mqtt/src/main/java/org/thingsboard/server/transport/mqtt/limits/ProxyIpFilter.java
@@ -38,23 +38,31 @@ public class ProxyIpFilter extends ChannelInboundHandlerAdapter {
 
     @Override
     public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
-        if(msg instanceof HAProxyMessage){
+        log.trace("[{}] Received msg: {}", ctx.channel().id(), msg);
+        if (msg instanceof HAProxyMessage) {
             HAProxyMessage proxyMsg = (HAProxyMessage) msg;
-            if(proxyMsg.sourceAddress() != null && proxyMsg.sourcePort() > 0) {
+            if (proxyMsg.sourceAddress() != null && proxyMsg.sourcePort() > 0) {
                 InetSocketAddress address = new InetSocketAddress(proxyMsg.sourceAddress(), proxyMsg.sourcePort());
-                if(!context.checkAddress(address)){
-                    ctx.close();
+                if (!context.checkAddress(address)) {
+                    closeChannel(ctx);
                 } else {
+                    log.trace("[{}] Setting address: {}", ctx.channel().id(), address);
                     ctx.channel().attr(MqttTransportService.ADDRESS).set(address);
                     // We no longer need this channel in the pipeline. Similar to HAProxyMessageDecoder
                     ctx.pipeline().remove(this);
                 }
             } else {
-                log.debug("Received local health-check connection message: {}", proxyMsg);
-                ctx.close();
+                log.trace("Received local health-check connection message: {}", proxyMsg);
+                closeChannel(ctx);
             }
-        } else {
-            super.channelRead(ctx, msg);
         }
     }
+
+    private void closeChannel(ChannelHandlerContext ctx) {
+        while (ctx.pipeline().last() != this) {
+            ctx.pipeline().removeLast();
+        }
+        ctx.pipeline().remove(this);
+        ctx.close();
+    }
 }