We Keep Coding

Quarkus WebSocket server async MDC Context propagation

I'm trying to create a Quarkus WebSocket server with some async workaround.
I want to process the incoming messages asynchronously by publishing an event in the Vertx EventBus and then process them in a different 'service' class. At the same time, I want to be able to propagate the MDC context.
Here is an example of what I'm trying to do, but so far the MDC context propagation is not working.
// WebSocket endpoint controller
#Slf4j
#ApplicationScoped
#ServerEndpoint(value = "/users/{userId}")
class UserWebSocketController {
private final WebsocketConnectionService websocketConnectionService;
private final Vertx vertx;
UserWebSocketController(WebsocketConnectionService websocketConnectionService, Vertx vertx) {
this.websocketConnectionService = websocketConnectionService;
this.vertx = vertx;
}
#OnOpen
void onOpen(Session session, #PathParam("userId") String userId) {
MDC.put("websocket.sessionId", session.getId());
MDC.put("user.id", userId);
log.info("New WebSocket Session opened.");
websocketConnectionService.addConnection(userId, session);
}
#OnMessage
void onMessage(Session session, String message, #PathParam("userId") String userId) {
// How do I get the same MDC context here?
//MDC.get("user.id") is null here
log.info("New message received.");
vertx.eventBus().send("websocket.message.new", message);
}
#OnClose
void onClose(Session session, #PathParam("userId") String userId) {
log.info("WebSocket Session closed.");
websocketConnectionService.removeSession(userId);
}
#OnError
void onError(Session session, #PathParam("userId") String userId, Throwable throwable) {
log.error("There was an error in the WebSocket Session.");
websocketConnectionService.removeSession(userId);
}
}
// Service class
#Slf4j
#ApplicationScoped
class UserService {
private final WebsocketConnectionService websocketConnectionService;
private final Vertx vertx;
UserService(WebsocketConnectionService websocketConnectionService, Vertx vertx) {
this.websocketConnectionService = websocketConnectionService;
this.vertx = vertx;
}
#ConsumeEvent("websocket.message.new")
Uni<Void> handleWebSocketMessages(String message) {
// How do I get the same MDC context here?
final var userId = MDC.get("user.id"); // this is null
log.info("'userId' exists in the MDC Context (userId=%s)".formatted(userId));
// do some business with the userId
return Uni.createFrom().voidItem();
}
}
Do you have any idea how can I make this context propagation work?

The MDC is, if I remember correctly, ThreadLocal, and only filled/relevant during the initial #OnOpen call.
Store MDC context map somewhere, maybe using the connection's Session as a key?
class UserWebSocketController {
class AsyncMessage implements Serializable {
AsyncMessage(String sessionId, String message) {
...
}
}
private static final Map<String,Map<String,String>> ACTIVE_MDCS = new ...
. . .
void onOpen(. . .) {
MDC.put("websocket.sessionId", session.getId());
MDC.put("user.id", userId);
storeSessionMDC(session.getId());
. . .
}
void onMessage(. . .) {
restoreSessionMDC(session.getId());
. . .
vertx.eventBus().send("websocket.message.new",
new AsyncMessage(session.getId(),message));
}
void onClose(. . .) {
restoreSessionMDC(session.getId());
. . .
removeSessionMDC(session.getId());
. . .
}
// etc.
private static void storeSessionMDC(String sessionId) {
ACTIVE_MDCS.put(sessionId, MDC.getCopyOfContextMap());
}
public static void restoreSessionMDC(String sessionId) {
MDC.setContextMap(ACTIVE_MDCS.get(sessionId));
}
private static void removeSessionMDC(String sessionId) {
ACTIVE_MDCS.remove(sessionId);
}
}
and
class UserService {
Uni<Void> handleWebSocketMessages(AsyncMessage asyncMessage) {
UserWebSocketController
.restoreSessionMDC(asyncMessage.getSessionId());
. . .
}
}

How to retrieve the session ID from Spring Websocket on initial handshake and to be used later on?

