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.lang.Nullable; 026import org.springframework.scheduling.TaskScheduler; 027import org.springframework.util.Assert; 028import org.springframework.util.MultiValueMap; 029import org.springframework.web.HttpRequestHandler; 030import org.springframework.web.servlet.handler.AbstractHandlerMapping; 031import org.springframework.web.socket.WebSocketHandler; 032import org.springframework.web.socket.handler.WebSocketHandlerDecorator; 033import org.springframework.web.socket.messaging.StompSubProtocolErrorHandler; 034import org.springframework.web.socket.messaging.StompSubProtocolHandler; 035import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler; 036import org.springframework.web.socket.server.support.WebSocketHandlerMapping; 037import org.springframework.web.util.UrlPathHelper; 038 039/** 040 * A registry for STOMP over WebSocket endpoints that maps the endpoints with a 041 * {@link org.springframework.web.servlet.HandlerMapping} for use in Spring MVC. 042 * 043 * @author Rossen Stoyanchev 044 * @author Artem Bilan 045 * @since 4.0 046 */ 047public class WebMvcStompEndpointRegistry implements StompEndpointRegistry { 048 049 private final WebSocketHandler webSocketHandler; 050 051 private final TaskScheduler sockJsScheduler; 052 053 private int order = 1; 054 055 @Nullable 056 private UrlPathHelper urlPathHelper; 057 058 private final SubProtocolWebSocketHandler subProtocolWebSocketHandler; 059 060 private final StompSubProtocolHandler stompHandler; 061 062 private final List<WebMvcStompWebSocketEndpointRegistration> registrations = new ArrayList<>(); 063 064 065 public WebMvcStompEndpointRegistry(WebSocketHandler webSocketHandler, 066 WebSocketTransportRegistration transportRegistration, TaskScheduler defaultSockJsTaskScheduler) { 067 068 Assert.notNull(webSocketHandler, "WebSocketHandler is required "); 069 Assert.notNull(transportRegistration, "WebSocketTransportRegistration is required"); 070 071 this.webSocketHandler = webSocketHandler; 072 this.subProtocolWebSocketHandler = unwrapSubProtocolWebSocketHandler(webSocketHandler); 073 074 if (transportRegistration.getSendTimeLimit() != null) { 075 this.subProtocolWebSocketHandler.setSendTimeLimit(transportRegistration.getSendTimeLimit()); 076 } 077 if (transportRegistration.getSendBufferSizeLimit() != null) { 078 this.subProtocolWebSocketHandler.setSendBufferSizeLimit(transportRegistration.getSendBufferSizeLimit()); 079 } 080 if (transportRegistration.getTimeToFirstMessage() != null) { 081 this.subProtocolWebSocketHandler.setTimeToFirstMessage(transportRegistration.getTimeToFirstMessage()); 082 } 083 084 this.stompHandler = new StompSubProtocolHandler(); 085 if (transportRegistration.getMessageSizeLimit() != null) { 086 this.stompHandler.setMessageSizeLimit(transportRegistration.getMessageSizeLimit()); 087 } 088 089 this.sockJsScheduler = defaultSockJsTaskScheduler; 090 } 091 092 private static SubProtocolWebSocketHandler unwrapSubProtocolWebSocketHandler(WebSocketHandler handler) { 093 WebSocketHandler actual = WebSocketHandlerDecorator.unwrap(handler); 094 if (!(actual instanceof SubProtocolWebSocketHandler)) { 095 throw new IllegalArgumentException("No SubProtocolWebSocketHandler in " + handler); 096 } 097 return (SubProtocolWebSocketHandler) actual; 098 } 099 100 101 @Override 102 public StompWebSocketEndpointRegistration addEndpoint(String... paths) { 103 this.subProtocolWebSocketHandler.addProtocolHandler(this.stompHandler); 104 WebMvcStompWebSocketEndpointRegistration registration = 105 new WebMvcStompWebSocketEndpointRegistration(paths, this.webSocketHandler, this.sockJsScheduler); 106 this.registrations.add(registration); 107 return registration; 108 } 109 110 /** 111 * Set the order for the resulting 112 * {@link org.springframework.web.servlet.HandlerMapping} 113 * relative to other handler mappings configured in Spring MVC. 114 * <p>The default value is 1. 115 */ 116 @Override 117 public void setOrder(int order) { 118 this.order = order; 119 } 120 121 protected int getOrder() { 122 return this.order; 123 } 124 125 /** 126 * Set the UrlPathHelper to configure on the {@code HandlerMapping} 127 * used to map handshake requests. 128 */ 129 @Override 130 public void setUrlPathHelper(@Nullable UrlPathHelper urlPathHelper) { 131 this.urlPathHelper = urlPathHelper; 132 } 133 134 @Nullable 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<>(); 154 for (WebMvcStompWebSocketEndpointRegistration registration : this.registrations) { 155 MultiValueMap<HttpRequestHandler, String> mappings = registration.getMappings(); 156 mappings.forEach((httpHandler, patterns) -> { 157 for (String pattern : patterns) { 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}