Expose WebSocket session info in interceptor

Expose information about the WebSocketSession consistently in all
methods of WebSocketGraphQlInterceptor.

See gh-268

Author: rstoyanchev

spring-graphql/src/main/java/org/springframework/graphql/server/WebGraphQlInterceptor.java

@@ -22,6 +22,7 @@
 
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.graphql.ExecutionGraphQlService;
+import org.springframework.util.Assert;
 
 
 /**
@@ -44,7 +45,8 @@ public interface WebGraphQlInterceptor {
 	/**
 	 * Intercept a request and delegate to the rest of the chain including other
 	 * interceptors and a {@link ExecutionGraphQlService}.
-	 * @param request the request to execute
+	 * @param request the request which may be a {@link WebSocketGraphQlRequest}
+	 * when intercepting a GraphQL request over WebSocket
 	 * @param chain the rest of the chain to execute the request
 	 * @return a {@link Mono} with the response
 	 */
@@ -57,7 +59,13 @@ public interface WebGraphQlInterceptor {
 	 * @return a new interceptor that chains the two
 	 */
 	default WebGraphQlInterceptor andThen(WebGraphQlInterceptor nextInterceptor) {
-		return (request, chain) -> intercept(request, nextRequest -> nextInterceptor.intercept(nextRequest, chain));
+		return (request, chain) -> intercept(request, nextRequest -> {
+			if (request instanceof WebSocketGraphQlRequest) {
+				Assert.isTrue(nextRequest instanceof WebSocketGraphQlRequest,
+						"Expected WebSocketGraphQlRequest but was: " + nextRequest.getClass().getName());
+			}
+			return nextInterceptor.intercept(nextRequest, chain);
+		});
 	}
 
 	/**

spring-graphql/src/main/java/org/springframework/graphql/server/WebSocketGraphQlInterceptor.java

@@ -40,12 +40,12 @@ default Mono<WebGraphQlResponse> intercept(WebGraphQlRequest request, Chain chai
 	 * Handle the {@code "connection_init"} message at the start of a GraphQL over
 	 * WebSocket session and return an optional payload for the
 	 * {@code "connection_ack"} message to send back.
-	 * @param sessionId the id of the WebSocket session
+	 * @param sessionInfo information about the underlying WebSocket session
 	 * @param connectionInitPayload the payload from the {@code "connection_init"} message
 	 * @return the payload for the {@code "connection_ack"}, or empty
 	 */
 	default Mono<Object> handleConnectionInitialization(
-			String sessionId, Map<String, Object> connectionInitPayload) {
+			WebSocketSessionInfo sessionInfo, Map<String, Object> connectionInitPayload) {
 
 		return Mono.empty();
 	}
@@ -55,25 +55,25 @@ default Mono<Object> handleConnectionInitialization(
 	 * subscription stream. The underlying {@link org.reactivestreams.Publisher}
 	 * for the subscription is automatically cancelled. This callback is for any
 	 * additional, or more centralized handling across subscriptions.
-	 * @param sessionId the id of the WebSocket session
+	 * @param sessionInfo information about the underlying WebSocket session
 	 * @param subscriptionId the unique id for the subscription; correlates to the
 	 * {@link WebGraphQlRequest#getId() requestId} from the original {@code "subscribe"}
 	 * message that started the subscription
 	 * @return {@code Mono} for the completion of handling
 	 */
-	default Mono<Void> handleCancelledSubscription(String sessionId, String subscriptionId) {
+	default Mono<Void> handleCancelledSubscription(WebSocketSessionInfo sessionInfo, String subscriptionId) {
 		return Mono.empty();
 	}
 
 	/**
 	 * Invoked when the WebSocket session is closed, from either side.
-	 * @param sessionId the id of the WebSocket session
+	 * @param sessionInfo information about the underlying WebSocket session
 	 * @param statusCode the WebSocket "close" status code
 	 * @param connectionInitPayload the payload from the {@code "connect_init"}
 	 * message received at the start of the connection
 	 */
 	default void handleConnectionClosed(
-			String sessionId, int statusCode, Map<String, Object> connectionInitPayload) {
+			WebSocketSessionInfo sessionInfo, int statusCode, Map<String, Object> connectionInitPayload) {
 	}
 
 }

spring-graphql/src/main/java/org/springframework/graphql/server/WebSocketGraphQlRequest.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2020-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.graphql.server;
+
+
+import java.net.URI;
+import java.util.Locale;
+import java.util.Map;
+
+import org.springframework.http.HttpHeaders;
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+
+
+/**
+ * {@link org.springframework.graphql.server.WebGraphQlRequest} extension for
+ * server handling of GraphQL over WebSocket requests.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.0.0
+ */
+public class WebSocketGraphQlRequest extends WebGraphQlRequest {
+
+	private final WebSocketSessionInfo sessionInfo;
+
+
+	/**
+	 * Create an instance.
+	 * @param uri the URL for the HTTP request or WebSocket handshake
+	 * @param headers the HTTP request headers
+	 * @param body the deserialized content of the GraphQL request
+	 * @param id the id from the GraphQL over WebSocket {@code "subscribe"} message
+	 * @param locale the locale from the HTTP request, if any
+	 * @param sessionInfo the WebSocket session id
+	 */
+	public WebSocketGraphQlRequest(
+			URI uri, HttpHeaders headers, Map<String, Object> body, String id, @Nullable Locale locale,
+			WebSocketSessionInfo sessionInfo) {
+
+		super(uri, headers, body, id, locale);
+		Assert.notNull(sessionInfo, "WebSocketSessionInfo is required");
+		this.sessionInfo = sessionInfo;
+	}
+
+
+	/**
+	 * Return information about the underlying WebSocket session.
+	 */
+	public WebSocketSessionInfo getSessionInfo() {
+		return this.sessionInfo;
+	}
+
+}

spring-graphql/src/main/java/org/springframework/graphql/server/WebSocketSessionInfo.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2002-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.graphql.server;
+
+import java.net.InetSocketAddress;
+import java.net.URI;
+import java.security.Principal;
+import java.util.Map;
+
+import reactor.core.publisher.Mono;
+
+import org.springframework.http.HttpHeaders;
+import org.springframework.lang.Nullable;
+
+/**
+ * Expose information about the underlying WebSocketSession including the
+ * session id, the attributes, and HTTP handshake request.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.0.0
+ */
+public interface WebSocketSessionInfo {
+
+	/**
+	 * Return the id for the WebSocketSession.
+	 */
+	String getId();
+
+	/**
+	 * Return the map with attributes associated with the WebSocket session.
+	 */
+	Map<String, Object> getAttributes();
+
+	/**
+	 * Return the URL for the WebSocket endpoint.
+	 */
+	URI getUri();
+
+	/**
+	 * Return the HTTP headers from the handshake request.
+	 */
+	HttpHeaders getHeaders();
+
+	/**
+	 * Return the principal associated with the handshake request, if any.
+	 */
+	Mono<Principal> getPrincipal();
+
+	/**
+	 * For a server session this is the remote address where the handshake
+	 * request came from. For a client session, it is {@code null}.
+	 */
+	@Nullable
+	InetSocketAddress getRemoteAddress();
+
+}

spring-graphql/src/main/java/org/springframework/graphql/server/webflux/GraphQlWebSocketHandler.java

@@ -16,6 +16,9 @@
 
 package org.springframework.graphql.server.webflux;
 
+import java.net.InetSocketAddress;
+import java.net.URI;
+import java.security.Principal;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
@@ -33,10 +36,12 @@
 import reactor.core.publisher.Mono;
 
 import org.springframework.graphql.server.WebGraphQlHandler;
-import org.springframework.graphql.server.WebGraphQlRequest;
 import org.springframework.graphql.server.WebGraphQlResponse;
 import org.springframework.graphql.server.WebSocketGraphQlInterceptor;
+import org.springframework.graphql.server.WebSocketGraphQlRequest;
+import org.springframework.graphql.server.WebSocketSessionInfo;
 import org.springframework.graphql.server.support.GraphQlWebSocketMessage;
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.codec.CodecConfigurer;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
@@ -106,6 +111,7 @@ public Mono<Void> handle(WebSocketSession session) {
 		}
 
 		// Session state
+		WebSocketSessionInfo sessionInfo = new WebFluxSessionInfo(session);
 		AtomicReference<Map<String, Object>> connectionInitPayloadRef = new AtomicReference<>();
 		Map<String, Subscription> subscriptions = new ConcurrentHashMap<>();
 
@@ -123,7 +129,7 @@ public Mono<Void> handle(WebSocketSession session) {
 						return;
 					}
 					int statusCode = (closeStatus != null ? closeStatus.getCode() : 1005);
-					this.webSocketInterceptor.handleConnectionClosed(session.getId(), statusCode, connectionInitPayload);
+					this.webSocketInterceptor.handleConnectionClosed(sessionInfo, statusCode, connectionInitPayload);
 				})
 				.subscribe();
 
@@ -139,8 +145,8 @@ public Mono<Void> handle(WebSocketSession session) {
 					if (id == null) {
 						return GraphQlStatus.close(session, GraphQlStatus.INVALID_MESSAGE_STATUS);
 					}
-					WebGraphQlRequest request = new WebGraphQlRequest(
-							handshakeInfo.getUri(), handshakeInfo.getHeaders(), payload, id, null);
+					WebSocketGraphQlRequest request = new WebSocketGraphQlRequest(
+							handshakeInfo.getUri(), handshakeInfo.getHeaders(), payload, id, null, sessionInfo);
 					if (logger.isDebugEnabled()) {
 						logger.debug("Executing: " + request);
 					}
@@ -155,15 +161,15 @@ public Mono<Void> handle(WebSocketSession session) {
 						if (subscription != null) {
 							subscription.cancel();
 						}
-						return this.webSocketInterceptor.handleCancelledSubscription(session.getId(), id)
+						return this.webSocketInterceptor.handleCancelledSubscription(sessionInfo, id)
 								.thenMany(Flux.empty());
 					}
 					return Flux.empty();
 				case CONNECTION_INIT:
 					if (!connectionInitPayloadRef.compareAndSet(null, payload)) {
 						return GraphQlStatus.close(session, GraphQlStatus.TOO_MANY_INIT_REQUESTS_STATUS);
 					}
-					return this.webSocketInterceptor.handleConnectionInitialization(session.getId(), payload)
+					return this.webSocketInterceptor.handleConnectionInitialization(sessionInfo, payload)
 							.defaultIfEmpty(Collections.emptyMap())
 							.map(ackPayload -> this.codecDelegate.encodeConnectionAck(session, ackPayload))
 							.flux()
@@ -232,6 +238,46 @@ static <V> Flux<V> close(WebSocketSession session, CloseStatus status) {
 	}
 
 
+	private static class WebFluxSessionInfo implements WebSocketSessionInfo {
+
+		private final WebSocketSession session;
+
+		private WebFluxSessionInfo(WebSocketSession session) {
+			this.session = session;
+		}
+
+		@Override
+		public String getId() {
+			return this.session.getId();
+		}
+
+		@Override
+		public Map<String, Object> getAttributes() {
+			return this.session.getAttributes();
+		}
+
+		@Override
+		public URI getUri() {
+			return this.session.getHandshakeInfo().getUri();
+		}
+
+		@Override
+		public HttpHeaders getHeaders() {
+			return this.session.getHandshakeInfo().getHeaders();
+		}
+
+		@Override
+		public Mono<Principal> getPrincipal() {
+			return this.session.getHandshakeInfo().getPrincipal();
+		}
+
+		@Override
+		public InetSocketAddress getRemoteAddress() {
+			return this.session.getHandshakeInfo().getRemoteAddress();
+		}
+	}
+
+
 	@SuppressWarnings("serial")
 	private static class SubscriptionExistsException extends RuntimeException {
 	}