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.handler; 018 019import java.util.Map; 020import java.util.concurrent.ConcurrentHashMap; 021 022import org.apache.commons.logging.Log; 023import org.apache.commons.logging.LogFactory; 024 025import org.springframework.beans.factory.BeanFactory; 026import org.springframework.beans.factory.BeanFactoryAware; 027import org.springframework.web.socket.CloseStatus; 028import org.springframework.web.socket.WebSocketHandler; 029import org.springframework.web.socket.WebSocketMessage; 030import org.springframework.web.socket.WebSocketSession; 031 032/** 033 * A {@link WebSocketHandler} that initializes and destroys a {@link WebSocketHandler} 034 * instance for each WebSocket connection and delegates all other methods to it. 035 * 036 * <p>Essentially create an instance of this class once, providing the type of 037 * {@link WebSocketHandler} class to create for each connection, and then pass it to any 038 * API method that expects a {@link WebSocketHandler}. 039 * 040 * <p>If initializing the target {@link WebSocketHandler} type requires a Spring 041 * BeanFactory, then the {@link #setBeanFactory(BeanFactory)} property accordingly. Simply 042 * declaring this class as a Spring bean will do that. Otherwise, {@link WebSocketHandler} 043 * instances of the target type will be created using the default constructor. 044 * 045 * @author Rossen Stoyanchev 046 * @since 4.0 047 */ 048public class PerConnectionWebSocketHandler implements WebSocketHandler, BeanFactoryAware { 049 050 private static final Log logger = LogFactory.getLog(PerConnectionWebSocketHandler.class); 051 052 053 private final BeanCreatingHandlerProvider<WebSocketHandler> provider; 054 055 private final Map<WebSocketSession, WebSocketHandler> handlers = 056 new ConcurrentHashMap<>(); 057 058 private final boolean supportsPartialMessages; 059 060 061 public PerConnectionWebSocketHandler(Class<? extends WebSocketHandler> handlerType) { 062 this(handlerType, false); 063 } 064 065 public PerConnectionWebSocketHandler(Class<? extends WebSocketHandler> handlerType, boolean supportsPartialMessages) { 066 this.provider = new BeanCreatingHandlerProvider<>(handlerType); 067 this.supportsPartialMessages = supportsPartialMessages; 068 } 069 070 071 @Override 072 public void setBeanFactory(BeanFactory beanFactory) { 073 this.provider.setBeanFactory(beanFactory); 074 } 075 076 077 @Override 078 public void afterConnectionEstablished(WebSocketSession session) throws Exception { 079 WebSocketHandler handler = this.provider.getHandler(); 080 this.handlers.put(session, handler); 081 handler.afterConnectionEstablished(session); 082 } 083 084 @Override 085 public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception { 086 getHandler(session).handleMessage(session, message); 087 } 088 089 @Override 090 public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception { 091 getHandler(session).handleTransportError(session, exception); 092 } 093 094 @Override 095 public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception { 096 try { 097 getHandler(session).afterConnectionClosed(session, closeStatus); 098 } 099 finally { 100 destroyHandler(session); 101 } 102 } 103 104 @Override 105 public boolean supportsPartialMessages() { 106 return this.supportsPartialMessages; 107 } 108 109 110 private WebSocketHandler getHandler(WebSocketSession session) { 111 WebSocketHandler handler = this.handlers.get(session); 112 if (handler == null) { 113 throw new IllegalStateException("WebSocketHandler not found for " + session); 114 } 115 return handler; 116 } 117 118 private void destroyHandler(WebSocketSession session) { 119 WebSocketHandler handler = this.handlers.remove(session); 120 try { 121 if (handler != null) { 122 this.provider.destroy(handler); 123 } 124 } 125 catch (Throwable ex) { 126 if (logger.isWarnEnabled()) { 127 logger.warn("Error while destroying " + handler, ex); 128 } 129 } 130 } 131 132 133 @Override 134 public String toString() { 135 return "PerConnectionWebSocketHandlerProxy[handlerType=" + this.provider.getHandlerType() + "]"; 136 } 137 138}