package io.hosuaby.restful.config; import io.hosuaby.restful.repositories.WebSocketSessionRepository; import io.hosuaby.restful.websocket.handlers.ConsoleHandler; import io.hosuaby.restful.websocket.handlers.NotificationHandler; import io.hosuaby.restful.websocket.handlers.TeapotHandler; import java.util.Map; import javax.servlet.http.HttpServletRequest; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.http.server.ServerHttpRequest; import org.springframework.http.server.ServerHttpResponse; import org.springframework.http.server.ServletServerHttpRequest; import org.springframework.web.servlet.HandlerMapping; import org.springframework.web.socket.WebSocketHandler; import org.springframework.web.socket.config.annotation.EnableWebSocket; import org.springframework.web.socket.config.annotation.WebSocketConfigurer; import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry; import org.springframework.web.socket.handler.PerConnectionWebSocketHandler; import org.springframework.web.socket.server.HandshakeInterceptor; /** * Websocket server configuration. */ @Configuration @EnableWebSocket public class WebSocketConfig implements WebSocketConfigurer { @Override public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) { registry .addHandler(notificationHandler(), "/teapots") .addHandler(teapotHandler(), "/teapots/register/{teapotId}") .addHandler(consoleHandler(), "/teapots/{teapotId}") .setAllowedOrigins("*") // allow all origins .addInterceptors(new UriTemplateHandshakeInterceptor()); } /** * @return one notification handler bean for all WebSocket connections. */ @Bean public WebSocketHandler notificationHandler() { return new NotificationHandler(); } /** * @return one teapot handler bean per WebSocket connection. */ @Bean public WebSocketHandler teapotHandler() { return new PerConnectionWebSocketHandler(TeapotHandler.class); } /** * @return one console handler bean per WebSocket connection. */ @Bean public WebSocketHandler consoleHandler() { return new PerConnectionWebSocketHandler(ConsoleHandler.class); } /** * @return repository for sessions of teapots. */ @Bean(name = "teapotSessionsRepository") public WebSocketSessionRepository teapotSessionsRepository() { return new WebSocketSessionRepository(); } /** * Handshake interceptor that copies URI template variables to attributes * of the WebSocket session. Working together with SimpleUrlHandlerMapping * of Spring MVC that copy path variables into attributes of the * HttpServletRequest. */ private class UriTemplateHandshakeInterceptor implements HandshakeInterceptor { @Override public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception { /* Retrieve original HTTP request */ HttpServletRequest origRequest = ((ServletServerHttpRequest) request).getServletRequest(); /* Retrieve template variables */ Map<String, String> uriTemplateVars = (Map<String, String>) origRequest .getAttribute( HandlerMapping.URI_TEMPLATE_VARIABLES_ATTRIBUTE); /* Put template variables into WebSocket session attributes */ if (uriTemplateVars != null) { attributes.putAll(uriTemplateVars); } return true; } @Override public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {} } }