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