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}