diff --git a/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/MQTTConnectHandler.java b/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/MQTTConnectHandler.java index 25342e408..5c09056b7 100644 --- a/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/MQTTConnectHandler.java +++ b/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/MQTTConnectHandler.java @@ -158,6 +158,26 @@ public final void channelRead(ChannelHandlerContext ctx, Object msg) { LWT willMessage = connMsg.variableHeader().isWillFlag() ? getWillMessage(connMsg, clientInfo) : null; + CompletableFuture willPermission = willMessage == null + ? CompletableFuture.completedFuture(AuthResult.ok(okOrGoAway.successInfo)) + : checkWillPermission(connMsg, willMessage, okOrGoAway.successInfo); + return willPermission.thenApply( + permissionResult -> new SessionPreparation(permissionResult, settings, willMessage)); + } + }, ctx.executor()) + .thenComposeAsync(sessionPreparation -> { + if (sessionPreparation == null) { + return CompletableFuture.completedFuture(null); + } + AuthResult okOrGoAway = sessionPreparation.permissionResult; + if (okOrGoAway.goAway != null) { + handleGoAway(okOrGoAway.goAway); + return CompletableFuture.completedFuture(null); + } else { + ClientInfo clientInfo = okOrGoAway.successInfo.clientInfo; + TenantSettings settings = sessionPreparation.settings; + LWT willMessage = sessionPreparation.willMessage; + int keepAliveSeconds = keepAliveSeconds(connMsg.variableHeader().keepAliveTimeSeconds(), settings); String userSessionId = userSessionId(clientInfo); @@ -449,6 +469,10 @@ private CompletableFuture expireInbox(long reqId, protected abstract CompletableFuture checkConnectPermission(MqttConnectMessage message, SuccessInfo successInfo); + protected abstract CompletableFuture checkWillPermission(MqttConnectMessage message, + LWT willMessage, + SuccessInfo successInfo); + protected abstract void handleMqttMessage(MqttMessage message); protected abstract GoAway onNoEnoughResources(MqttConnectMessage message, TenantResourceType resourceType, @@ -661,6 +685,11 @@ private enum ExpireResult { ERROR } + private record SessionPreparation(AuthResult permissionResult, + TenantSettings settings, + LWT willMessage) { + } + public record SuccessInfo(ClientInfo clientInfo, Optional responseInfo, // mqtt5 Optional authData, // mqtt5 diff --git a/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/MQTTSessionHandler.java b/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/MQTTSessionHandler.java index f464af894..97c859075 100644 --- a/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/MQTTSessionHandler.java +++ b/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/MQTTSessionHandler.java @@ -403,7 +403,7 @@ public final void channelInactive(ChannelHandlerContext ctx) { resendTask.cancel(true); } if (noDelayLWT != null) { - addBgTask(pubWillMessage(noDelayLWT)); + addBgTask(doPubLastWill(noDelayLWT)); } cancelStallTask(); Sets.newHashSet(fgTasks).forEach(t -> t.cancel(true)); @@ -1513,27 +1513,6 @@ private CompletableFuture handleQoS2Pub(long reqId, }); } - private CompletableFuture pubWillMessage(LWT willMessage) { - return authProvider.checkPermission(clientInfo(), buildPubAction(willMessage.getTopic(), - willMessage.getMessage() - .getPubQoS(), - willMessage.getMessage().getIsRetain())) - .thenCompose(checkResult -> { - assert ctx.executor().inEventLoop(); - if (checkResult.hasGranted()) { - return doPubLastWill(willMessage); - } else { - sessionCtx.eventCollector.report(getLocal(PubActionDisallow.class) - .isLastWill(true) - .topic(willMessage.getTopic()) - .qos(willMessage.getMessage().getPubQoS()) - .isRetain(willMessage.getMessage().getIsRetain()) - .clientInfo(clientInfo)); - return CompletableFuture.completedFuture(null); - } - }); - } - private void checkIdle() { if (sessionCtx.nanoTime() - lastActiveAtNanos > idleTimeoutNanos) { idleTimeoutTask.cancel(true); diff --git a/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/v3/MQTT3ConnectHandler.java b/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/v3/MQTT3ConnectHandler.java index 2f407daae..603a8143d 100644 --- a/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/v3/MQTT3ConnectHandler.java +++ b/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/v3/MQTT3ConnectHandler.java @@ -30,6 +30,7 @@ import static org.apache.bifromq.mqtt.handler.condition.ORCondition.or; import static org.apache.bifromq.mqtt.handler.v3.MQTT3MessageUtils.toWillMessage; import static org.apache.bifromq.mqtt.utils.AuthUtil.buildConnAction; +import static org.apache.bifromq.mqtt.utils.AuthUtil.buildPubAction; import static org.apache.bifromq.plugin.eventcollector.ThreadLocalEventPool.getLocal; import static org.apache.bifromq.type.MQTTClientInfoConstants.MQTT_CHANNEL_ID_KEY; import static org.apache.bifromq.type.MQTTClientInfoConstants.MQTT_CLIENT_ADDRESS_KEY; @@ -72,6 +73,7 @@ import org.apache.bifromq.plugin.clientbalancer.IClientBalancer; import org.apache.bifromq.plugin.clientbalancer.Redirection; import org.apache.bifromq.plugin.eventcollector.OutOfTenantResource; +import org.apache.bifromq.plugin.eventcollector.mqttbroker.accessctrl.PubActionDisallow; import org.apache.bifromq.plugin.eventcollector.mqttbroker.channelclosed.AuthError; import org.apache.bifromq.plugin.eventcollector.mqttbroker.channelclosed.IdentifierRejected; import org.apache.bifromq.plugin.eventcollector.mqttbroker.channelclosed.MalformedClientIdentifier; @@ -258,6 +260,49 @@ protected CompletableFuture checkConnectPermission(MqttConnectMessag }); } + @Override + protected CompletableFuture checkWillPermission(MqttConnectMessage message, + LWT willMessage, + SuccessInfo successInfo) { + ClientInfo clientInfo = successInfo.clientInfo(); + return authProvider.checkPermission(clientInfo, buildPubAction(willMessage.getTopic(), + willMessage.getMessage().getPubQoS(), willMessage.getMessage().getIsRetain())) + .handle((checkResult, e) -> { + if (e != null) { + return goAway(MqttMessageBuilders.connAck() + .returnCode(CONNECTION_REFUSED_SERVER_UNAVAILABLE) + .build(), + getLocal(AuthError.class) + .cause("Failed to check Will publish permission") + .peerAddress(ChannelAttrs.socketAddress(ctx.channel()))); + } + switch (checkResult.getTypeCase()) { + case GRANTED -> { + return AuthResult.ok(successInfo); + } + case DENIED -> { + return goAway(MqttMessageBuilders.connAck() + .returnCode(CONNECTION_REFUSED_NOT_AUTHORIZED) + .build(), + getLocal(PubActionDisallow.class) + .isLastWill(true) + .topic(willMessage.getTopic()) + .qos(willMessage.getMessage().getPubQoS()) + .isRetain(willMessage.getMessage().getIsRetain()) + .clientInfo(clientInfo)); + } + default -> { + return goAway(MqttMessageBuilders.connAck() + .returnCode(CONNECTION_REFUSED_SERVER_UNAVAILABLE) + .build(), + getLocal(AuthError.class) + .cause("Failed to check Will publish permission") + .peerAddress(ChannelAttrs.socketAddress(ctx.channel()))); + } + } + }); + } + @Override protected void handleMqttMessage(MqttMessage message) { // never happen in MQTT3 diff --git a/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/v5/MQTT5ConnectHandler.java b/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/v5/MQTT5ConnectHandler.java index 5d3934ba8..e9034780e 100644 --- a/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/v5/MQTT5ConnectHandler.java +++ b/bifromq-mqtt/bifromq-mqtt-server/src/main/java/org/apache/bifromq/mqtt/handler/v5/MQTT5ConnectHandler.java @@ -55,6 +55,7 @@ import static org.apache.bifromq.mqtt.handler.v5.MQTT5MessageUtils.toWillMessage; import static org.apache.bifromq.mqtt.handler.v5.MQTT5MessageUtils.topicAliasMaximum; import static org.apache.bifromq.mqtt.utils.AuthUtil.buildConnAction; +import static org.apache.bifromq.mqtt.utils.AuthUtil.buildPubAction; import static org.apache.bifromq.mqtt.utils.MQTT5MessageSizer.MIN_CONTROL_PACKET_SIZE; import static org.apache.bifromq.plugin.eventcollector.ThreadLocalEventPool.getLocal; import static org.apache.bifromq.type.MQTTClientInfoConstants.MQTT_CHANNEL_ID_KEY; @@ -103,6 +104,7 @@ import org.apache.bifromq.plugin.clientbalancer.IClientBalancer; import org.apache.bifromq.plugin.clientbalancer.Redirection; import org.apache.bifromq.plugin.eventcollector.OutOfTenantResource; +import org.apache.bifromq.plugin.eventcollector.mqttbroker.accessctrl.PubActionDisallow; import org.apache.bifromq.plugin.eventcollector.mqttbroker.channelclosed.AuthError; import org.apache.bifromq.plugin.eventcollector.mqttbroker.channelclosed.EnhancedAuthAbortByClient; import org.apache.bifromq.plugin.eventcollector.mqttbroker.channelclosed.MalformedClientIdentifier; @@ -379,6 +381,59 @@ protected CompletableFuture checkConnectPermission(MqttConnectMessag }); } + @Override + protected CompletableFuture checkWillPermission(MqttConnectMessage message, + LWT willMessage, + SuccessInfo successInfo) { + ClientInfo clientInfo = successInfo.clientInfo(); + return authProvider.checkPermission(clientInfo, buildPubAction(willMessage.getTopic(), + willMessage.getMessage().getPubQoS(), willMessage.getMessage().getIsRetain(), + toUserProperties(message.payload().willProperties()))) + .handle((checkResult, e) -> { + if (e != null) { + return goAway(MqttMessageBuilders.connAck() + .properties(MQTT5MessageBuilders.connAckProperties() + .reasonString("Failed to check Will publish permission") + .build()) + .returnCode(CONNECTION_REFUSED_UNSPECIFIED_ERROR) + .build(), + getLocal(AuthError.class) + .cause("Failed to check Will publish permission") + .peerAddress(ChannelAttrs.socketAddress(ctx.channel()))); + } + switch (checkResult.getTypeCase()) { + case GRANTED -> { + return AuthResult.ok(successInfo); + } + case DENIED -> { + return goAway(MqttMessageBuilders.connAck() + .properties(MQTT5MessageBuilders.connAckProperties() + .reasonString("Will publish not authorized") + .build()) + .returnCode(CONNECTION_REFUSED_NOT_AUTHORIZED_5) + .build(), + getLocal(PubActionDisallow.class) + .isLastWill(true) + .topic(willMessage.getTopic()) + .qos(willMessage.getMessage().getPubQoS()) + .isRetain(willMessage.getMessage().getIsRetain()) + .clientInfo(clientInfo)); + } + default -> { + return goAway(MqttMessageBuilders.connAck() + .properties(MQTT5MessageBuilders.connAckProperties() + .reasonString("Failed to check Will publish permission") + .build()) + .returnCode(CONNECTION_REFUSED_UNSPECIFIED_ERROR) + .build(), + getLocal(AuthError.class) + .cause("Failed to check Will publish permission") + .peerAddress(ChannelAttrs.socketAddress(ctx.channel()))); + } + } + }); + } + private void extendedAuth(MQTT5ExtendedAuthData authData) { this.isAuthing = true; authProvider.extendedAuth(authData) diff --git a/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v3/BaseMQTTTest.java b/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v3/BaseMQTTTest.java index a08a8ae49..01c457dc6 100644 --- a/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v3/BaseMQTTTest.java +++ b/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v3/BaseMQTTTest.java @@ -276,6 +276,10 @@ protected void mockAuthPass(String... attrsKeyValues) { .thenReturn(CompletableFuture.completedFuture(CheckResult.newBuilder() .setGranted(Granted.getDefaultInstance()) .build())); + when(authProvider.checkPermission(any(ClientInfo.class), argThat(MQTTAction::hasPub))) + .thenReturn(CompletableFuture.completedFuture(CheckResult.newBuilder() + .setGranted(Granted.getDefaultInstance()) + .build())); } protected void mockAuthReject(Reject.Code code, String reason) { diff --git a/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v3/MQTTConnectTest.java b/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v3/MQTTConnectTest.java index 140e9d3bc..10314abe4 100644 --- a/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v3/MQTTConnectTest.java +++ b/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v3/MQTTConnectTest.java @@ -29,8 +29,12 @@ import static org.apache.bifromq.plugin.eventcollector.EventType.MQTT_SESSION_START; import static org.apache.bifromq.plugin.eventcollector.EventType.PING_REQ; import static org.apache.bifromq.type.MQTTClientInfoConstants.MQTT_PROTOCOL_VER_KEY; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotEquals; import static org.testng.Assert.assertNull; @@ -39,17 +43,24 @@ import io.netty.handler.codec.mqtt.MqttConnAckMessage; import io.netty.handler.codec.mqtt.MqttConnectMessage; import io.netty.handler.codec.mqtt.MqttMessage; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeUnit; import lombok.extern.slf4j.Slf4j; import org.apache.bifromq.inbox.rpc.proto.AttachReply; import org.apache.bifromq.inbox.rpc.proto.DetachReply; import org.apache.bifromq.mqtt.utils.MQTTMessageUtils; +import org.apache.bifromq.plugin.authprovider.type.CheckResult; +import org.apache.bifromq.plugin.authprovider.type.Denied; +import org.apache.bifromq.plugin.authprovider.type.Error; +import org.apache.bifromq.plugin.authprovider.type.MQTTAction; import org.apache.bifromq.plugin.authprovider.type.Reject; import org.apache.bifromq.plugin.eventcollector.Event; import org.apache.bifromq.plugin.eventcollector.EventType; import org.apache.bifromq.plugin.eventcollector.mqttbroker.clientconnected.ClientConnected; +import org.apache.bifromq.type.ClientInfo; import org.mockito.ArgumentCaptor; import org.testng.Assert; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; @Slf4j @@ -253,12 +264,58 @@ public void validWillTopic() { mockInboxDetach(DetachReply.Code.OK); MqttConnectMessage connectMessage = MQTTMessageUtils.qoSWillMqttConnectMessage(1, true); channel.writeInbound(connectMessage); + channel.runPendingTasks(); MqttConnAckMessage ackMessage = channel.readOutbound(); // verifications assertEquals(ackMessage.variableHeader().connectReturnCode(), CONNECTION_ACCEPTED); verifyEvent(MQTT_SESSION_START, CLIENT_CONNECTED); } + @Test + public void willPermissionDeniedBeforeSessionCalls() { + mockAuthPass(); + when(authProvider.checkPermission(any(ClientInfo.class), argThat(MQTTAction::hasPub))) + .thenReturn(CompletableFuture.completedFuture(CheckResult.newBuilder() + .setDenied(Denied.getDefaultInstance()) + .build())); + + channel.writeInbound(MQTTMessageUtils.qoSWillMqttConnectMessage(1, true)); + channel.advanceTimeBy(disconnectDelay, TimeUnit.MILLISECONDS); + channel.runScheduledPendingTasks(); + channel.runPendingTasks(); + + MqttConnAckMessage connAck = channel.readOutbound(); + assertEquals(connAck.variableHeader().connectReturnCode(), CONNECTION_REFUSED_NOT_AUTHORIZED); + verifyNoInteractions(inboxClient, sessionDictClient); + } + + @DataProvider + public Object[][] willPermissionErrors() { + return new Object[][] { + {CompletableFuture.failedFuture(new RuntimeException("auth unavailable"))}, + {CompletableFuture.completedFuture(CheckResult.newBuilder() + .setError(Error.getDefaultInstance()) + .build())}, + {CompletableFuture.completedFuture(CheckResult.getDefaultInstance())} + }; + } + + @Test(dataProvider = "willPermissionErrors") + public void willPermissionErrorBeforeSessionCalls(CompletableFuture permissionResult) { + mockAuthPass(); + when(authProvider.checkPermission(any(ClientInfo.class), argThat(MQTTAction::hasPub))) + .thenReturn(permissionResult); + + channel.writeInbound(MQTTMessageUtils.qoSWillMqttConnectMessage(1, true)); + channel.advanceTimeBy(disconnectDelay, TimeUnit.MILLISECONDS); + channel.runScheduledPendingTasks(); + channel.runPendingTasks(); + + MqttConnAckMessage connAck = channel.readOutbound(); + assertEquals(connAck.variableHeader().connectReturnCode(), CONNECTION_REFUSED_SERVER_UNAVAILABLE); + verifyNoInteractions(inboxClient, sessionDictClient); + } + @Test public void pingAndPingResp() { mockAuthPass(); @@ -267,6 +324,7 @@ public void pingAndPingResp() { mockInboxDetach(DetachReply.Code.OK); MqttConnectMessage connectMessage = MQTTMessageUtils.qoSWillMqttConnectMessage(1, true); channel.writeInbound(connectMessage); + channel.runPendingTasks(); MqttConnAckMessage ackMessage = channel.readOutbound(); // verifications assertEquals(ackMessage.variableHeader().connectReturnCode(), CONNECTION_ACCEPTED); diff --git a/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v3/MQTTWillMessageTest.java b/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v3/MQTTWillMessageTest.java index 941711e2e..2aae24fb6 100644 --- a/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v3/MQTTWillMessageTest.java +++ b/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v3/MQTTWillMessageTest.java @@ -28,7 +28,6 @@ import static org.apache.bifromq.plugin.eventcollector.EventType.MQTT_SESSION_STOP; import static org.apache.bifromq.plugin.eventcollector.EventType.MSG_RETAINED; import static org.apache.bifromq.plugin.eventcollector.EventType.MSG_RETAINED_ERROR; -import static org.apache.bifromq.plugin.eventcollector.EventType.PUB_ACTION_DISALLOW; import static org.apache.bifromq.plugin.eventcollector.EventType.RETAIN_MSG_CLEARED; import static org.apache.bifromq.plugin.eventcollector.EventType.WILL_DISTED; import static org.apache.bifromq.plugin.eventcollector.EventType.WILL_DIST_ERROR; @@ -36,8 +35,7 @@ import static org.apache.bifromq.retain.rpc.proto.RetainReply.Result.ERROR; import static org.apache.bifromq.retain.rpc.proto.RetainReply.Result.RETAINED; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyLong; -import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.testng.Assert.assertEquals; @@ -47,6 +45,7 @@ import java.util.concurrent.TimeUnit; import lombok.extern.slf4j.Slf4j; import org.apache.bifromq.mqtt.utils.MQTTMessageUtils; +import org.apache.bifromq.plugin.authprovider.type.MQTTAction; import org.apache.bifromq.plugin.eventcollector.EventType; import org.apache.bifromq.sessiondict.rpc.proto.ServerRedirection; import org.apache.bifromq.type.ClientInfo; @@ -109,15 +108,16 @@ public void willWhenNotSelfKick() { } @Test - public void willAuthCheckFailed() { + public void willPermissionNotRechecked() { setupTransientSessionWithLWT(false); mockAuthCheck(false); + mockDistDist(true); channel.advanceTimeBy(50, TimeUnit.SECONDS); testTicker.advanceTimeBy(50, TimeUnit.SECONDS); channel.runPendingTasks(); Assert.assertFalse(channel.isActive()); - verifyEvent(MQTT_SESSION_START, CLIENT_CONNECTED, IDLE, PUB_ACTION_DISALLOW, MQTT_SESSION_STOP); - verify(distClient, times(0)).pub(anyLong(), anyString(), any(), any(ClientInfo.class)); + verifyEvent(MQTT_SESSION_START, CLIENT_CONNECTED, IDLE, MQTT_SESSION_STOP, WILL_DISTED); + verify(authProvider, times(1)).checkPermission(any(ClientInfo.class), argThat(MQTTAction::hasPub)); } @Test diff --git a/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v5/ConnectHandlerTest.java b/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v5/ConnectHandlerTest.java index 0e8a13c30..1c6684bfd 100644 --- a/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v5/ConnectHandlerTest.java +++ b/bifromq-mqtt/bifromq-mqtt-server/src/test/java/org/apache/bifromq/mqtt/handler/v5/ConnectHandlerTest.java @@ -21,13 +21,16 @@ import static io.netty.handler.codec.mqtt.MqttConnectReturnCode.CONNECTION_REFUSED_CLIENT_IDENTIFIER_NOT_VALID; import static io.netty.handler.codec.mqtt.MqttConnectReturnCode.CONNECTION_REFUSED_MALFORMED_PACKET; +import static io.netty.handler.codec.mqtt.MqttConnectReturnCode.CONNECTION_REFUSED_NOT_AUTHORIZED_5; import static io.netty.handler.codec.mqtt.MqttConnectReturnCode.CONNECTION_REFUSED_PACKET_TOO_LARGE; import static io.netty.handler.codec.mqtt.MqttConnectReturnCode.CONNECTION_REFUSED_PROTOCOL_ERROR; import static io.netty.handler.codec.mqtt.MqttConnectReturnCode.CONNECTION_REFUSED_QUOTA_EXCEEDED; import static io.netty.handler.codec.mqtt.MqttConnectReturnCode.CONNECTION_REFUSED_SERVER_BUSY; import static io.netty.handler.codec.mqtt.MqttConnectReturnCode.CONNECTION_REFUSED_SERVER_MOVED; +import static io.netty.handler.codec.mqtt.MqttConnectReturnCode.CONNECTION_REFUSED_UNSPECIFIED_ERROR; import static io.netty.handler.codec.mqtt.MqttConnectReturnCode.CONNECTION_REFUSED_USE_ANOTHER_SERVER; import static io.netty.handler.codec.mqtt.MqttProperties.MqttPropertyType.SERVER_REFERENCE; +import static io.netty.handler.codec.mqtt.MqttProperties.MqttPropertyType.WILL_DELAY_INTERVAL; import static org.apache.bifromq.plugin.eventcollector.EventType.MALFORMED_CLIENT_IDENTIFIER; import static org.apache.bifromq.plugin.eventcollector.EventType.MALFORMED_USERNAME; import static org.apache.bifromq.plugin.eventcollector.EventType.MALFORMED_WILL_TOPIC; @@ -40,6 +43,7 @@ import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; @@ -54,6 +58,7 @@ import io.netty.handler.codec.mqtt.MqttConnAckMessage; import io.netty.handler.codec.mqtt.MqttConnectMessage; import io.netty.handler.codec.mqtt.MqttMessageBuilders; +import io.netty.handler.codec.mqtt.MqttProperties; import io.netty.handler.codec.mqtt.MqttVersion; import java.net.InetSocketAddress; import java.util.Collections; @@ -70,6 +75,8 @@ import org.apache.bifromq.mqtt.spi.IUserPropsCustomizer; import org.apache.bifromq.plugin.authprovider.IAuthProvider; import org.apache.bifromq.plugin.authprovider.type.CheckResult; +import org.apache.bifromq.plugin.authprovider.type.Denied; +import org.apache.bifromq.plugin.authprovider.type.Error; import org.apache.bifromq.plugin.authprovider.type.Granted; import org.apache.bifromq.plugin.authprovider.type.MQTT5AuthData; import org.apache.bifromq.plugin.authprovider.type.MQTT5AuthResult; @@ -88,9 +95,11 @@ import org.apache.bifromq.plugin.settingprovider.ISettingProvider; import org.apache.bifromq.plugin.settingprovider.Setting; import org.apache.bifromq.type.ClientInfo; +import org.apache.bifromq.type.QoS; import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.testng.annotations.BeforeMethod; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; public class ConnectHandlerTest extends MockableTest { @@ -431,7 +440,62 @@ public void inboxExistCallBackPressureRejected() { verify(eventCollector).report(argThat(e -> e.type() == EventType.SERVER_BUSY)); } - + @DataProvider + public Object[][] willDelaySeconds() { + return new Object[][] {{0}, {7}}; + } + + @Test(dataProvider = "willDelaySeconds") + public void willPermissionDeniedBeforeInbox(int willDelaySeconds) { + mockAuthPass(); + when(authProvider.checkPermission(any(ClientInfo.class), argThat(MQTTAction::hasPub))) + .thenReturn(CompletableFuture.completedFuture(CheckResult.newBuilder() + .setDenied(Denied.getDefaultInstance()) + .build())); + + channel.writeInbound(willConnect(willDelaySeconds)); + channel.advanceTimeBy(6, TimeUnit.SECONDS); + channel.runScheduledPendingTasks(); + channel.runPendingTasks(); + + MqttConnAckMessage connAck = channel.readOutbound(); + assertEquals(connAck.variableHeader().connectReturnCode(), CONNECTION_REFUSED_NOT_AUTHORIZED_5); + verifyNoInteractions(inboxClient); + verify(authProvider).checkPermission(any(ClientInfo.class), argThat(action -> action.hasPub() + && action.getPub().getTopic().equals("will/topic") + && action.getPub().getQos() == QoS.AT_MOST_ONCE + && action.getPub().getIsRetained() + && action.getPub().getUserProps().getUserPropertiesCount() == 1 + && action.getPub().getUserProps().getUserProperties(0).getKey().equals("key") + && action.getPub().getUserProps().getUserProperties(0).getValue().equals("value"))); + } + + @DataProvider + public Object[][] willPermissionErrors() { + return new Object[][] { + {CompletableFuture.failedFuture(new RuntimeException("auth unavailable"))}, + {CompletableFuture.completedFuture(CheckResult.newBuilder() + .setError(Error.getDefaultInstance()) + .build())}, + {CompletableFuture.completedFuture(CheckResult.getDefaultInstance())} + }; + } + + @Test(dataProvider = "willPermissionErrors") + public void willPermissionErrorBeforeInbox(CompletableFuture permissionResult) { + mockAuthPass(); + when(authProvider.checkPermission(any(ClientInfo.class), argThat(MQTTAction::hasPub))) + .thenReturn(permissionResult); + + channel.writeInbound(willConnect(0)); + channel.advanceTimeBy(6, TimeUnit.SECONDS); + channel.runScheduledPendingTasks(); + channel.runPendingTasks(); + + MqttConnAckMessage connAck = channel.readOutbound(); + assertEquals(connAck.variableHeader().connectReturnCode(), CONNECTION_REFUSED_UNSPECIFIED_ERROR); + verifyNoInteractions(inboxClient); + } @Test public void receiveMaximumZeroIsProtocolError() { @@ -490,4 +554,32 @@ public void attachCallBackPressureRejected() { verify(eventCollector).report(argThat(e -> e.type() == EventType.SERVER_BUSY)); } + private void mockAuthPass() { + when(authProvider.auth(any(MQTT5AuthData.class))).thenReturn(CompletableFuture.completedFuture( + MQTT5AuthResult.newBuilder().setSuccess(Success.newBuilder().setTenantId("tenantId").build()).build())); + when(authProvider.checkPermission(any(ClientInfo.class), argThat(MQTTAction::hasConn))).thenReturn( + CompletableFuture.completedFuture( + CheckResult.newBuilder().setGranted(Granted.getDefaultInstance()).build())); + when(authProvider.checkPermission(any(ClientInfo.class), argThat(MQTTAction::hasPub))).thenReturn( + CompletableFuture.completedFuture( + CheckResult.newBuilder().setGranted(Granted.getDefaultInstance()).build())); + } + + private MqttConnectMessage willConnect(int willDelaySeconds) { + MqttProperties willProperties = MQTT5MessageUtils.mqttProps() + .addUserProperty("key", "value") + .build(); + willProperties.add(new MqttProperties.IntegerProperty(WILL_DELAY_INTERVAL.value(), willDelaySeconds)); + return MqttMessageBuilders.connect() + .clientId("client") + .cleanSession(true) + .protocolVersion(MqttVersion.MQTT_5) + .willFlag(true) + .willTopic("will/topic") + .willRetain(true) + .willMessage("will") + .willProperties(willProperties) + .build(); + } + }