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