diff --git a/application/src/test/java/org/thingsboard/server/controller/ApiKeyWebSocketApiTest.java b/application/src/test/java/org/thingsboard/server/controller/ApiKeyWebSocketApiTest.java index b4e96beee2..375d37ae95 100644 --- a/application/src/test/java/org/thingsboard/server/controller/ApiKeyWebSocketApiTest.java +++ b/application/src/test/java/org/thingsboard/server/controller/ApiKeyWebSocketApiTest.java @@ -17,13 +17,23 @@ package org.thingsboard.server.controller; import lombok.extern.slf4j.Slf4j; import org.junit.After; +import org.junit.Assert; import org.junit.Before; import org.junit.Test; import org.thingsboard.server.common.data.pat.ApiKey; import org.thingsboard.server.common.data.pat.ApiKeyInfo; +import org.thingsboard.server.common.data.query.DeviceTypeFilter; +import org.thingsboard.server.common.data.query.EntityCountQuery; import org.thingsboard.server.dao.service.DaoSqlTest; +import org.thingsboard.server.service.ws.telemetry.cmd.v2.EntityCountCmd; +import org.thingsboard.server.service.ws.telemetry.cmd.v2.EntityCountUpdate; +import java.net.URI; import java.net.URISyntaxException; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; import static org.assertj.core.api.Assertions.assertThat; import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; @@ -59,6 +69,51 @@ public class ApiKeyWebSocketApiTest extends WebSocketApiTest { TbTestWebSocketClient client = buildAndConnectWebSocketClientWithApiKeyHeader(apiKey.getValue()); try { assertThat(client.isOpen()).isTrue(); + + DeviceTypeFilter dtf = new DeviceTypeFilter(List.of("default"), "Device"); + EntityCountQuery ecq = new EntityCountQuery(dtf, Collections.emptyList()); + EntityCountCmd cmd = new EntityCountCmd(1, ecq); + client.send(cmd); + + EntityCountUpdate update = client.parseCountReply(client.waitForReply()); + Assert.assertEquals(1, update.getCmdId()); + Assert.assertTrue(update.getCount() >= 0); + } finally { + client.close(); + } + } + + @Test + public void testInvalidApiKeyHeader_connectionRejected() throws Exception { + TbTestWebSocketClient client = new TbTestWebSocketClient( + new URI(WS_URL + wsPort + "/api/ws"), Map.of("X-API-Key", "invalid-key")); + try { + boolean connected = client.connectBlocking(TIMEOUT, TimeUnit.SECONDS); + assertThat(connected).isFalse(); + } finally { + client.close(); + } + } + + @Test + public void testEmptyApiKeyHeader_connectionRejected() throws Exception { + TbTestWebSocketClient client = new TbTestWebSocketClient( + new URI(WS_URL + wsPort + "/api/ws"), Map.of("X-API-Key", "")); + try { + boolean connected = client.connectBlocking(TIMEOUT, TimeUnit.SECONDS); + assertThat(connected).isFalse(); + } finally { + client.close(); + } + } + + @Test + public void testInvalidApiKeyAuthCmd_connectionClosed() throws Exception { + TbTestWebSocketClient client = new TbTestWebSocketClient(new URI(WS_URL + wsPort + "/api/ws")); + assertThat(client.connectBlocking(TIMEOUT, TimeUnit.SECONDS)).isTrue(); + try { + client.authenticateWithApiKey("invalid-key"); + assertThat(client.waitForClose()).isTrue(); } finally { client.close(); } diff --git a/application/src/test/java/org/thingsboard/server/controller/TbTestWebSocketClient.java b/application/src/test/java/org/thingsboard/server/controller/TbTestWebSocketClient.java index 35536eaecc..d663b2c9ac 100644 --- a/application/src/test/java/org/thingsboard/server/controller/TbTestWebSocketClient.java +++ b/application/src/test/java/org/thingsboard/server/controller/TbTestWebSocketClient.java @@ -53,6 +53,8 @@ public class TbTestWebSocketClient extends WebSocketClient { private static final long TIMEOUT = TimeUnit.SECONDS.toMillis(30); + private final CountDownLatch closeLatch = new CountDownLatch(1); + @Getter private volatile String lastMsg; private volatile CountDownLatch reply; @@ -98,6 +100,16 @@ public class TbTestWebSocketClient extends WebSocketClient { @Override public void onClose(int i, String s, boolean b) { log.info("CLOSED."); + closeLatch.countDown(); + } + + public boolean waitForClose() { + try { + return closeLatch.await(TIMEOUT, TimeUnit.MILLISECONDS); + } catch (InterruptedException e) { + log.warn("Failed to await close", e); + return false; + } } @Override