From ef33e5ea520a4269358026663a6c39cd59aeb0c0 Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@vmware.com>
Date: Mon, 18 Jan 2021 21:32:02 +0000
Subject: [PATCH] Minor refactoring in WebSocketMessageInput

Decouple from HandshakeInfo which is WebFlux specific.
Rename the class to WebSocketMessageInput since it's per message.
Use more specific name for id field based on the protocol.
---
 ...tInput.java => WebSocketMessageInput.java} | 35 ++++++++-----------
 .../webflux/GraphQLWebSocketHandler.java      |  8 ++---
 2 files changed, 18 insertions(+), 25 deletions(-)
 rename spring-graphql-web/src/main/java/org/springframework/graphql/{WebSocketInput.java => WebSocketMessageInput.java} (56%)

diff --git a/spring-graphql-web/src/main/java/org/springframework/graphql/WebSocketInput.java b/spring-graphql-web/src/main/java/org/springframework/graphql/WebSocketMessageInput.java
similarity index 56%
rename from spring-graphql-web/src/main/java/org/springframework/graphql/WebSocketInput.java
rename to spring-graphql-web/src/main/java/org/springframework/graphql/WebSocketMessageInput.java
index 06a8b95d..f19c4d59 100644
--- a/spring-graphql-web/src/main/java/org/springframework/graphql/WebSocketInput.java
+++ b/spring-graphql-web/src/main/java/org/springframework/graphql/WebSocketMessageInput.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -15,46 +15,39 @@
  */
 package org.springframework.graphql;
 
+import java.net.URI;
 import java.util.Map;
 
-import org.springframework.web.reactive.socket.HandshakeInfo;
+import org.springframework.http.HttpHeaders;
 
 /**
  * Extension of {@link WebInput} that contains a GraphQL subscription received
- * over a WebSocket connection.
+ * as a message over a WebSocket connection.
  */
-public class WebSocketInput extends WebInput {
+public class WebSocketMessageInput extends WebInput {
 
-	private final HandshakeInfo handshakeInfo;
-
-	private final String id;
+	private final String requestId;
 
 
-	public WebSocketInput(HandshakeInfo handshakeInfo, String id, Map<String, Object> payload) {
-		super(handshakeInfo.getUri(), handshakeInfo.getHeaders(), payload);
-		this.handshakeInfo = handshakeInfo;
-		this.id = id;
+	public WebSocketMessageInput(
+			URI uri, HttpHeaders headers, String subscribeId, Map<String, Object> payload) {
+
+		super(uri, headers, payload);
+		this.requestId = subscribeId;
 	}
 
 
-	/**
-	 * Return information about the WebSocket handshake.
-	 */
-	public HandshakeInfo handshakeInfo() {
-		return this.handshakeInfo;
-	}
-
 	/**
 	 * Return the id that will correlate server responses to client requests
 	 * within a multiplexed WebSocket connection.
 	 */
-	public String id() {
-		return this.id;
+	public String requestId() {
+		return this.requestId;
 	}
 
 	@Override
 	public String toString() {
-		return "id='" + id() + "', " + super.toString();
+		return "requestId='" + requestId() + "', " + super.toString();
 	}
 
 }
diff --git a/spring-graphql-web/src/main/java/org/springframework/graphql/webflux/GraphQLWebSocketHandler.java b/spring-graphql-web/src/main/java/org/springframework/graphql/webflux/GraphQLWebSocketHandler.java
index 2e311475..eccef838 100644
--- a/spring-graphql-web/src/main/java/org/springframework/graphql/webflux/GraphQLWebSocketHandler.java
+++ b/spring-graphql-web/src/main/java/org/springframework/graphql/webflux/GraphQLWebSocketHandler.java
@@ -43,7 +43,7 @@ import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.graphql.WebInterceptor;
 import org.springframework.graphql.WebInterceptorExecutionChain;
 import org.springframework.graphql.WebOutput;
-import org.springframework.graphql.WebSocketInput;
+import org.springframework.graphql.WebSocketMessageInput;
 import org.springframework.http.MediaType;
 import org.springframework.http.codec.DecoderHttpMessageReader;
 import org.springframework.http.codec.EncoderHttpMessageWriter;
@@ -152,13 +152,13 @@ public class GraphQLWebSocketHandler implements WebSocketHandler {
 							if (id == null) {
 								return GraphQLStatus.closeWithInvalidMessage(session);
 							}
-							HandshakeInfo handshakeInfo = session.getHandshakeInfo();
-							WebSocketInput input = new WebSocketInput(handshakeInfo, id, getPayload(map));
+							HandshakeInfo info = session.getHandshakeInfo();
+							WebSocketMessageInput input = new WebSocketMessageInput(info.getUri(), info.getHeaders(), id, getPayload(map));
 							if (logger.isDebugEnabled()) {
 								logger.debug("Executing: " + input);
 							}
 							return executionChain.execute(input).flatMapMany(output ->
-									handleWebOutput(session, input.id(), subscriptions, output));
+									handleWebOutput(session, input.requestId(), subscriptions, output));
 						case COMPLETE:
 							if (id != null) {
 								Subscription subscription = subscriptions.remove(id);