001/*
002 * Copyright 2002-2015 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 * BeanFctory, 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<WebSocketSession, WebSocketHandler>();
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<WebSocketHandler>(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}