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