Quarkus WebSocket server async MDC Context propagation - java

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());
. . .
}
}

Related

Spring boot Restcontroller and WebSocket without STOMP

I have the following problem.
I'm trying to develop a Spring Boot application that serves as RestController and also uses Websockets (w/o STOMP). The RestController has an universal GetMapping method ("/{name}/**) that fetches, depending on name variable, content out of a template database.
My websocket handler should react as broadcast message broker for calls at the endpoint ("/broadcast").
When I test it with Postman, the broadcast websocket call just calls my Restcontroller, what is not intended. It should connect to the Websocket handler.
My code looks like this:
WebSocketConfig:
#Configuration
#EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
#Override
public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
registry.addHandler(webSocketBroadcastHandler(), "/broadcast").setAllowedOrigins("*");
}
#Bean
public WebSocketHandler webSocketBroadcastHandler() {
CfWebSocketBroadcastHandler swsh = new CfWebSocketBroadcastHandler();
return swsh;
}
}
Broadcast handler:
public class CfWebSocketBroadcastHandler extends TextWebSocketHandler {
private static Set<WebSocketSession> sessions = null;
public CfWebSocketBroadcastHandler() {
sessions = new CopyOnWriteArraySet<>();
}
#Override
public void afterConnectionEstablished(WebSocketSession session) throws Exception {
sessions.add(session);
}
#Override
public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
sessions.remove(session);
}
#Override
public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
if (!sessions.contains(session)) {
sessions.add(session);
}
String request = message.getPayload();
WebSocketBroadcastMessage bcm = new Gson().fromJson(request, WebSocketBroadcastMessage.class);
broadcast(bcm);
}
public void broadcast(WebSocketBroadcastMessage bcm) {
for (WebSocketSession session : sessions) {
if (session.isOpen()) {
try {
session.sendMessage(new TextMessage(new Gson().toJson(bcm)));
} catch (IOException ex) {
Logger.getLogger(CfWebSocketBroadcastHandler.class.getName()).log(Level.SEVERE, null, ex);
}
}
}
}
}
RestController:
#RestController
#Component
#Configuration
public class MyRestcontroller {
#GetMapping(path = "/{name}/**")
public void universalGet(#PathVariable("name") String name, #Context HttpServletRequest request, #Context HttpServletResponse response) {
// get template from databse with name variable
}
}
How can I make sure, that the websocket handler gets called instead of the restcontroller?
Further infos:
I'm using spring-boot 2.6.7 and embedded Tomcat 9.0.0.M6.
The maven dependencies are included.
Thanks for any help.

How to retrieve the current logged in user in a websocket controller

I am trying to obtain the currently authenticated user in the controller for websockets. The problem is, I cannot access the user using SecurityContextHolder.getContext().getAuthentication().getPrincipal()).getId().
I have tried to give Principal as a parameter to the method but it returns principal null.
Security configuration:
#Configuration
#EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {
#Override
public void registerStompEndpoints(StompEndpointRegistry registry) {
registry.addEndpoint("/connect").setAllowedOrigins("*");
}
#Override
public void configureMessageBroker(MessageBrokerRegistry registry) {
registry.enableSimpleBroker("/topic/messages");
registry.setApplicationDestinationPrefixes("/ws");
}
}
Controller for websocket:
#Controller
public class MessageController {
#Autowired
private Consumer consumer;
#Autowired
private Utils utils;
#Autowired
private PersonService personService;
#Autowired
SimpMessagingTemplate simpMessagingTemplate;
String destination = "/topic/messages";
ExecutorService executorService =
Executors.newFixedThreadPool(1);
Future<?> submittedTask;
#MessageMapping("/start")
public void startTask(Principal principal){
// Here, I would like to get the logged in user
// If I use principal like this: principal.getName() => NullPointerException
if ( submittedTask != null ){
simpMessagingTemplate.convertAndSend(destination,
"Task already started");
return;
}
simpMessagingTemplate.convertAndSendToUser(sha.getUser().getName(), destination,
"Started task");
submittedTask = executorService.submit(() -> {
while(true){
simpMessagingTemplate.convertAndSend(destination,
// "The calculated value " + val + " is equal to : " + max);
}
});
}
How can I get the authenticated user? I needed it to check when to start the task for the web socket
Try to implement ChannelInterceptor, that need to be registrated in Config file (class that implements WebSocketMessageBrokerConfigurer)
#Configuration
#EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {
private final ChannelInterceptor serverPushInBoundInterceptor;
#Autowired
public WebSocketConfig(#Qualifier("serverPushInBoundInterceptor") ChannelInterceptor serverPushInBoundInterceptor) {
this.serverPushInBoundInterceptor = serverPushInBoundInterceptor;
}
....
#Override
public void configureClientInboundChannel(ChannelRegistration registration) {
registration.interceptors(serverPushInBoundInterceptor);
}
}
#Component("serverPushInBoundInterceptor")
public class ServerPushInBoundInterceptor implements ChannelInterceptor {
private static final Logger log = LoggerFactory.getLogger(ServerPushInBoundInterceptor.class);
#Override
#SuppressWarnings("NullableProblems")
public Message<?> postReceive(Message<?> message, MessageChannel channel) {
StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
if (StompCommand.CONNECT.equals(Objects.requireNonNull(accessor).getCommand())) {
List<String> authorization = accessor.getNativeHeader("Authorization");
if (authorization != null && !authorization.isEmpty()) {
String auth = authorization.get(0).split(" ")[1];
System.out.println(auth);
try {
// find Principal
Principal principal = ...
accessor.setUser(new UsernamePasswordAuthenticationToken(principal, principal.getCredentials(), principal.getAuthorities()));
} catch (Exception exc) {
log.error("preSend", exc);
}
}
}
return message;
}
}

How to guarantee that subscription was happened in stomp client?

I have the following class:
#Autowired
private WebSocketStompClient client;
private volatile StompSession stompSession;
public ListenableFuture<StompSession> connect(String token) {
.....
//connecting
return client.connect(settingsBean.getMarketPlaceUrl(),
new WebSocketHttpHeaders(headers),
stompHeadersConnect,
new StompSessionHandlerAdapter() {
#Override
public void afterConnected(StompSession session, StompHeaders connectedHeaders) {
stompSession = session;
...
session.subscribe(stompHeadersSubscribe, myFrameHandler);
// As I understand I don't have guarantees that subscription was completed successfully
}
}
public void send(String url, Object obj) {
stompSession.send(url, obj);
}
Sometimes send method invokes when connect has already been established, but subscribtion - was not yet.
How to await moment when subscription have been established?

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.

Categories