org.springframework.web.socket.WebSocketHandler Java Examples

@Test
public void supportedSubProtocols() {
	this.handshakeHandler.setSupportedProtocols("stomp", "mqtt");
	given(this.upgradeStrategy.getSupportedVersions()).willReturn(new String[] {"13"});
	this.servletRequest.setMethod("GET");

	WebSocketHttpHeaders headers = new WebSocketHttpHeaders(this.request.getHeaders());
	headers.setUpgrade("WebSocket");
	headers.setConnection("Upgrade");
	headers.setSecWebSocketVersion("13");
	headers.setSecWebSocketKey("82/ZS2YHjEnUN97HLL8tbw==");
	headers.setSecWebSocketProtocol("STOMP");

	WebSocketHandler handler = new TextWebSocketHandler();
	Map<String, Object> attributes = Collections.emptyMap();
	this.handshakeHandler.doHandshake(this.request, this.response, handler, attributes);

	verify(this.upgradeStrategy).upgrade(this.request, this.response, "STOMP",
			Collections.emptyList(), null, handler, attributes);
}
 

@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;
}
 

@Test
public void constructorWithAttributeNames() throws Exception {
	Map<String, Object> attributes = new HashMap<>();
	WebSocketHandler wsHandler = Mockito.mock(WebSocketHandler.class);

	this.servletRequest.setSession(new MockHttpSession(null, "123"));
	this.servletRequest.getSession().setAttribute("foo", "bar");
	this.servletRequest.getSession().setAttribute("bar", "baz");

	Set<String> names = Collections.singleton("foo");
	HttpSessionHandshakeInterceptor interceptor = new HttpSessionHandshakeInterceptor(names);
	interceptor.beforeHandshake(this.request, this.response, wsHandler, attributes);

	assertEquals(2, attributes.size());
	assertEquals("bar", attributes.get("foo"));
	assertEquals("123", attributes.get(HttpSessionHandshakeInterceptor.HTTP_SESSION_ID_ATTR_NAME));
}
 

@Test
public void supportedExtensions() {
	WebSocketExtension extension1 = new WebSocketExtension("ext1");
	WebSocketExtension extension2 = new WebSocketExtension("ext2");

	given(this.upgradeStrategy.getSupportedVersions()).willReturn(new String[] {"13"});
	given(this.upgradeStrategy.getSupportedExtensions(this.request)).willReturn(Collections.singletonList(extension1));

	this.servletRequest.setMethod("GET");

	WebSocketHttpHeaders headers = new WebSocketHttpHeaders(this.request.getHeaders());
	headers.setUpgrade("WebSocket");
	headers.setConnection("Upgrade");
	headers.setSecWebSocketVersion("13");
	headers.setSecWebSocketKey("82/ZS2YHjEnUN97HLL8tbw==");
	headers.setSecWebSocketExtensions(Arrays.asList(extension1, extension2));

	WebSocketHandler handler = new TextWebSocketHandler();
	Map<String, Object> attributes = Collections.<String, Object>emptyMap();
	this.handshakeHandler.doHandshake(this.request, this.response, handler, attributes);

	verify(this.upgradeStrategy).upgrade(this.request, this.response, null,
			Collections.singletonList(extension1), null, handler, attributes);
}
 

@Override
public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
    try {
        if (ipBlackListEnabled.typedValue()) {
            String ip = InetSocketAddressHelper.getRemoteIP(request);
            for (String blackIp : ipBlackList.typedValue()) {
                if (ip.equalsIgnoreCase(blackIp)) {
                    return false;
                }
            }
        }
    } catch (Throwable ex) {
        logger.warn("process WebSocket ip black list failed" + ex.getMessage(), ex);
    }
    return true;
}
 

@Override
public final ListenableFuture<WebSocketSession> doHandshake(
		WebSocketHandler handler, @Nullable WebSocketHttpHeaders headers, URI url) {

	Assert.notNull(handler, "WebSocketHandler is required");
	Assert.notNull(url, "URL is required");

	String scheme = url.getScheme();
	if (!supportedProtocols.contains(scheme)) {
		throw new IllegalArgumentException("Invalid scheme: '" + scheme + "'");
	}

	SettableListenableFuture<WebSocketSession> connectFuture = new SettableListenableFuture<>();
	try {
		SockJsUrlInfo sockJsUrlInfo = new SockJsUrlInfo(url);
		ServerInfo serverInfo = getServerInfo(sockJsUrlInfo, getHttpRequestHeaders(headers));
		createRequest(sockJsUrlInfo, headers, serverInfo).connect(handler, connectFuture);
	}
	catch (Throwable exception) {
		if (logger.isErrorEnabled()) {
			logger.error("Initial SockJS \"Info\" request to server failed, url=" + url, exception);
		}
		connectFuture.setException(exception);
	}
	return connectFuture;
}
 

