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}