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.client;
018
019import java.util.List;
020
021import org.springframework.context.Lifecycle;
022import org.springframework.http.HttpHeaders;
023import org.springframework.util.concurrent.ListenableFuture;
024import org.springframework.util.concurrent.ListenableFutureCallback;
025import org.springframework.web.socket.WebSocketHandler;
026import org.springframework.web.socket.WebSocketHttpHeaders;
027import org.springframework.web.socket.WebSocketSession;
028import org.springframework.web.socket.handler.LoggingWebSocketHandlerDecorator;
029
030/**
031 * A WebSocket connection manager that is given a URI, a {@link WebSocketClient}, and a
032 * {@link WebSocketHandler}, connects to a WebSocket server through {@link #start()} and
033 * {@link #stop()} methods. If {@link #setAutoStartup(boolean)} is set to {@code true}
034 * this will be done automatically when the Spring ApplicationContext is refreshed.
035 *
036 * @author Rossen Stoyanchev
037 * @since 4.0
038 */
039public class WebSocketConnectionManager extends ConnectionManagerSupport {
040
041        private final WebSocketClient client;
042
043        private final WebSocketHandler webSocketHandler;
044
045        private WebSocketSession webSocketSession;
046
047        private WebSocketHttpHeaders headers = new WebSocketHttpHeaders();
048
049
050        public WebSocketConnectionManager(WebSocketClient client,
051                        WebSocketHandler webSocketHandler, String uriTemplate, Object... uriVariables) {
052
053                super(uriTemplate, uriVariables);
054                this.client = client;
055                this.webSocketHandler = decorateWebSocketHandler(webSocketHandler);
056        }
057
058
059        /**
060         * Decorate the WebSocketHandler provided to the class constructor.
061         * <p>By default {@link LoggingWebSocketHandlerDecorator} is added.
062         */
063        protected WebSocketHandler decorateWebSocketHandler(WebSocketHandler handler) {
064                return new LoggingWebSocketHandlerDecorator(handler);
065        }
066
067        /**
068         * Set the sub-protocols to use. If configured, specified sub-protocols will be
069         * requested in the handshake through the {@code Sec-WebSocket-Protocol} header. The
070         * resulting WebSocket session will contain the protocol accepted by the server, if
071         * any.
072         */
073        public void setSubProtocols(List<String> protocols) {
074                this.headers.setSecWebSocketProtocol(protocols);
075        }
076
077        /**
078         * Return the configured sub-protocols to use.
079         */
080        public List<String> getSubProtocols() {
081                return this.headers.getSecWebSocketProtocol();
082        }
083
084        /**
085         * Set the origin to use.
086         */
087        public void setOrigin(String origin) {
088                this.headers.setOrigin(origin);
089        }
090
091        /**
092         * Return the configured origin.
093         */
094        public String getOrigin() {
095                return this.headers.getOrigin();
096        }
097
098        /**
099         * Provide default headers to add to the WebSocket handshake request.
100         */
101        public void setHeaders(HttpHeaders headers) {
102                this.headers.clear();
103                this.headers.putAll(headers);
104        }
105
106        /**
107         * Return the default headers for the WebSocket handshake request.
108         */
109        public HttpHeaders getHeaders() {
110                return this.headers;
111        }
112
113
114        @Override
115        public void startInternal() {
116                if (this.client instanceof Lifecycle && !((Lifecycle) client).isRunning()) {
117                        ((Lifecycle) client).start();
118                }
119                super.startInternal();
120        }
121
122        @Override
123        public void stopInternal() throws Exception {
124                if (this.client instanceof Lifecycle && ((Lifecycle) client).isRunning()) {
125                        ((Lifecycle) client).stop();
126                }
127                super.stopInternal();
128        }
129
130        @Override
131        protected void openConnection() {
132                if (logger.isInfoEnabled()) {
133                        logger.info("Connecting to WebSocket at " + getUri());
134                }
135
136                ListenableFuture<WebSocketSession> future =
137                                this.client.doHandshake(this.webSocketHandler, this.headers, getUri());
138
139                future.addCallback(new ListenableFutureCallback<WebSocketSession>() {
140                        @Override
141                        public void onSuccess(WebSocketSession result) {
142                                webSocketSession = result;
143                                logger.info("Successfully connected");
144                        }
145                        @Override
146                        public void onFailure(Throwable ex) {
147                                logger.error("Failed to connect", ex);
148                        }
149                });
150        }
151
152        @Override
153        protected void closeConnection() throws Exception {
154                if (this.webSocketSession != null) {
155                        this.webSocketSession.close();
156                }
157        }
158
159        @Override
160        protected boolean isConnected() {
161                return (this.webSocketSession != null && this.webSocketSession.isOpen());
162        }
163
164}