Fix the auth data is NPE error (#1531)

streamnative · Nov 14, 2024 · 7d76b42 · 7d76b42
1 parent e78bcfa
commit 7d76b42
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 0 deletions.
diff --git a/...oker/src/main/java/io/streamnative/pulsar/handlers/mqtt/broker/channel/MQTTServerCnx.java b/...oker/src/main/java/io/streamnative/pulsar/handlers/mqtt/broker/channel/MQTTServerCnx.java
@@ -13,12 +13,14 @@
  */
 package io.streamnative.pulsar.handlers.mqtt.broker.channel;
 
+import static io.streamnative.pulsar.handlers.mqtt.common.Constants.AUTH_DATA_ATTRIBUTE_KEY;
 import io.netty.channel.ChannelHandlerContext;
 import io.streamnative.pulsar.handlers.mqtt.broker.impl.consumer.MQTTConsumer;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.pulsar.broker.PulsarService;
+import org.apache.pulsar.broker.authentication.AuthenticationDataSource;
 import org.apache.pulsar.broker.loadbalance.extensions.data.BrokerLookupData;
 import org.apache.pulsar.broker.service.Consumer;
 import org.apache.pulsar.broker.service.ServerCnx;
@@ -65,4 +67,9 @@ private void safelyRemoveConsumer(Consumer consumer) {
             });
         }
     }
+
+    @Override
+    public AuthenticationDataSource getAuthenticationData() {
+        return ctx.channel().attr(AUTH_DATA_ATTRIBUTE_KEY).get();
+    }
 }
diff --git a/mqtt-common/src/main/java/io/streamnative/pulsar/handlers/mqtt/common/Connection.java b/mqtt-common/src/main/java/io/streamnative/pulsar/handlers/mqtt/common/Connection.java
@@ -16,6 +16,7 @@
 import static io.streamnative.pulsar.handlers.mqtt.common.Connection.ConnectionState.CONNECT_ACK;
 import static io.streamnative.pulsar.handlers.mqtt.common.Connection.ConnectionState.DISCONNECTED;
 import static io.streamnative.pulsar.handlers.mqtt.common.Connection.ConnectionState.ESTABLISHED;
+import static io.streamnative.pulsar.handlers.mqtt.common.Constants.AUTH_DATA_ATTRIBUTE_KEY;
 import static io.streamnative.pulsar.handlers.mqtt.common.utils.MqttMessageUtils.getAuthMethod;
 import static io.streamnative.pulsar.handlers.mqtt.common.utils.NettyUtils.ATTR_KEY_CONNECTION;
 import static java.util.concurrent.atomic.AtomicReferenceFieldUpdater.newUpdater;
@@ -101,6 +102,7 @@ public class Connection {
 
     static ChannelException channelInactiveException = new ChannelException("Channel is inactive");
 
+
     Connection(ConnectionBuilder builder) {
         this.clientId = builder.clientId;
         this.protocolVersion = builder.protocolVersion;
@@ -116,6 +118,7 @@ public class Connection {
         this.processor = builder.processor;
         this.fromProxy = builder.fromProxy;
         this.authData = builder.authData;
+        this.channel.attr(AUTH_DATA_ATTRIBUTE_KEY).set(authData);
         this.addIdleStateHandler();
         this.manager.addConnection(this);
         this.topicAliasManager = new TopicAliasManager(clientRestrictions.getTopicAliasMaximum());
@@ -162,6 +165,7 @@ public CompletableFuture<Void> sendAck(MqttAck mqttAck) {
 
     public void updateAuthData(AuthenticationDataSource authData) {
         this.authData = authData;
+        this.channel.attr(AUTH_DATA_ATTRIBUTE_KEY).set(authData);
     }
 
     public CompletableFuture<Void> sendAckThenClose(MqttAck mqttAck) {

diff --git a/mqtt-common/src/main/java/io/streamnative/pulsar/handlers/mqtt/common/Constants.java b/mqtt-common/src/main/java/io/streamnative/pulsar/handlers/mqtt/common/Constants.java
@@ -13,6 +13,9 @@
  */
 package io.streamnative.pulsar.handlers.mqtt.common;
 
+import io.netty.util.AttributeKey;
+import org.apache.pulsar.broker.authentication.AuthenticationDataSource;
+
 /**
  * Server constants keeper.
  */
@@ -47,6 +50,9 @@ public final class Constants {
 
     public static final String MQTT_SUB_PROTOCOL_CSV_LIST = "mqtt, mqttv3.1, mqttv3.1.1, mqttv5.0";
 
+    public static final AttributeKey<AuthenticationDataSource> AUTH_DATA_ATTRIBUTE_KEY =
+            AttributeKey.valueOf("authData");
+
     private Constants() {
     }
 }