@Test
public void handshakeHandlerPassedToSockJsRegistration() {
	WebSocketHandler handler = new TextWebSocketHandler();
	HandshakeHandler handshakeHandler = new DefaultHandshakeHandler();

	this.registration.addHandler(handler, "/foo").setHandshakeHandler(handshakeHandler).withSockJS();

	List<Mapping> mappings = this.registration.getMappings();
	assertEquals(1, mappings.size());

	Mapping mapping = mappings.get(0);
	assertEquals(handler, mapping.webSocketHandler);
	assertEquals("/foo/**", mapping.path);
	assertNotNull(mapping.sockJsService);

	WebSocketTransportHandler transportHandler =
			(WebSocketTransportHandler) mapping.sockJsService.getTransportHandlers().get(TransportType.WEBSOCKET);
	assertSame(handshakeHandler, transportHandler.getHandshakeHandler());
}
 

@Test
public void handshakeHandlerPassedToSockJsRegistration() {
	WebSocketHandler handler = new TextWebSocketHandler();
	HandshakeHandler handshakeHandler = new DefaultHandshakeHandler();

	this.registration.addHandler(handler, "/foo").setHandshakeHandler(handshakeHandler).withSockJS();
	this.registration.getSockJsServiceRegistration().setTaskScheduler(this.taskScheduler);

	List<Mapping> mappings = this.registration.getMappings();
	assertEquals(1, mappings.size());

	Mapping mapping = mappings.get(0);
	assertEquals(handler, mapping.webSocketHandler);
	assertEquals("/foo/**", mapping.path);
	assertNotNull(mapping.sockJsService);

	WebSocketTransportHandler transportHandler =
			(WebSocketTransportHandler) mapping.sockJsService.getTransportHandlers().get(TransportType.WEBSOCKET);
	assertSame(handshakeHandler, transportHandler.getHandshakeHandler());
}
 

@Override
public ListenableFuture<WebSocketSession> connect(TransportRequest request, WebSocketHandler handler) {
	final SettableListenableFuture<WebSocketSession> future = new SettableListenableFuture<WebSocketSession>();
	WebSocketClientSockJsSession session = new WebSocketClientSockJsSession(request, handler, future);
	handler = new ClientSockJsWebSocketHandler(session);
	request.addTimeoutTask(session.getTimeoutTask());

	URI url = request.getTransportUrl();
	WebSocketHttpHeaders headers = new WebSocketHttpHeaders(request.getHandshakeHeaders());
	if (logger.isDebugEnabled()) {
		logger.debug("Starting WebSocket session url=" + url);
	}
	this.webSocketClient.doHandshake(handler, headers, url).addCallback(
			new ListenableFutureCallback<WebSocketSession>() {
				@Override
				public void onSuccess(WebSocketSession webSocketSession) {
					// WebSocket session ready, SockJS Session not yet
				}
				@Override
				public void onFailure(Throwable ex) {
					future.setException(ex);
				}
			});
	return future;
}
 

@Override
public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
                               Map<String, Object> attributes) throws Exception {
    if (request.getHeaders().containsKey("Sec-WebSocket-Extensions")) {
        request.getHeaders().set("Sec-WebSocket-Extensions", "permessage-deflate");
    }
    String ip = ((ServletServerHttpRequest) request).getServletRequest().getParameter("ip");
    String containerId = ((ServletServerHttpRequest) request).getServletRequest().getParameter("containerId");
    String width = ((ServletServerHttpRequest) request).getServletRequest().getParameter("width");
    String height = ((ServletServerHttpRequest) request).getServletRequest().getParameter("height");
    attributes.put("ip",ip);
    attributes.put("containerId",containerId);
    attributes.put("width",width);
    attributes.put("height",height);
    return super.beforeHandshake(request, response, wsHandler, attributes);
}
 

@Test
public void interceptorsWithAllowedOrigins() {
	WebSocketHandler handler = new TextWebSocketHandler();
	HttpSessionHandshakeInterceptor interceptor = new HttpSessionHandshakeInterceptor();

	this.registration.addHandler(handler, "/foo").addInterceptors(interceptor).setAllowedOrigins("http://mydomain1.com");

	List<Mapping> mappings = this.registration.getMappings();
	assertEquals(1, mappings.size());

	Mapping mapping = mappings.get(0);
	assertEquals(handler, mapping.webSocketHandler);
	assertEquals("/foo", mapping.path);
	assertNotNull(mapping.interceptors);
	assertEquals(2, mapping.interceptors.length);
	assertEquals(interceptor, mapping.interceptors[0]);
	assertEquals(OriginHandshakeInterceptor.class, mapping.interceptors[1].getClass());
}
 

@Test
public void delegateMessageException() throws Exception {
	StubSockJsServiceConfig sockJsConfig = new StubSockJsServiceConfig();
	this.servletRequest.setContent("[\"x\"]".getBytes("UTF-8"));

	WebSocketHandler wsHandler = mock(WebSocketHandler.class);
	TestHttpSockJsSession session = new TestHttpSockJsSession("1", sockJsConfig, wsHandler, null);
	session.delegateConnectionEstablished();

	willThrow(new Exception()).given(wsHandler).handleMessage(session, new TextMessage("x"));

	try {
		XhrReceivingTransportHandler transportHandler = new XhrReceivingTransportHandler();
		transportHandler.initialize(sockJsConfig);
		transportHandler.handleRequest(this.request, this.response, wsHandler, session);
		fail("Expected exception");
	}
	catch (SockJsMessageDeliveryException ex) {
		assertNull(session.getCloseStatus());
	}
}
 

@Test
public void interceptors() {
	WebSocketHandler handler = new TextWebSocketHandler();
	HttpSessionHandshakeInterceptor interceptor = new HttpSessionHandshakeInterceptor();

	this.registration.addHandler(handler, "/foo").addInterceptors(interceptor);

	List<Mapping> mappings = this.registration.getMappings();
	assertEquals(1, mappings.size());

	Mapping mapping = mappings.get(0);
	assertEquals(handler, mapping.webSocketHandler);
	assertEquals("/foo", mapping.path);
	assertNotNull(mapping.interceptors);
	assertEquals(2, mapping.interceptors.length);
	assertEquals(interceptor, mapping.interceptors[0]);
	assertEquals(OriginHandshakeInterceptor.class, mapping.interceptors[1].getClass());
}
 

@Override
public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
        WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {

    // Set ip attribute to WebSocket session
    attributes.put("ip", request.getRemoteAddress());

    // Set servlet request attribute to WebSocket session
    if (request instanceof ServletServerHttpRequest) {
        ServletServerHttpRequest servletServerRequest = (ServletServerHttpRequest) request;
        HttpServletRequest servletRequest = servletServerRequest.getServletRequest();
        attributes.put(UNDERLYING_SERVLET_REQUEST, servletRequest);
    }

    return true;
}
 

@Test
public void openConnection() throws Exception {
	List<String> subprotocols = Arrays.asList("abc");

	TestLifecycleWebSocketClient client = new TestLifecycleWebSocketClient(false);
	WebSocketHandler handler = new TextWebSocketHandler();

	WebSocketConnectionManager manager = new WebSocketConnectionManager(client, handler , "/path/{id}", "123");
	manager.setSubProtocols(subprotocols);
	manager.openConnection();

	WebSocketHttpHeaders expectedHeaders = new WebSocketHttpHeaders();
	expectedHeaders.setSecWebSocketProtocol(subprotocols);

	assertEquals(expectedHeaders, client.headers);
	assertEquals(new URI("/path/123"), client.uri);

	WebSocketHandlerDecorator loggingHandler = (WebSocketHandlerDecorator) client.webSocketHandler;
	assertEquals(LoggingWebSocketHandlerDecorator.class, loggingHandler.getClass());

	assertSame(handler, loggingHandler.getDelegate());
}
 

@Override
protected void connectInternal(TransportRequest request, WebSocketHandler handler, URI receiveUrl,
		HttpHeaders handshakeHeaders, XhrClientSockJsSession session,
		SettableListenableFuture<WebSocketSession> connectFuture) {

	this.actualHandshakeHeaders = handshakeHeaders;
	this.actualSession = session;
}
 

@Test
public void originValueNoMatch() throws Exception {
	Map<String, Object> attributes = new HashMap<>();
	WebSocketHandler wsHandler = Mockito.mock(WebSocketHandler.class);
	this.servletRequest.addHeader(HttpHeaders.ORIGIN, "https://mydomain1.com");
	List<String> allowed = Collections.singletonList("https://mydomain2.com");
	OriginHandshakeInterceptor interceptor = new OriginHandshakeInterceptor(allowed);
	assertFalse(interceptor.beforeHandshake(request, response, wsHandler, attributes));
	assertEquals(servletResponse.getStatus(), HttpStatus.FORBIDDEN.value());
}
 

protected AbstractClientSockJsSession(TransportRequest request, WebSocketHandler handler,
		SettableListenableFuture<WebSocketSession> connectFuture) {

	Assert.notNull(request, "'request' is required");
	Assert.notNull(handler, "'handler' is required");
	Assert.notNull(connectFuture, "'connectFuture' is required");
	this.request = request;
	this.webSocketHandler = handler;
	this.connectFuture = connectFuture;
}
 

@Bean
@SuppressWarnings("deprecation")
public HandlerMapping stompWebSocketHandlerMapping() {
	WebSocketHandler handler = decorateWebSocketHandler(subProtocolWebSocketHandler());
	WebMvcStompEndpointRegistry registry = new WebMvcStompEndpointRegistry(handler,
			getTransportRegistration(), userSessionRegistry(), messageBrokerTaskScheduler());
	registry.setApplicationContext(getApplicationContext());
	registerStompEndpoints(registry);
	return registry.getHandlerMapping();
}
 

@Bean
public HandlerMapping stompWebSocketHandlerMapping() {
	WebSocketHandler handler = decorateWebSocketHandler(subProtocolWebSocketHandler());
	WebMvcStompEndpointRegistry registry = new WebMvcStompEndpointRegistry(
			handler, getTransportRegistration(), messageBrokerTaskScheduler());
	ApplicationContext applicationContext = getApplicationContext();
	if (applicationContext != null) {
		registry.setApplicationContext(applicationContext);
	}
	registerStompEndpoints(registry);
	return registry.getHandlerMapping();
}
 

@Override
protected void connectInternal(final TransportRequest transportRequest, final WebSocketHandler handler,
		final URI receiveUrl, final HttpHeaders handshakeHeaders, final XhrClientSockJsSession session,
		final SettableListenableFuture<WebSocketSession> connectFuture) {

	getTaskExecutor().execute(new Runnable() {
		@Override
		public void run() {
			HttpHeaders httpHeaders = transportRequest.getHttpRequestHeaders();
			XhrRequestCallback requestCallback = new XhrRequestCallback(handshakeHeaders);
			XhrRequestCallback requestCallbackAfterHandshake = new XhrRequestCallback(httpHeaders);
			XhrReceiveExtractor responseExtractor = new XhrReceiveExtractor(session);
			while (true) {
				if (session.isDisconnected()) {
					session.afterTransportClosed(null);
					break;
				}
				try {
					if (logger.isTraceEnabled()) {
						logger.trace("Starting XHR receive request, url=" + receiveUrl);
					}
					getRestTemplate().execute(receiveUrl, HttpMethod.POST, requestCallback, responseExtractor);
					requestCallback = requestCallbackAfterHandshake;
				}
				catch (Throwable ex) {
					if (!connectFuture.isDone()) {
						connectFuture.setException(ex);
					}
					else {
						session.handleTransportError(ex);
						session.afterTransportClosed(new CloseStatus(1006, ex.getMessage()));
					}
					break;
				}
			}
		}
	});
}
 

@Override
protected void addSockJsServiceMapping(MultiValueMap<HttpRequestHandler, String> mappings,
		SockJsService sockJsService, WebSocketHandler handler, String pathPattern) {

	SockJsHttpRequestHandler httpHandler = new SockJsHttpRequestHandler(sockJsService, handler);
	mappings.add(httpHandler, pathPattern);
}
 

@Before
public void setup() {
	this.session = mock(Session.class);
	given(this.session.getUpgradeRequest()).willReturn(Mockito.mock(UpgradeRequest.class));
	given(this.session.getUpgradeResponse()).willReturn(Mockito.mock(UpgradeResponse.class));

	this.webSocketHandler = mock(WebSocketHandler.class);
	this.webSocketSession = new JettyWebSocketSession(null, null);
	this.adapter = new JettyWebSocketHandlerAdapter(this.webSocketHandler, this.webSocketSession);
}
 

@Override
public ListenableFuture<WebSocketSession> doHandshake(WebSocketHandler handler,
		WebSocketHttpHeaders headers, URI uri) {

	this.webSocketHandler = handler;
	this.headers = headers;
	this.uri = uri;
	return new ListenableFutureTask<>(() -> null);
}
 

@Bean
public HandlerMapping stompWebSocketHandlerMapping() {
	WebSocketHandler handler = decorateWebSocketHandler(subProtocolWebSocketHandler());
	WebMvcStompEndpointRegistry registry = new WebMvcStompEndpointRegistry(
			handler, getTransportRegistration(), messageBrokerTaskScheduler());
	ApplicationContext applicationContext = getApplicationContext();
	if (applicationContext != null) {
		registry.setApplicationContext(applicationContext);
	}
	registerStompEndpoints(registry);
	return registry.getHandlerMapping();
}
 

@Test
public void clientLifecycle() throws Exception {
	TestLifecycleWebSocketClient client = new TestLifecycleWebSocketClient(false);
	WebSocketHandler handler = new TextWebSocketHandler();
	WebSocketConnectionManager manager = new WebSocketConnectionManager(client, handler , "/a");

	manager.startInternal();
	assertTrue(client.isRunning());

	manager.stopInternal();
	assertFalse(client.isRunning());
}
 

/**
 * Perform the sub-protocol negotiation based on requested and supported sub-protocols.
 * For the list of supported sub-protocols, this method first checks if the target
 * WebSocketHandler is a {@link SubProtocolCapable} and then also checks if any
 * sub-protocols have been explicitly configured with
 * {@link #setSupportedProtocols(String...)}.
 * @param requestedProtocols the requested sub-protocols
 * @param webSocketHandler the WebSocketHandler that will be used
 * @return the selected protocols or {@code null}
 * @see #determineHandlerSupportedProtocols(WebSocketHandler)
 */
@Nullable
protected String selectProtocol(List<String> requestedProtocols, WebSocketHandler webSocketHandler) {
	List<String> handlerProtocols = determineHandlerSupportedProtocols(webSocketHandler);
	for (String protocol : requestedProtocols) {
		if (handlerProtocols.contains(protocol.toLowerCase())) {
			return protocol;
		}
		if (this.supportedProtocols.contains(protocol.toLowerCase())) {
			return protocol;
		}
	}
	return null;
}
 

@Override
protected void connectInternal(TransportRequest request, WebSocketHandler handler, URI receiveUrl,
		HttpHeaders handshakeHeaders, XhrClientSockJsSession session,
		SettableListenableFuture<WebSocketSession> connectFuture) {

	executeReceiveRequest(request, receiveUrl, handshakeHeaders, session, connectFuture);
}
 

@Test
public void sameOriginMatchWithEmptyAllowedOrigins() throws Exception {
	Map<String, Object> attributes = new HashMap<>();
	WebSocketHandler wsHandler = Mockito.mock(WebSocketHandler.class);
	this.servletRequest.addHeader(HttpHeaders.ORIGIN, "http://mydomain2.com");
	this.servletRequest.setServerName("mydomain2.com");
	OriginHandshakeInterceptor interceptor = new OriginHandshakeInterceptor(Collections.emptyList());
	assertTrue(interceptor.beforeHandshake(request, response, wsHandler, attributes));
	assertNotEquals(servletResponse.getStatus(), HttpStatus.FORBIDDEN.value());
}
 

private void handleRequestAndExpectFailure() throws Exception {
	resetResponse();

	WebSocketHandler wsHandler = mock(WebSocketHandler.class);
	AbstractSockJsSession session = new TestHttpSockJsSession("1", new StubSockJsServiceConfig(), wsHandler, null);

	new XhrReceivingTransportHandler().handleRequest(this.request, this.response, wsHandler, session);

	assertEquals(500, this.servletResponse.getStatus());
	verifyNoMoreInteractions(wsHandler);
}