Simplify use of headers for SockJsClient requests

Before this change, XhrTransport implementations had to be configured
with the headers to use for HTTP requests other than the initial
handshake.

After this change the handshake headers passed to SockJsClient by
default are used for all other HTTP requests related to the SockJS
connection (e.g. info request, xhr send/receive). A property on
SockJsClient allows restricting the headers to use for other HTTP
requests to a subset of the handshake headers.

Issue: SPR-13254

Author: rstoyanchev

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/client/AbstractXhrTransport.java

@@ -26,7 +26,6 @@
 
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
-import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
 import org.springframework.util.concurrent.ListenableFuture;
 import org.springframework.util.concurrent.SettableListenableFuture;
@@ -61,8 +60,6 @@ public abstract class AbstractXhrTransport implements XhrTransport {
 
 	private HttpHeaders requestHeaders = new HttpHeaders();
 
-	private HttpHeaders xhrSendRequestHeaders = new HttpHeaders();
-
 
 	@Override
 	public List<TransportType> getTransportTypes() {
@@ -97,24 +94,25 @@ public boolean isXhrStreamingDisabled() {
 	/**
 	 * Configure headers to be added to every executed HTTP request.
 	 * @param requestHeaders the headers to add to requests
+	 * @deprecated as of 4.2 in favor of {@link SockJsClient#setHttpHeaderNames}.
 	 */
+	@Deprecated
 	public void setRequestHeaders(HttpHeaders requestHeaders) {
 		this.requestHeaders.clear();
-		this.xhrSendRequestHeaders.clear();
 		if (requestHeaders != null) {
 			this.requestHeaders.putAll(requestHeaders);
-			this.xhrSendRequestHeaders.putAll(requestHeaders);
-			this.xhrSendRequestHeaders.setContentType(MediaType.APPLICATION_JSON);
 		}
 	}
 
+	@Deprecated
 	public HttpHeaders getRequestHeaders() {
 		return this.requestHeaders;
 	}
 
 
 	// Transport methods
 
+	@SuppressWarnings("deprecation")
 	@Override
 	public ListenableFuture<WebSocketSession> connect(TransportRequest request, WebSocketHandler handler) {
 		SettableListenableFuture<WebSocketSession> connectFuture = new SettableListenableFuture<WebSocketSession>();
@@ -128,8 +126,8 @@ public ListenableFuture<WebSocketSession> connect(TransportRequest request, WebS
 		}
 
 		HttpHeaders handshakeHeaders = new HttpHeaders();
-		handshakeHeaders.putAll(request.getHandshakeHeaders());
 		handshakeHeaders.putAll(getRequestHeaders());
+		handshakeHeaders.putAll(request.getHandshakeHeaders());
 
 		connectInternal(request, handler, receiveUrl, handshakeHeaders, session, connectFuture);
 		return connectFuture;
@@ -142,11 +140,17 @@ protected abstract void connectInternal(TransportRequest request, WebSocketHandl
 	// InfoReceiver methods
 
 	@Override
-	public String executeInfoRequest(URI infoUrl) {
+	@SuppressWarnings("deprecation")
+	public String executeInfoRequest(URI infoUrl, HttpHeaders headers) {
 		if (logger.isDebugEnabled()) {
 			logger.debug("Executing SockJS Info request, url=" + infoUrl);
 		}
-		ResponseEntity<String> response = executeInfoRequestInternal(infoUrl);
+		HttpHeaders infoRequestHeaders = new HttpHeaders();
+		infoRequestHeaders.putAll(getRequestHeaders());
+		if (headers != null) {
+			infoRequestHeaders.putAll(headers);
+		}
+		ResponseEntity<String> response = executeInfoRequestInternal(infoUrl, infoRequestHeaders);
 		if (response.getStatusCode() != HttpStatus.OK) {
 			if (logger.isErrorEnabled()) {
 				logger.error("SockJS Info request (url=" + infoUrl + ") failed: " + response);
@@ -159,16 +163,16 @@ public String executeInfoRequest(URI infoUrl) {
 		return response.getBody();
 	}
 
-	protected abstract ResponseEntity<String> executeInfoRequestInternal(URI infoUrl);
+	protected abstract ResponseEntity<String> executeInfoRequestInternal(URI infoUrl, HttpHeaders headers);
 
 	// XhrTransport methods
 
 	@Override
-	public void executeSendRequest(URI url, TextMessage message) {
+	public void executeSendRequest(URI url, HttpHeaders headers, TextMessage message) {
 		if (logger.isTraceEnabled()) {
 			logger.trace("Starting XHR send, url=" + url);
 		}
-		ResponseEntity<String> response = executeSendRequestInternal(url, this.xhrSendRequestHeaders, message);
+		ResponseEntity<String> response = executeSendRequestInternal(url, headers, message);
 		if (response.getStatusCode() != HttpStatus.NO_CONTENT) {
 			if (logger.isErrorEnabled()) {
 				logger.error("XHR send request (url=" + url + ") failed: " + response);
@@ -180,7 +184,8 @@ public void executeSendRequest(URI url, TextMessage message) {
 		}
 	}
 
-	protected abstract ResponseEntity<String> executeSendRequestInternal(URI url, HttpHeaders headers, TextMessage message);
+	protected abstract ResponseEntity<String> executeSendRequestInternal(URI url,
+			HttpHeaders headers, TextMessage message);
 
 
 	@Override

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/client/DefaultTransportRequest.java

@@ -52,6 +52,8 @@ class DefaultTransportRequest implements TransportRequest {
 
 	private final HttpHeaders handshakeHeaders;
 
+	private final HttpHeaders httpRequestHeaders;
+
 	private final Transport transport;
 
 	private final TransportType serverTransportType;
@@ -69,7 +71,8 @@ class DefaultTransportRequest implements TransportRequest {
 	private DefaultTransportRequest fallbackRequest;
 
 
-	public DefaultTransportRequest(SockJsUrlInfo sockJsUrlInfo, HttpHeaders handshakeHeaders,
+	public DefaultTransportRequest(SockJsUrlInfo sockJsUrlInfo,
+			HttpHeaders handshakeHeaders, HttpHeaders httpRequestHeaders,
 			Transport transport, TransportType serverTransportType, SockJsMessageCodec codec) {
 
 		Assert.notNull(sockJsUrlInfo, "'sockJsUrlInfo' is required");
@@ -78,6 +81,7 @@ public DefaultTransportRequest(SockJsUrlInfo sockJsUrlInfo, HttpHeaders handshak
 		Assert.notNull(codec, "'codec' is required");
 		this.sockJsUrlInfo = sockJsUrlInfo;
 		this.handshakeHeaders = (handshakeHeaders != null ? handshakeHeaders : new HttpHeaders());
+		this.httpRequestHeaders = (httpRequestHeaders != null ? httpRequestHeaders : new HttpHeaders());
 		this.transport = transport;
 		this.serverTransportType = serverTransportType;
 		this.codec = codec;
@@ -94,6 +98,11 @@ public HttpHeaders getHandshakeHeaders() {
 		return this.handshakeHeaders;
 	}
 
+	@Override
+	public HttpHeaders getHttpRequestHeaders() {
+		return this.httpRequestHeaders;
+	}
+
 	@Override
 	public URI getTransportUrl() {
 		return this.sockJsUrlInfo.getTransportUrl(this.serverTransportType);

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/client/InfoReceiver.java

@@ -17,6 +17,8 @@
 
 import java.net.URI;
 
+import org.springframework.http.HttpHeaders;
+
 /**
  * A component that can execute the SockJS "Info" request that needs to be
  * performed before the SockJS session starts in order to check server endpoint
@@ -34,10 +36,11 @@ public interface InfoReceiver {
 	/**
 	 * Perform an HTTP request to the SockJS "Info" URL.
 	 * and return the resulting JSON response content, or raise an exception.
-	 *
+	 * <p>Note that as of 4.2 this method accepts a {@code headers} parameter.
 	 * @param infoUrl the URL to obtain SockJS server information from
+	 * @param headers the headers to use for the request
 	 * @return the body of the response
 	 */
-	String executeInfoRequest(URI infoUrl);
+	String executeInfoRequest(URI infoUrl, HttpHeaders headers);
 
 }

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/client/JettyXhrTransport.java

@@ -106,11 +106,12 @@ public boolean isRunning() {
 
 
 	@Override
-	protected void connectInternal(TransportRequest request, WebSocketHandler handler,
+	protected void connectInternal(TransportRequest transportRequest, WebSocketHandler handler,
 			URI url, HttpHeaders handshakeHeaders, XhrClientSockJsSession session,
 			SettableListenableFuture<WebSocketSession> connectFuture) {
 
-		SockJsResponseListener listener = new SockJsResponseListener(url, getRequestHeaders(), session, connectFuture);
+		HttpHeaders httpHeaders = transportRequest.getHttpRequestHeaders();
+		SockJsResponseListener listener = new SockJsResponseListener(url, httpHeaders, session, connectFuture);
 		executeReceiveRequest(url, handshakeHeaders, listener);
 	}
 
@@ -124,8 +125,8 @@ private void executeReceiveRequest(URI url, HttpHeaders headers, SockJsResponseL
 	}
 
 	@Override
-	protected ResponseEntity<String> executeInfoRequestInternal(URI infoUrl) {
-		return executeRequest(infoUrl, HttpMethod.GET, getRequestHeaders(), null);
+	protected ResponseEntity<String> executeInfoRequestInternal(URI infoUrl, HttpHeaders headers) {
+		return executeRequest(infoUrl, HttpMethod.GET, headers, null);
 	}
 
 	@Override

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/client/RestTemplateXhrTransport.java

@@ -94,15 +94,16 @@ public TaskExecutor getTaskExecutor() {
 
 
 	@Override
-	protected void connectInternal(final TransportRequest request, final WebSocketHandler handler,
+	protected void connectInternal(final TransportRequest transportRequest, final WebSocketHandler handler,
 			final URI receiveUrl, final HttpHeaders handshakeHeaders, final XhrClientSockJsSession session,
 			final SettableListenableFuture<WebSocketSession> connectFuture) {
 
 		getTaskExecutor().execute(new Runnable() {
 			@Override
 			public void run() {
+				HttpHeaders httpHeaders = transportRequest.getHttpRequestHeaders();
 				XhrRequestCallback requestCallback = new XhrRequestCallback(handshakeHeaders);
-				XhrRequestCallback requestCallbackAfterHandshake = new XhrRequestCallback(getRequestHeaders());
+				XhrRequestCallback requestCallbackAfterHandshake = new XhrRequestCallback(httpHeaders);
 				XhrReceiveExtractor responseExtractor = new XhrReceiveExtractor(session);
 				while (true) {
 					if (session.isDisconnected()) {
@@ -132,8 +133,8 @@ public void run() {
 	}
 
 	@Override
-	public ResponseEntity<String> executeInfoRequestInternal(URI infoUrl) {
-		RequestCallback requestCallback = new XhrRequestCallback(getRequestHeaders());
+	protected ResponseEntity<String> executeInfoRequestInternal(URI infoUrl, HttpHeaders headers) {
+		RequestCallback requestCallback = new XhrRequestCallback(headers);
 		return this.restTemplate.execute(infoUrl, HttpMethod.GET, requestCallback, textResponseExtractor);
 	}
 

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/client/SockJsClient.java

@@ -78,6 +78,8 @@ public class SockJsClient implements WebSocketClient, Lifecycle {
 
 	private final List<Transport> transports;
 
+	private String[] httpHeaderNames;
+
 	private InfoReceiver infoReceiver;
 
 	private SockJsMessageCodec messageCodec;
@@ -116,6 +118,30 @@ private static InfoReceiver initInfoReceiver(List<Transport> transports) {
 	}
 
 
+	/**
+	 * The names of HTTP headers that should be copied from the handshake headers
+	 * of each call to {@link SockJsClient#doHandshake(WebSocketHandler, WebSocketHttpHeaders, URI)}
+	 * and also used with other HTTP requests issued as part of that SockJS
+	 * connection, e.g. the initial info request, XHR send or receive requests.
+	 *
+	 * <p>By default if this property is not set, all handshake headers are also
+	 * used for other HTTP requests. Set it if you want only a subset of handshake
+	 * headers (e.g. auth headers) to be used for other HTTP requests.
+	 *
+	 * @param httpHeaderNames HTTP header names
+	 */
+	public void setHttpHeaderNames(String... httpHeaderNames) {
+		this.httpHeaderNames = httpHeaderNames;
+	}
+
+	/**
+	 * The configured HTTP header names to be copied from the handshake
+	 * headers and also included in other HTTP requests.
+	 */
+	public String[] getHttpHeaderNames() {
+		return this.httpHeaderNames;
+	}
+
 	/**
 	 * Configure the {@code InfoReceiver} to use to perform the SockJS "Info"
 	 * request before the SockJS session starts.
@@ -225,7 +251,7 @@ public final ListenableFuture<WebSocketSession> doHandshake(
 		SettableListenableFuture<WebSocketSession> connectFuture = new SettableListenableFuture<WebSocketSession>();
 		try {
 			SockJsUrlInfo sockJsUrlInfo = new SockJsUrlInfo(url);
-			ServerInfo serverInfo = getServerInfo(sockJsUrlInfo);
+			ServerInfo serverInfo = getServerInfo(sockJsUrlInfo, getHttpRequestHeaders(headers));
 			createRequest(sockJsUrlInfo, headers, serverInfo).connect(handler, connectFuture);
 		}
 		catch (Throwable exception) {
@@ -237,12 +263,27 @@ public final ListenableFuture<WebSocketSession> doHandshake(
 		return connectFuture;
 	}
 
-	private ServerInfo getServerInfo(SockJsUrlInfo sockJsUrlInfo) {
+	private HttpHeaders getHttpRequestHeaders(HttpHeaders webSocketHttpHeaders) {
+		if (getHttpHeaderNames() == null) {
+			return webSocketHttpHeaders;
+		}
+		else {
+			HttpHeaders httpHeaders = new HttpHeaders();
+			for (String name : getHttpHeaderNames()) {
+				if (webSocketHttpHeaders.containsKey(name)) {
+					httpHeaders.put(name, webSocketHttpHeaders.get(name));
+				}
+			}
+			return httpHeaders;
+		}
+	}
+
+	private ServerInfo getServerInfo(SockJsUrlInfo sockJsUrlInfo, HttpHeaders headers) {
 		URI infoUrl = sockJsUrlInfo.getInfoUrl();
 		ServerInfo info = this.serverInfoCache.get(infoUrl);
 		if (info == null) {
 			long start = System.currentTimeMillis();
-			String response = this.infoReceiver.executeInfoRequest(infoUrl);
+			String response = this.infoReceiver.executeInfoRequest(infoUrl, headers);
 			long infoRequestTime = System.currentTimeMillis() - start;
 			info = new ServerInfo(response, infoRequestTime);
 			this.serverInfoCache.put(infoUrl, info);
@@ -255,7 +296,8 @@ private DefaultTransportRequest createRequest(SockJsUrlInfo urlInfo, HttpHeaders
 		for (Transport transport : this.transports) {
 			for (TransportType type : transport.getTransportTypes()) {
 				if (serverInfo.isWebSocketEnabled() || !TransportType.WEBSOCKET.equals(type)) {
-					requests.add(new DefaultTransportRequest(urlInfo, headers, transport, type, getMessageCodec()));
+					requests.add(new DefaultTransportRequest(urlInfo, headers, getHttpRequestHeaders(headers),
+							transport, type, getMessageCodec()));
 				}
 			}
 		}

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/client/TransportRequest.java

@@ -47,6 +47,13 @@ public interface TransportRequest {
 	 */
 	HttpHeaders getHandshakeHeaders();
 
+	/**
+	 * Return the headers to add to all other HTTP requests besides the handshake
+	 * request such XHR receive and send requests.
+	 * @since 4.2
+	 */
+	HttpHeaders getHttpRequestHeaders();
+
 	/**
 	 * Return the transport URL for the given transport.
 	 * For an {@link XhrTransport} this is the URL for receiving messages.