001/* 002 * Copyright 2002-2018 the original author or authors. 003 * 004 * Licensed under the Apache License, Version 2.0 (the "License"); 005 * you may not use this file except in compliance with the License. 006 * You may obtain a copy of the License at 007 * 008 * https://www.apache.org/licenses/LICENSE-2.0 009 * 010 * Unless required by applicable law or agreed to in writing, software 011 * distributed under the License is distributed on an "AS IS" BASIS, 012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 013 * See the License for the specific language governing permissions and 014 * limitations under the License. 015 */ 016 017package org.springframework.web.socket.config.annotation; 018 019import java.util.ArrayList; 020import java.util.LinkedHashMap; 021import java.util.List; 022import java.util.Map; 023 024import org.springframework.context.ApplicationContext; 025import org.springframework.scheduling.TaskScheduler; 026import org.springframework.util.Assert; 027import org.springframework.util.MultiValueMap; 028import org.springframework.web.HttpRequestHandler; 029import org.springframework.web.servlet.handler.AbstractHandlerMapping; 030import org.springframework.web.socket.WebSocketHandler; 031import org.springframework.web.socket.handler.WebSocketHandlerDecorator; 032import org.springframework.web.socket.messaging.StompSubProtocolErrorHandler; 033import org.springframework.web.socket.messaging.StompSubProtocolHandler; 034import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler; 035import org.springframework.web.socket.server.support.WebSocketHandlerMapping; 036import org.springframework.web.util.UrlPathHelper; 037 038/** 039 * A registry for STOMP over WebSocket endpoints that maps the endpoints with a 040 * {@link org.springframework.web.servlet.HandlerMapping} for use in Spring MVC. 041 * 042 * @author Rossen Stoyanchev 043 * @author Artem Bilan 044 * @since 4.0 045 */ 046public class WebMvcStompEndpointRegistry implements StompEndpointRegistry { 047 048 private final WebSocketHandler webSocketHandler; 049 050 private final TaskScheduler sockJsScheduler; 051 052 private int order = 1; 053 054 private UrlPathHelper urlPathHelper; 055 056 private final SubProtocolWebSocketHandler subProtocolWebSocketHandler; 057 058 private final StompSubProtocolHandler stompHandler; 059 060 private final List<WebMvcStompWebSocketEndpointRegistration> registrations = 061 new ArrayList<WebMvcStompWebSocketEndpointRegistration>(); 062 063 064 @SuppressWarnings("deprecation") 065 public WebMvcStompEndpointRegistry(WebSocketHandler webSocketHandler, 066 WebSocketTransportRegistration transportRegistration, 067 org.springframework.messaging.simp.user.UserSessionRegistry userSessionRegistry, 068 TaskScheduler defaultSockJsTaskScheduler) { 069 070 Assert.notNull(webSocketHandler, "WebSocketHandler is required "); 071 Assert.notNull(transportRegistration, "WebSocketTransportRegistration is required"); 072 073 this.webSocketHandler = webSocketHandler; 074 this.subProtocolWebSocketHandler = unwrapSubProtocolWebSocketHandler(webSocketHandler); 075 076 if (transportRegistration.getSendTimeLimit() != null) { 077 this.subProtocolWebSocketHandler.setSendTimeLimit(transportRegistration.getSendTimeLimit()); 078 } 079 if (transportRegistration.getSendBufferSizeLimit() != null) { 080 this.subProtocolWebSocketHandler.setSendBufferSizeLimit(transportRegistration.getSendBufferSizeLimit()); 081 } 082 083 this.stompHandler = new StompSubProtocolHandler(); 084 this.stompHandler.setUserSessionRegistry(userSessionRegistry); 085 086 if (transportRegistration.getMessageSizeLimit() != null) { 087 this.stompHandler.setMessageSizeLimit(transportRegistration.getMessageSizeLimit()); 088 } 089 090 this.sockJsScheduler = defaultSockJsTaskScheduler; 091 } 092 093 private static SubProtocolWebSocketHandler unwrapSubProtocolWebSocketHandler(WebSocketHandler handler) { 094 WebSocketHandler actual = WebSocketHandlerDecorator.unwrap(handler); 095 if (!(actual instanceof SubProtocolWebSocketHandler)) { 096 throw new IllegalArgumentException("No SubProtocolWebSocketHandler in " + handler); 097 } 098 return (SubProtocolWebSocketHandler) actual; 099 } 100 101 102 @Override 103 public StompWebSocketEndpointRegistration addEndpoint(String... paths) { 104 this.subProtocolWebSocketHandler.addProtocolHandler(this.stompHandler); 105 WebMvcStompWebSocketEndpointRegistration registration = 106 new WebMvcStompWebSocketEndpointRegistration(paths, this.webSocketHandler, this.sockJsScheduler); 107 this.registrations.add(registration); 108 return registration; 109 } 110 111 /** 112 * Set the order for the resulting 113 * {@link org.springframework.web.servlet.HandlerMapping} 114 * relative to other handler mappings configured in Spring MVC. 115 * <p>The default value is 1. 116 */ 117 @Override 118 public void setOrder(int order) { 119 this.order = order; 120 } 121 122 protected int getOrder() { 123 return this.order; 124 } 125 126 /** 127 * Set the UrlPathHelper to configure on the {@code HandlerMapping} 128 * used to map handshake requests. 129 */ 130 @Override 131 public void setUrlPathHelper(UrlPathHelper urlPathHelper) { 132 this.urlPathHelper = urlPathHelper; 133 } 134 135 protected UrlPathHelper getUrlPathHelper() { 136 return this.urlPathHelper; 137 } 138 139 @Override 140 public WebMvcStompEndpointRegistry setErrorHandler(StompSubProtocolErrorHandler errorHandler) { 141 this.stompHandler.setErrorHandler(errorHandler); 142 return this; 143 } 144 145 protected void setApplicationContext(ApplicationContext applicationContext) { 146 this.stompHandler.setApplicationEventPublisher(applicationContext); 147 } 148 149 /** 150 * Return a handler mapping with the mapped ViewControllers. 151 */ 152 public AbstractHandlerMapping getHandlerMapping() { 153 Map<String, Object> urlMap = new LinkedHashMap<String, Object>(); 154 for (WebMvcStompWebSocketEndpointRegistration registration : this.registrations) { 155 MultiValueMap<HttpRequestHandler, String> mappings = registration.getMappings(); 156 for (HttpRequestHandler httpHandler : mappings.keySet()) { 157 for (String pattern : mappings.get(httpHandler)) { 158 urlMap.put(pattern, httpHandler); 159 } 160 } 161 } 162 WebSocketHandlerMapping hm = new WebSocketHandlerMapping(); 163 hm.setUrlMap(urlMap); 164 hm.setOrder(this.order); 165 if (this.urlPathHelper != null) { 166 hm.setUrlPathHelper(this.urlPathHelper); 167 } 168 return hm; 169 } 170 171}