Spring's WebSocket implementation provides a WebSocketSession interface that allows you to access the authenticated user's information from the WebSocketSession.getPrincipal() method. When the user is not authenticated, the result of this method will be null. We can use this to determine whether the user has been authenticated or not, which is very convenient.

However, when we need to separate the real-time-related feature into a standalone application, we cannot use the WebSocketSession.getPrincipla() method to retrieve the authenticated user's information unless we set up a clustered session using Spring Session. But that's not the only way. We can also use JWT to perform the authentication.

In this section, we will introduce how to generate a JWT real-time token to perform authentication on a WebSocket connection in our application.

As mentioned earlier, the real-time token will be sent to the client in the response of the /api/me API. It is where we will generate the token with the user ID as the JWT subject. We will use the jjwt library (https://github.com/jwtk/jjwt) to generate and verify the token.

We will create TokenManager to isolate the logic for generating the JWT from the rest of the application. Here is how com.taskagile.domain.common.security.TokenManager looks:

@Component
public class TokenManager {
  private Key secretKey;
  public TokenManager(
      @Value("${app.token-secret-key}") String secretKey) {
    this.secretKey = 
       Keys.hmacShaKeyFor(Decoders.BASE64.decode(secretKey));
  }

  public String jwt(UserId userId) {
    return Jwts.builder()
      .setSubject(String.valueOf(userId.value()))
      .signWith(secretKey).compact();
  }
  
  public UserId verifyJwt(String jws) {
    String userIdValue = Jwts.parser().setSigningKey(secretKey)
       .parseClaimsJws(jws).getBody().getSubject();
    return new UserId(Long.valueOf(userIdValue));
  }
}

As you can see, we will inject a secret key that will be added to application.properties into TokenManager. There are two methods in this class, jwt(), which is used to generate a JWT string based on UserId and the verifyJwt() method, which is used to verify the token we receive from the client.

Here is the change to MeApiController:

@Controller
public class MeApiController {
  ...
  private TokenManager tokenManager;
  public MeApiController(...TokenManager tokenManager) {
    ...
    this.tokenManager = tokenManager;
  }

  @GetMapping("/api/me")
  public ResponseEntity<ApiResult> getMyData(@CurrentUser SimpleUser 
  currentUser) {
    ...
    String realTimeToken = tokenManager.jwt(user.getId());
    return MyDataResult.build(user, teams, boards, realTimeServerUrl, 
    realTimeToken);
  }
}

As you can see, we inject the instance of TokenManager and use it to generate a JWT string as a real-time token.

Once the real-time client initializes the connection with the token, our implementation of WebSocketHandler will receive the request in its afterConnectionEstablished() method. That's where we will perform the authentication by verifying the real-time token.

Our implementation of WebSocketHandler is a request dispatcher, called WebSocketRequestDispatcher. It has the following responsibilities:

• Authenticating the request once the connection is established
• Dispatching requests to channel handlers, which we will discuss in detail in the next section
• Cleaning up the session once a connection is closed

For now, we will only implement the first one. Here is how WebSocketRequestDispatcher looks:

@Component
public class WebSocketRequestDispatcher extends TextWebSocketHandler {
  private TokenManager tokenManager;
  public WebSocketRequestDispatcher(TokenManager tokenManager,     
      ChannelHandlerResolver channelHandlerResolver) {
    this.tokenManager = tokenManager;
  }

  @Override
  public void afterConnectionEstablished(WebSocketSession 
  webSocketSession) {
    log.debug("WebSocket connection established");
    RealTimeSession session = new RealTimeSession(webSocketSession);
    String token = session.getToken();

    try {
      UserId userId = tokenManager.verifyJwt(token);
      session.addAttribute("userId", userId);
      session.reply("authenticated");
    } catch (JwtException exception) {
      log.debug("Invalid JWT token value: {}", token);
      session.fail("authentication failed");
    }
  }
}

As you can see, we inject the instance of TokenManager into the request dispatcher, and use it to verify the token received inside the afterConnectionEstablished() method. RealTimeSession is a wrapper over WebSocketSession to provide some convenient methods.

With the implementation of WebSocketRequestDispatcher, our real-time client can be authenticated. Let's commit the code, and here is the commit record: