Skip to main content

Spring Websocket integration test with JWT Authentication

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

Popular posts from this blog

Spring WebSocket authentication with JWT Spring Security

1. Overview In this tutorial, we will talk about how to authenticate your Spring websocket server using JWT Spring Secutiry. Suppose that you already have a stateless microservice with an authorisation server up and running. Now you would want that all of your websocket service instances have to connect to the authorisation service to validate user's access token before allowing them to establish a websocket connection to your server. The official Spring document [1] has mentioned about this. However, they don't clarify how can we obtain the Authentication object. One can write a custom token validation mechanism and validate the user header token manually. But in our opinion, it is best to avoid doing this, as it may introduce some security vulnerabilities. In this tutorial, we will walk you though step-by-step of how to validate user's access token by reusing the JwtAuthentication provided by Spring Security framework. 2. Gradle Dependencies Since this is a Gradle-ba...