1. Overview
In our previous tutorial we have introduced an approach to authenticate a user when connecting to our websocket server using JWT authentication. In this article, we will see how can we perform integration tests to check the behaviours of our security layers in different scenarios, as well as ensuring the correctness of our implementation for future updates.
2. Gradle dependencies
Since this is a Gradle-based project, we add the required testing dependencies to the build.gradle:
testImplementation 'org.springframework.boot:spring-boot-starter-test'
testImplementation 'com.github.tomakehurst:wiremock-jre8:2.34.0'
testImplementation 'org.awaitility:awaitility-kotlin:4.2.0'
testImplementation 'io.projectreactor:reactor-test'
3. Defining a Controller endpoint for testing
In this example, we will write some integration tests to validate the requests to access a controller endpoint, as defined below:
@Controller
public class WebsocketController {
private static final Logger LOGGER = LoggerFactory.getLogger(WebsocketController.class);
@MessageMapping("/message")
public Message incoming(Principal principal, @Payload Message message) {
LOGGER.info("Sending a message from user {} to user {}", principal.getName(), message.getUsername());
return message;
}
}
As can be seen in the code snippets above, we will try to validate the requests which send to the /app/message endpoint, and also see if we can obtain the message which is broadcasted into the /topic/message channel or not.
4. Configuring Inmemory MessageBroker for Testing
In our previous tutorial, we have used an inmemory broker just for the sake of simplicity. However, in a production application, you should want to use a real full-features message broker such as RabbitMQ or ActiveMQ instead. Hence, your websocket configuration should look like:
@Configuration
@EnableWebSocketMessageBroker
@Order(Ordered.HIGHEST_PRECEDENCE + 99)
@Profile("prod")
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {
@Value("${broker.relay.host}")
private String brokerRelayHost;
@Value("${broker.relay.port}")
private int brokerRelayPort;
@Value("${broker.relay.user}")
private String brokerUsername;
@Value("${broker.relay.password}")
private String brokerPassword;
@Override
public void configureMessageBroker(MessageBrokerRegistry registry) {
registry.enableStompBrokerRelay("/queue", "/topic")
.setRelayHost(brokerRelayHost)
.setRelayPort(brokerRelayPort)
.setSystemLogin(brokerUsername)
.setSystemPasscode(brokerPassword)
.setClientLogin(brokerUsername)
.setClientPasscode(brokerPassword)
.setUserDestinationBroadcast("/topic/unresolved-user")
.setUserRegistryBroadcast("/topic/log-user-registry");
registry.setApplicationDestinationPrefixes("/app");
}
...
}
However, by enabling StompBrokerRelay, it requires you to have a message broker instance up and running before starting up this server, which could be complicated at times and is not ideal for CI/CD pipeline process. Instead, we can specify an inmemory broker, only in testing environment by overiding the original websocket configuration as follow:
@Configuration
@Profile("test")
public class WebSocketConfigTest extends WebSocketConfig {
@Override
public void configureMessageBroker(MessageBrokerRegistry registry) {
registry.enableSimpleBroker("/queue", "/topic");
registry.setApplicationDestinationPrefixes("/app");
}
}
Here we have created a WebSocketConfigTest configuration file that annotated with @Profile("test"). This configuration allow us to use Spring's inmemory broker, while keeping all other configurations defined in WebSocketConfig configuration class. Since we have associated this configuration class with "test" profile, it will be initialised instead of the original WebSocketConfig annotated with @Profile("prod") configuration. 5. Defining an Authorisation Mocking Server
In this example, we mock up the response from the Authorisation server as follow:
@ActiveProfiles("test")
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
class WebSocketAuthenticationTests {
private static final String JWKS_RESPONSE = "Paste a mock JWKs here...";
@LocalServerPort
private int rdmServerPort;
private static WireMockServer mockServer;
@BeforeEach
public void setupMockServer() {
mockServer = new WireMockServer(wireMockConfig().port(8081));
mockServer.start();
mockServer.stubFor(get(urlEqualTo("/.well-known/jwks.json"))
.willReturn(aResponse()
.withStatus(200)
.withBody(JWKS_RESPONSE)
.withHeader("Content-Type", "application/json")));
}
@AfterEach
public void tearDownServer() {
mockServer.stop();
}
...
}As can be seen in the code implementation above, we create a new Http WireMockServer on port 8081, whenever we perform a get request to http://localhost:8081/.well-known/jwks.json, we will receive a JWKs response as defined in the JWKS_RESPONSE variable. This will ensure the JWTProcessor bean (that we defined in the previous tutorial) will be able to send and validate user's access token.6. Configuring a websocket stomp client
In order to establish a websocket connection and perform different requests, as well as listen to different websocket channels, we can use a tool provided by Spring framwork which is
WebSocketStompClient to perform and manage the connection:@ActiveProfiles("test")
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
class WebSocketAuthenticationTests {
private final String messageReadToken = "Paste a mock token here";
private WebSocketStompClient stompClient;
private StompSession session;
@LocalServerPort
private int rdmServerPort;
@BeforeEach
public void setupMockServer() {
var webSocketClient = new StandardWebSocketClient();
var sockJsClient = new SockJsClient(Collections.singletonList(new WebSocketTransport(webSocketClient)));
stompClient = new WebSocketStompClient(sockJsClient);
stompClient.setMessageConverter(new MappingJackson2MessageConverter());
...
}
@AfterEach
public void tearDownServer() {
if (session != null) {
session.disconnect();
}
...
}
...
}
One thing to note here is that we use MappingJackson2MessageConverter instead of the default SimpleMessageConverter to convert the payload of incoming and outgoing messages, as for convenience.
7. Integration tests
For the sake of simplicity, we demonstrate two simple but crucial tests to ensure that our service is working properly, including the case when the user not provides access token and when the user provides a correct access token.
7.1 User tries to establish a connection without access token
@Test
void performWhenNoBearerTokenThenDeniesAccess() {
try {
System.out.print("Current port: " + rdmServerPort);
session = stompClient.connect(String.format("ws://localhost:%d/websocket", rdmServerPort),
new StompSessionHandlerAdapter(){}).get(5, SECONDS);
Assert.isTrue(false, "An exception should be thrown before reaching this line");
} catch (Exception ex) {
String message = ex.getMessage();
Assert.isTrue(message.contains("Connection closed"), "Connection should be closed due to AccessDenied");
}
}
As can be seen in the code snippet above, we expect that if a user tries to connect to our websocket server without authentication header, the websocket server should reject and close the connection immediately.
7.2 User tries to establish a connect with a proper access token
@Test
void performWhenHasValidBearerTokenThenAllowsAccess() throws Exception {
var handshakeHeaders = new WebSocketHttpHeaders();
var connectHeaders = new StompHeaders();
connectHeaders.add("Authorization", "Bearer " + messageReadToken);
var message = "myMessage";
var receivedMessages = new LinkedBlockingDeque<Message>();
session = stompClient.connect(String.format("ws://localhost:%d/websocket", rdmServerPort), handshakeHeaders,
connectHeaders, new LoggingStompSessionHandler(Message.class) {
@Override
public void afterConnected(@NotNull StompSession session, @NotNull StompHeaders connectedHeaders) {
session.subscribe("/topic/message", this);
session.send("/app/message", new Message("test", message));
}
@Override
public void handleFrame(@NotNull StompHeaders headers, Object payload) {
Assert.isTrue(receivedMessages.offer((Message) payload), "Should be added");
}
}).get(5, SECONDS);
var response = receivedMessages.poll(10, SECONDS);
Assert.notNull(response, "No response retrieved");
Assert.isTrue(message.equals(response.getMessage()), message + " vs " + response.getMessage());
}
Here, once we have provided the correct access token, the server should allow us to establish a websocket connection. And once the connection has been established, we subcribe to the /topic/message channel and send a message to the /app/message endpoint that we have defined in section 3 above. If everything is working correctly, we should be able to get back the message that we have sent previously, which is indicated that the message have gone through back and forth between our websocket client and server.
One thing to note that the LoggingStompSessionHandler is just an utility class that extend the StompSessionHandlerAdapter class, for convenience:
public class LoggingStompSessionHandler extends StompSessionHandlerAdapter {
private static final Logger LOGGER = LoggerFactory.getLogger(LoggingStompSessionHandler.class);
private final Type payloadType;
public LoggingStompSessionHandler(Type payloadType) {
this.payloadType = payloadType;
}
@Override
public Type getPayloadType(StompHeaders headers) {
return payloadType;
}
@Override
public void handleException(StompSession session, StompCommand command, StompHeaders headers, byte[] payload, Throwable exception) {
LOGGER.error("Exception occurred when handling WebSocket message.", exception);
}
@Override
public void handleTransportError(StompSession session, Throwable exception) {
LOGGER.error("Exception occurred when transporting data.", exception);
}
}
A simple implementation of this tutorial can be found here: https://github.com/dispatcher-servlet/jwt-websocket-integration-test
Comments
Post a Comment