001/*
002 * Copyright 2002-2019 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.Map;
020import javax.servlet.ServletContext;
021
022import org.springframework.context.Lifecycle;
023import org.springframework.http.server.ServerHttpRequest;
024import org.springframework.http.server.ServerHttpResponse;
025import org.springframework.util.Assert;
026import org.springframework.web.context.ServletContextAware;
027import org.springframework.web.socket.CloseStatus;
028import org.springframework.web.socket.WebSocketHandler;
029import org.springframework.web.socket.server.HandshakeFailureException;
030import org.springframework.web.socket.server.HandshakeHandler;
031import org.springframework.web.socket.sockjs.SockJsException;
032import org.springframework.web.socket.sockjs.SockJsTransportFailureException;
033import org.springframework.web.socket.sockjs.transport.SockJsSession;
034import org.springframework.web.socket.sockjs.transport.SockJsSessionFactory;
035import org.springframework.web.socket.sockjs.transport.TransportHandler;
036import org.springframework.web.socket.sockjs.transport.TransportType;
037import org.springframework.web.socket.sockjs.transport.session.AbstractSockJsSession;
038import org.springframework.web.socket.sockjs.transport.session.WebSocketServerSockJsSession;
039
040/**
041 * WebSocket-based {@link TransportHandler}. Uses {@link SockJsWebSocketHandler} and
042 * {@link WebSocketServerSockJsSession} to add SockJS processing.
043 *
044 * <p>Also implements {@link HandshakeHandler} to support raw WebSocket communication at
045 * SockJS URL "/websocket".
046 *
047 * @author Rossen Stoyanchev
048 * @since 4.0
049 */
050public class WebSocketTransportHandler extends AbstractTransportHandler
051                implements SockJsSessionFactory, HandshakeHandler, Lifecycle, ServletContextAware {
052
053        private final HandshakeHandler handshakeHandler;
054
055        private volatile boolean running;
056
057
058        public WebSocketTransportHandler(HandshakeHandler handshakeHandler) {
059                Assert.notNull(handshakeHandler, "HandshakeHandler must not be null");
060                this.handshakeHandler = handshakeHandler;
061        }
062
063
064        @Override
065        public TransportType getTransportType() {
066                return TransportType.WEBSOCKET;
067        }
068
069        public HandshakeHandler getHandshakeHandler() {
070                return this.handshakeHandler;
071        }
072
073        @Override
074        public void setServletContext(ServletContext servletContext) {
075                if (this.handshakeHandler instanceof ServletContextAware) {
076                        ((ServletContextAware) this.handshakeHandler).setServletContext(servletContext);
077                }
078        }
079
080
081        @Override
082        public void start() {
083                if (!isRunning()) {
084                        this.running = true;
085                        if (this.handshakeHandler instanceof Lifecycle) {
086                                ((Lifecycle) this.handshakeHandler).start();
087                        }
088                }
089        }
090
091        @Override
092        public void stop() {
093                if (isRunning()) {
094                        this.running = false;
095                        if (this.handshakeHandler instanceof Lifecycle) {
096                                ((Lifecycle) this.handshakeHandler).stop();
097                        }
098                }
099        }
100
101        @Override
102        public boolean isRunning() {
103                return this.running;
104        }
105
106
107        @Override
108        public boolean checkSessionType(SockJsSession session) {
109                return (session instanceof WebSocketServerSockJsSession);
110        }
111
112        @Override
113        public AbstractSockJsSession createSession(String id, WebSocketHandler handler, Map<String, Object> attrs) {
114                return new WebSocketServerSockJsSession(id, getServiceConfig(), handler, attrs);
115        }
116
117        @Override
118        public void handleRequest(ServerHttpRequest request, ServerHttpResponse response,
119                        WebSocketHandler wsHandler, SockJsSession wsSession) throws SockJsException {
120
121                WebSocketServerSockJsSession sockJsSession = (WebSocketServerSockJsSession) wsSession;
122                try {
123                        wsHandler = new SockJsWebSocketHandler(getServiceConfig(), wsHandler, sockJsSession);
124                        this.handshakeHandler.doHandshake(request, response, wsHandler, sockJsSession.getAttributes());
125                }
126                catch (Throwable ex) {
127                        sockJsSession.tryCloseWithSockJsTransportError(ex, CloseStatus.SERVER_ERROR);
128                        throw new SockJsTransportFailureException("WebSocket handshake failure", wsSession.getId(), ex);
129                }
130        }
131
132        @Override
133        public boolean doHandshake(ServerHttpRequest request, ServerHttpResponse response,
134                        WebSocketHandler handler, Map<String, Object> attributes) throws HandshakeFailureException {
135
136                return this.handshakeHandler.doHandshake(request, response, handler, attributes);
137        }
138}