001/*
002 * Copyright 2002-2014 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.sockjs.transport.handler;
018
019import java.util.ArrayList;
020import java.util.List;
021import java.util.concurrent.atomic.AtomicInteger;
022
023import org.springframework.util.Assert;
024import org.springframework.web.socket.CloseStatus;
025import org.springframework.web.socket.SubProtocolCapable;
026import org.springframework.web.socket.TextMessage;
027import org.springframework.web.socket.WebSocketHandler;
028import org.springframework.web.socket.WebSocketSession;
029import org.springframework.web.socket.handler.TextWebSocketHandler;
030import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
031import org.springframework.web.socket.sockjs.transport.SockJsServiceConfig;
032import org.springframework.web.socket.sockjs.transport.session.WebSocketServerSockJsSession;
033
034/**
035 * An implementation of {@link WebSocketHandler} that adds SockJS messages frames, sends
036 * SockJS heartbeat messages, and delegates lifecycle events and messages to a target
037 * {@link WebSocketHandler}.
038 *
039 * <p>Methods in this class allow exceptions from the wrapped {@link WebSocketHandler} to
040 * propagate. However, any exceptions resulting from SockJS message handling (e.g. while
041 * sending SockJS frames or heartbeat messages) are caught and treated as transport
042 * errors, i.e. routed to the
043 * {@link WebSocketHandler#handleTransportError(WebSocketSession, Throwable)
044 * handleTransportError} method of the wrapped handler and the session closed.
045 *
046 * @author Rossen Stoyanchev
047 * @since 4.0
048 */
049public class SockJsWebSocketHandler extends TextWebSocketHandler implements SubProtocolCapable {
050
051        private final SockJsServiceConfig sockJsServiceConfig;
052
053        private final WebSocketServerSockJsSession sockJsSession;
054
055        private final List<String> subProtocols;
056
057        private final AtomicInteger sessionCount = new AtomicInteger(0);
058
059
060        public SockJsWebSocketHandler(SockJsServiceConfig serviceConfig, WebSocketHandler webSocketHandler,
061                        WebSocketServerSockJsSession sockJsSession) {
062
063                Assert.notNull(serviceConfig, "serviceConfig must not be null");
064                Assert.notNull(webSocketHandler, "webSocketHandler must not be null");
065                Assert.notNull(sockJsSession, "session must not be null");
066
067                this.sockJsServiceConfig = serviceConfig;
068                this.sockJsSession = sockJsSession;
069
070                webSocketHandler = WebSocketHandlerDecorator.unwrap(webSocketHandler);
071                this.subProtocols = ((webSocketHandler instanceof SubProtocolCapable) ?
072                                new ArrayList<String>(((SubProtocolCapable) webSocketHandler).getSubProtocols()) : null);
073        }
074
075        @Override
076        public List<String> getSubProtocols() {
077                return this.subProtocols;
078        }
079
080        protected SockJsServiceConfig getSockJsConfig() {
081                return this.sockJsServiceConfig;
082        }
083
084        @Override
085        public void afterConnectionEstablished(WebSocketSession wsSession) throws Exception {
086                Assert.isTrue(this.sessionCount.compareAndSet(0, 1), "Unexpected connection");
087                this.sockJsSession.initializeDelegateSession(wsSession);
088        }
089
090        @Override
091        public void handleTextMessage(WebSocketSession wsSession, TextMessage message) throws Exception {
092                this.sockJsSession.handleMessage(message, wsSession);
093        }
094
095        @Override
096        public void afterConnectionClosed(WebSocketSession wsSession, CloseStatus status) throws Exception {
097                this.sockJsSession.delegateConnectionClosed(status);
098        }
099
100        @Override
101        public void handleTransportError(WebSocketSession webSocketSession, Throwable exception) throws Exception {
102                this.sockJsSession.delegateError(exception);
103        }
104
105}