How to retrieve the session ID from Spring Websocket on initial handshake and to be used later on?
Goal
The Goal is whenever a Websocket client connects to the websocket server. The websocket client will pass in a parameter. It will pass in their tenantID. Upon successfull connection to the websocket server. The websocket server usually generates a session ID. What I hope to achieve is associate this websocket sessionID to that specific tenant parameter. And later on, whenever a websocket server will send an update to the that client/tenant. it will send it to that specific tenant via its sessionID that the websocket server has created.
Here is my websocket server configuration..
#Configuration
#EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {
#Override
public void configureMessageBroker(MessageBrokerRegistry config) {
config.enableSimpleBroker("/queue");
config.setApplicationDestinationPrefixes("/app");
}
#Override
public void registerStompEndpoints(StompEndpointRegistry registry) {
//registry.addEndpoint("/client").addInterceptors(new WSHandshakeInterceptor()).withSockJS();
registry.addEndpoint("/agent").addInterceptors(new WSHandshakeInterceptor()).withSockJS();
}
}
WSHandhsakeInterceptor.java
public class WSHandshakeInterceptor implements HandshakeInterceptor {
#Override
public boolean beforeHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, Map<String, Object> map) throws Exception {
return true;
}
#Override
public void afterHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, Exception e) {
HttpHeaders header = serverHttpRequest.getHeaders();
String client = header.get("client-id").get(0);
String sessionId = null;
if (serverHttpRequest instanceof ServletServerHttpRequest) {
ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) serverHttpRequest;
HttpSession session = servletRequest.getServletRequest().getSession();
sessionId = session.getId();
System.out.println("Session ID + "+sessionId);
System.out.println("CLIENT ID "+client);
ClassChangeNotificationServiceImpl.clientSessionMap.put(client, sessionId);
}
}
}
Here is how my websocket client connects to the server.
WebSocketClient simpleWebSocketClient = new StandardWebSocketClient();
List<Transport> transports = new ArrayList<>(1);
transports.add(new WebSocketTransport(simpleWebSocketClient));
SockJsClient sockJsClient = new SockJsClient(transports);
stompClient = new WebSocketStompClient(sockJsClient);
stompClient.setMessageConverter(new MappingJackson2MessageConverter());
String url = "ws://localhost:8081/agent";
sessionHandler = new MyStompSessionHandler();
try {
WebSocketHttpHeaders wsHTTPheaders = new WebSocketHttpHeaders();
wsHTTPheaders.add("CLIENT-ID","XXXTESTCLIENTXXX");
session = stompClient.connect(url, wsHTTPheaders,sessionHandler).get();
} catch (InterruptedException e) {
e.printStackTrace();
} catch (ExecutionException e) {
e.printStackTrace();
}
MyStompSessionHandler.java
import org.springframework.messaging.simp.stomp.StompFrameHandler;
import org.springframework.messaging.simp.stomp.StompHeaders;
import org.springframework.messaging.simp.stomp.StompSession;
import org.springframework.messaging.simp.stomp.StompSessionHandlerAdapter;
import java.lang.reflect.Type;
import java.util.List;
import java.util.Map;
public class MyStompSessionHandler extends StompSessionHandlerAdapter {
private void subscribeTopic(String topic, StompSession session) {
session.subscribe(topic, new StompFrameHandler() {
#Override
public Type getPayloadType(StompHeaders headers) {
System.out.println("HANDLER");
return String.class;
}
#Override
public void handleFrame(StompHeaders headers, Object payload) {
System.out.println("TEST HEKHEk");
}
});
}
#Override
public Type getPayloadType(StompHeaders headers) {
return String.class;
}
/**
* This implementation is empty.
*/
#Override
public void handleFrame(StompHeaders headers, Object payload) {
String resp = (String) payload;
System.out.println("Received responses from websocket server: "+ resp);
}
#Override
public void afterConnected(StompSession session, StompHeaders connectedHeaders) {
subscribeTopic("/user/queue/response", session);
System.out.println("CONNECTEd");
for(Map.Entry<String, List<String>> entry: connectedHeaders.entrySet()) {
System.out.println(entry.getKey()+"S");
}
}
}
Now For the fun part. or what I've been trying to achieve. is I can send an update to the websocket client via this code
this.template.convertAndSendToUser(client, "/queue/reply", message);
Where client is the sessionID that spring has generated. and it will send a message to that specific client only
Question
How do I retrieve the sessionID after a successful handshake?? can I use this sessionID to send specific updates to that client. Here's my code..

Just add this method anywhere in your code to get the session ID after handshake:
#EventListener
private void handleSessionConnected(SessionConnectEvent event) {
String sessionId = SimpAttributesContextHolder.currentAttributes().getSessionId();
}

How to check if websocket user is authenticated when subscribe a channel?

I have websocket config class with :
#Override
public void configureMessageBroker(MessageBrokerRegistry config) {
config.enableSimpleBroker("/topics");
config.setApplicationDestinationPrefixes("/app");
}
#Override
public void registerStompEndpoints(StompEndpointRegistry registry) {
registry.addEndpoint("/notifications").setAllowedOrigins("*").withSockJS();
}
As well as ClientInboundChannel:
#Override
public void configureClientInboundChannel(ChannelRegistration registration) {
registration.setInterceptors(new ChannelInterceptorAdapter() {
#Override
public Message<?> preSend(Message<?> message, MessageChannel channel) {
StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
final String xAuthToken = accessor.getFirstNativeHeader(ManipulatesHeaders.X_AUTH_TOKEN_HEADER);
if (xAuthToken == null) {
return message;
}
final UserDetails userDetails = authService.getUserDetails(xAuthToken);
if (StompCommand.CONNECT == accessor.getCommand()) {
final WebSocketPrincipal principal = (...)
userRegistry.registerUser(principal, message);
accessor.setUser(principal);
}
return message;
}
}
}
Now, I would like to send a welcome messsage to every user that subscribe a particular channel. It's obvious that It might be achieved by creating a class that implements ApplicationListener<SessionSubscribeEvent> and provides
#Override
public void onApplicationEvent(final SessionSubscribeEventevent) {
StompHeaderAccessor sha = StompHeaderAccessor.wrap(event.getMessage());
}
Second approach is to perform similar thing in aforementioned ClientInboundChannel (it's the same).
The problem is, that there is no simpUser header in a message when a user send a stomp SUBSCRIBE.
Process is as follows: User sends stomp CONNECT, this line is executed: accessor.setUser(principal) and at this point message has properly setsimpUser header. But when stomp SUBSCRIBE message is received in preSend, the StompHeaderAccessor does not have simpUser header. Thus I am not able to asses whether user is already authenticated or not.
So how could I check if a user that sent a SUBSCRIBE message is already authenticated or not ? (and why authenticated user does not send simpUser header)

In your websocket configuration class change your extend to AbstractSecurityWebSocketMessageBrokerConfigurer
then you can set your security as follows:
#Override
protected void configureInbound(MessageSecurityMetadataSourceRegistry message) {
message
.nullDestMatcher().permitAll()
.simpDestMatchers("/app/**").authenticated()
.simpSubscribeDestMatchers("/topics/**").authenticated()
.anyMessage().denyAll();
}
if you get the error Could not verify the provided CSRF token because your session was not found you will have to override `sameOriginDisabled in the websocket configuration class
#Override
protected boolean sameOriginDisabled() {
return true;
}
for more information on websocket security:
https://docs.spring.io/spring-security/site/docs/current/reference/html/websocket.html

How to prevent subscription to spring socket /queue/private/* destination.

I have an Java socket API application, that handles socket requests from users and sends responses.
I have a configurer:
#Configuration
#EnableWebSocketMessageBroker
public class WebSocketConfig extends AbstractWebSocketMessageBrokerConfigurer {
private static final Logger LOGGER = Logger.getLogger(WebSocketConfig.class);
#Override
public void configureMessageBroker(MessageBrokerRegistry config) {
config.enableSimpleBroker("/queue");
config.setApplicationDestinationPrefixes("/server_in");
config.setUserDestinationPrefix("/user");
}
#Override
public void registerStompEndpoints(StompEndpointRegistry registry) {
registry.addEndpoint("/websocket").withSockJS();
}
}
When i send response to user i do the following:
this.simpMessagingTemplate.convertAndSend("/queue/private/user_"+secret_key, socketResponse);
On client i have the following code:
sc.subscribe('/queue/private/user_'+secret_key, function (greeting) {
console.log(greeting.body);
});
And the response is handled successfully.
But the problem is that some other user can also subscribe to "/queue/private/*" destination and handle private messages.
sc.subscribe('/queue/private/*', function (greeting) {
console.log(greeting.body);
});
How can I privent that behaviour?

If you want each user to have a socket and only get him messages, what you can do is :
Subscribe as you do to the endPoint but with "/user" infront for example
sc.subscribe('/user/queue/websocket, function (greeting) {
console.log(greeting.body);
});
and at the server side you should have a rest method:
#RequestMapping(value = "/test", method = RequestMethod.POST)
public void test(Principal principal) throws Exception {
this.template.convertAndSendToUser(principal.getName(), "/queue/click", "");
}
With this every user subscibers to each channel and only the user is notified about, when a rest call is made.
The rest call should be authenticated so the Principal has the username.
The user channel is auto managed from Spring so you have to add it like this.

You can extend ChannelInterceptorAdapter and manage each event individually:
public class AuthorizedChannelInterceptorAdapter extends ChannelInterceptorAdapter {
#Override
public Message<?> preSend(Message<?> message, MessageChannel messageChannel) throws AuthenticationException {
StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
if (StompCommand.CONNECT == accessor.getCommand())
setUserAuthenticationToken(accessor);
else if (StompCommand.SUBSCRIBE == accessor.getCommand())
validateSubscription((Authentication) accessor.getUser(), accessor.getDestination());
return message;
}
private void setUserAuthenticationToken(StompHeaderAccessor accessor) {
String token = accessor.getFirstNativeHeader(HttpHeaders.AUTHORIZATION);
accessor.setUser(loadAuthentication(token));
}
private Authentication loadAuthentication(String token){
return ....;
}
private void validateSubscription(Authentication authentication, String destination) {
if(...)
throw new AccessDeniedException("No permission to subscribe to this topic");
}
}
First of all you will need to store the authentication object provided by the client in connection event. After this, each event sent by the client will have this authentication object set so you can use it to validate if it is authorized to subscribe to a particular channel.

How to send a message through web socket to a connected user?

I want to send a message through a web socket to a specific user. So far I can open a web socket and read message from client like that:
#ServerEndpoint(value = "/wsep")
public class WebSocketEndpoint {
private static final Logger LOGGER = LoggerFactory.getLogger(WebSocketEndpoint.class);
private Session session;
#OnOpen
public void onOpen(Session session) {
this.session = session;
try {
session.getBasicRemote().sendText("You are connected. Your ID is " + session.getId());
} catch (Exception e) {
LOGGER.error("Error on open web socket", e);
}
}
#OnMessage
public void onClientMessage(String message, Session session) {
LOGGER.info("Message from {} is: {}", session.getId(), message);
}
#OnClose
public void onClose(Session session) {
this.session = null;
LOGGER.info("{} disconnected", session.getId());
}
}
I have an independent service which creates message in destination to a user. My Message class is a simple POJO:
public class Message {
private String fromUserName;
private String toUserName;
private String content;
...
}
When a new message is created in my MessageService, I want to inform the receiver if he is connected. I think I have to add a method WebSocketEndpoint.onServerMessage:
public void onServerMessage(Session session, Message message) {
session.getBasicRemote().sendText(message.getContent());
}
But I don't know how to do something like that which works.

There will be one instance of the ServerEndpoint for all your users. So, it should store all client sessions. As Vu.N suggested, one way you can do it is using a map:
Map<String, Session> sessions = new ConcurrentHashMap<>();
public void onOpen(Session session) {
String username = [...]
sessions.put(username, session);
}
Then, it will be easy to send the message to the user:
public void onServerMessage(Session session, Message message) {
sessions.get(message.getToUserName())
.getBasicRemote() // see also getAsyncRemote()
.sendText(message.getContent());
}
Now the hardest part is to get the username?
In my past works, I've done it in 3 ways:
The client connects to a URL that has some "key". This "key" would be used to find the right username. The WebSocket server endpoint would be like this:
#ServerEndpoint(value="/wsep/{key}") // the URL will have an extra "key"
public class WebSocketEndpoint {
[...]
#OnOpen
public void onOpen(Session session, #PathParam("key") String key) {
String username = getUserNameWithKey(key);
sessions.add(username, session);
}
The client sends some information in the first message. You just ignore the #OnOpen part:
#OnOpen
public void onOpen(Session session) {
LOGGER.info("Session open. ID: {}", session.getId());
}
#OnMessage
public void onMessage(Session session, String message) {
String username = getUserNameFromMessage(message);
sessions.add(username, session);
}
Some user info can be obtained from a cookie, JWT, or something. You'll need a ServerEndpointConfig.Configurator to get that information from the request. For example:
public class CookieServerConfigurator extends ServerEndpointConfig.Configurator {
#Override
public void modifyHandshake(ServerEndpointConfig sec, HandshakeRequest request, HandshakeResponse response) {
Map<String,List<String>> headers = request.getHeaders();
sec.getUserProperties().put("cookie", headers.get("cookie"));
}
}
Then, the server endpoint will point to that configurator:
#ServerEndpoint(value = "/example2/", configurator = CookieServerConfigurator.class)
public class WebSocketEndpoint {
and you can get the information like this:
#OnOpen
public void onOpen(Session session, EndpointConfig endpointConfig) {
String username = getUsername((List<String>)endpointConfig.getUserProperties().get("cookie"));
You can see some working examples here: https://github.com/matruskan/websocket-example
For more complex systems, you can also have a "tagging system", instead of using a Map. Then, each session can receive messages sent to any of its tags, and a message sent to a tag can be directed to many sessions.