Add cookies to WebGraphQlRequest

rstoyanchev · rstoyanchev · commit 66b3c8228062 · 2023-03-10T16:04:40.000Z
Closes gh-626
diff --git a/spring-graphql-test/src/main/java/org/springframework/graphql/test/tester/WebGraphQlHandlerGraphQlTransport.java b/spring-graphql-test/src/main/java/org/springframework/graphql/test/tester/WebGraphQlHandlerGraphQlTransport.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -77,14 +77,10 @@ public CodecConfigurer getCodecConfigurer() {
 
 	@Override
 	protected Mono<ExecutionGraphQlResponse> executeInternal(ExecutionGraphQlRequest executionRequest) {
-
 		String id = idGenerator.generateId().toString();
 		Map<String, Object> body = executionRequest.toMap();
-
-		WebGraphQlRequest webExecutionRequest =
-				new WebGraphQlRequest(this.url, this.headers, body, id, null);
-
-		return this.graphQlHandler.handleRequest(webExecutionRequest).cast(ExecutionGraphQlResponse.class);
+		WebGraphQlRequest webRequest = new WebGraphQlRequest(this.url, this.headers, null, body, id, null);
+		return this.graphQlHandler.handleRequest(webRequest).cast(ExecutionGraphQlResponse.class);
 	}
 
 }
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/WebGraphQlRequest.java b/spring-graphql/src/main/java/org/springframework/graphql/server/WebGraphQlRequest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020-2022 the original author or authors.
+ * Copyright 2020-2023 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.
@@ -22,9 +22,13 @@
 
 import org.springframework.graphql.ExecutionGraphQlRequest;
 import org.springframework.graphql.support.DefaultExecutionGraphQlRequest;
+import org.springframework.http.HttpCookie;
 import org.springframework.http.HttpHeaders;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.util.CollectionUtils;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
 import org.springframework.util.StringUtils;
 import org.springframework.web.server.ServerWebInputException;
 import org.springframework.web.util.UriComponents;
@@ -42,21 +46,39 @@
  */
 public class WebGraphQlRequest extends DefaultExecutionGraphQlRequest implements ExecutionGraphQlRequest {
 
+	private static final MultiValueMap<String, HttpCookie> EMPTY_COOKIES =
+			CollectionUtils.unmodifiableMultiValueMap(new LinkedMultiValueMap<>());
+
+
 	private final UriComponents uri;
 
 	private final HttpHeaders headers;
 
+	private final MultiValueMap<String, HttpCookie> cookies;
+
+
+	/**
+	 * Create an instance.
+	 * @deprecated as of 1.1.3 in favor of the constructor with cookies
+	 */
+	@Deprecated
+	public WebGraphQlRequest(URI uri, HttpHeaders headers, Map<String, Object> body, String id, @Nullable Locale locale) {
+		this(uri, headers, null, body, id, locale);
+	}
 
 	/**
 	 * Create an instance.
 	 * @param uri the URL for the HTTP request or WebSocket handshake
 	 * @param headers the HTTP request headers
+	 * @param cookies the request cookies
 	 * @param body the deserialized content of the GraphQL request
 	 * @param id an identifier for the GraphQL request
 	 * @param locale the locale from the HTTP request, if any
+	 * @since 1.1.3
 	 */
 	public WebGraphQlRequest(
-			URI uri, HttpHeaders headers, Map<String, Object> body, String id, @Nullable Locale locale) {
+			URI uri, HttpHeaders headers, @Nullable MultiValueMap<String, HttpCookie> cookies,
+			Map<String, Object> body, String id, @Nullable Locale locale) {
 
 		super(getKey("query", body), getKey("operationName", body), getKey("variables", body),
 				getKey("extensions", body), id, locale);
@@ -66,6 +88,7 @@ public WebGraphQlRequest(
 
 		this.uri = UriComponentsBuilder.fromUri(uri).build(true);
 		this.headers = headers;
+		this.cookies = (cookies != null ? cookies : EMPTY_COOKIES);
 	}
 
 	@SuppressWarnings("unchecked")
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/WebSocketGraphQlRequest.java b/spring-graphql/src/main/java/org/springframework/graphql/server/WebSocketGraphQlRequest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020-2022 the original author or authors.
+ * Copyright 2020-2023 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.
@@ -21,9 +21,11 @@
 import java.util.Locale;
 import java.util.Map;
 
+import org.springframework.http.HttpCookie;
 import org.springframework.http.HttpHeaders;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.util.MultiValueMap;
 
 
 /**
@@ -38,20 +40,34 @@ public class WebSocketGraphQlRequest extends WebGraphQlRequest {
 	private final WebSocketSessionInfo sessionInfo;
 
 
+	/**
+	 * Create an instance.
+	 * @deprecated as of 1.1.3 in favor of the constructor with cookies
+	 */
+	@Deprecated
+	public WebSocketGraphQlRequest(
+			URI uri, HttpHeaders headers, Map<String, Object> body, String id, @Nullable Locale locale,
+			WebSocketSessionInfo sessionInfo) {
+
+		this(uri, headers, null, body, id, locale, sessionInfo);
+	}
+
 	/**
 	 * Create an instance.
 	 * @param uri the URL for the HTTP request or WebSocket handshake
 	 * @param headers the HTTP request headers
+	 * @param cookies the request cookies
 	 * @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
+	 * @since 1.1.3
 	 */
 	public WebSocketGraphQlRequest(
-			URI uri, HttpHeaders headers, Map<String, Object> body, String id, @Nullable Locale locale,
-			WebSocketSessionInfo sessionInfo) {
+			URI uri, HttpHeaders headers, @Nullable MultiValueMap<String, HttpCookie> cookies,
+			Map<String, Object> body, String id, @Nullable Locale locale, WebSocketSessionInfo sessionInfo) {
 
-		super(uri, headers, body, id, locale);
+		super(uri, headers, cookies, body, id, locale);
 		Assert.notNull(sessionInfo, "WebSocketSessionInfo is required");
 		this.sessionInfo = sessionInfo;
 	}
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/webflux/GraphQlHttpHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/server/webflux/GraphQlHttpHandler.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020-2022 the original author or authors.
+ * Copyright 2020-2023 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.
@@ -74,7 +74,8 @@ public Mono<ServerResponse> handleRequest(ServerRequest serverRequest) {
 		return serverRequest.bodyToMono(MAP_PARAMETERIZED_TYPE_REF)
 				.flatMap(body -> {
 					WebGraphQlRequest graphQlRequest = new WebGraphQlRequest(
-							serverRequest.uri(), serverRequest.headers().asHttpHeaders(), body,
+							serverRequest.uri(), serverRequest.headers().asHttpHeaders(),
+							serverRequest.cookies(), body,
 							serverRequest.exchange().getRequest().getId(),
 							serverRequest.exchange().getLocaleContext().getLocale());
 					if (logger.isDebugEnabled()) {
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/webflux/GraphQlWebSocketHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/server/webflux/GraphQlWebSocketHandler.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -148,7 +148,8 @@ public Mono<Void> handle(WebSocketSession session) {
 						return GraphQlStatus.close(session, GraphQlStatus.INVALID_MESSAGE_STATUS);
 					}
 					WebSocketGraphQlRequest request = new WebSocketGraphQlRequest(
-							handshakeInfo.getUri(), handshakeInfo.getHeaders(), payload, id, null, sessionInfo);
+							handshakeInfo.getUri(), handshakeInfo.getHeaders(), handshakeInfo.getCookies(),
+							payload, id, null, sessionInfo);
 					if (logger.isDebugEnabled()) {
 						logger.debug("Executing: " + request);
 					}
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlHttpHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlHttpHandler.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020-2022 the original author or authors.
+ * Copyright 2020-2023 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.
@@ -21,8 +21,8 @@
 import java.util.List;
 import java.util.Map;
 
-
 import jakarta.servlet.ServletException;
+import jakarta.servlet.http.Cookie;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import reactor.core.publisher.Mono;
@@ -31,10 +31,13 @@
 import org.springframework.core.ParameterizedTypeReference;
 import org.springframework.graphql.server.WebGraphQlHandler;
 import org.springframework.graphql.server.WebGraphQlRequest;
+import org.springframework.http.HttpCookie;
 import org.springframework.http.MediaType;
 import org.springframework.util.AlternativeJdkIdGenerator;
 import org.springframework.util.Assert;
 import org.springframework.util.IdGenerator;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
 import org.springframework.web.HttpMediaTypeNotSupportedException;
 import org.springframework.web.server.ServerWebInputException;
 import org.springframework.web.servlet.function.ServerRequest;
@@ -86,8 +89,9 @@ public GraphQlHttpHandler(WebGraphQlHandler graphQlHandler) {
 	public ServerResponse handleRequest(ServerRequest serverRequest) throws ServletException {
 
 		WebGraphQlRequest graphQlRequest = new WebGraphQlRequest(
-				serverRequest.uri(), serverRequest.headers().asHttpHeaders(), readBody(serverRequest),
-				this.idGenerator.generateId().toString(), LocaleContextHolder.getLocale());
+				serverRequest.uri(), serverRequest.headers().asHttpHeaders(), initCookies(serverRequest),
+				readBody(serverRequest), this.idGenerator.generateId().toString(),
+				LocaleContextHolder.getLocale());
 
 		if (logger.isDebugEnabled()) {
 			logger.debug("Executing: " + graphQlRequest);
@@ -107,6 +111,16 @@ public ServerResponse handleRequest(ServerRequest serverRequest) throws ServletE
 		return ServerResponse.async(responseMono);
 	}
 
+	private static MultiValueMap<String, HttpCookie> initCookies(ServerRequest serverRequest) {
+		MultiValueMap<String, Cookie> source = serverRequest.cookies();
+		MultiValueMap<String, HttpCookie> target = new LinkedMultiValueMap<>(source.size());
+		source.values().forEach(cookieList -> cookieList.forEach(cookie -> {
+			HttpCookie httpCookie = new HttpCookie(cookie.getName(), cookie.getValue());
+			target.add(cookie.getName(), httpCookie);
+		}));
+		return target;
+	}
+
 	private static Map<String, Object> readBody(ServerRequest request) throws ServletException {
 		try {
 			return request.body(MAP_PARAMETERIZED_TYPE_REF);
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -204,7 +204,7 @@ private void handleInternal(WebSocketSession session, TextMessage webSocketMessa
 				Assert.notNull(uri, "Expected handshake url");
 				HttpHeaders headers = session.getHandshakeHeaders();
 				WebSocketGraphQlRequest request =
-						new WebSocketGraphQlRequest(uri, headers, payload, id, null, state.getSessionInfo());
+						new WebSocketGraphQlRequest(uri, headers, null, payload, id, null, state.getSessionInfo());
 				if (logger.isDebugEnabled()) {
 					logger.debug("Executing: " + request);
 				}
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/query/QuerydslDataFetcherTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/query/QuerydslDataFetcherTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -50,8 +50,8 @@
 import org.springframework.graphql.data.query.QuerydslDataFetcher.Builder;
 import org.springframework.graphql.data.query.QuerydslDataFetcher.QuerydslBuilderCustomizer;
 import org.springframework.graphql.execution.RuntimeWiringConfigurer;
-import org.springframework.graphql.server.WebGraphQlRequest;
 import org.springframework.graphql.server.WebGraphQlHandler;
+import org.springframework.graphql.server.WebGraphQlRequest;
 import org.springframework.graphql.server.WebGraphQlResponse;
 import org.springframework.http.HttpHeaders;
 import org.springframework.lang.Nullable;
@@ -302,7 +302,8 @@ private static GraphQlSetup initGraphQlSetup(
 
 	private WebGraphQlRequest request(String query) {
 		return new WebGraphQlRequest(
-				URI.create("/"), new HttpHeaders(), Collections.singletonMap("query", query), "1", Locale.ENGLISH);
+				URI.create("/"), new HttpHeaders(), null,
+				Collections.singletonMap("query", query), "1", Locale.ENGLISH);
 	}
 
 
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/query/jpa/QueryByExampleDataFetcherJpaTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/query/jpa/QueryByExampleDataFetcherJpaTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -25,15 +25,14 @@
 import java.util.stream.Collectors;
 
 import javax.sql.DataSource;
-import jakarta.persistence.EntityManagerFactory;
 
 import graphql.schema.DataFetcher;
+import jakarta.persistence.EntityManagerFactory;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Mono;
 
 import org.springframework.beans.factory.annotation.Autowired;
-import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.data.jpa.repository.config.EnableJpaRepositories;
@@ -189,7 +188,8 @@ private static GraphQlSetup initGraphQlSetup(@Nullable QueryByExampleExecutor<?>
 
 	private WebGraphQlRequest request(String query) {
 		return new WebGraphQlRequest(
-				URI.create("/"), new HttpHeaders(), Collections.singletonMap("query", query), "1", null);
+				URI.create("/"), new HttpHeaders(), null,
+				Collections.singletonMap("query", query), "1", null);
 	}
 
 
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/query/mongo/QueryByExampleDataFetcherMongoDbTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/query/mongo/QueryByExampleDataFetcherMongoDbTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -41,8 +41,8 @@
 import org.springframework.data.mongodb.repository.config.EnableMongoRepositories;
 import org.springframework.data.repository.query.QueryByExampleExecutor;
 import org.springframework.graphql.BookSource;
-import org.springframework.graphql.ResponseHelper;
 import org.springframework.graphql.GraphQlSetup;
+import org.springframework.graphql.ResponseHelper;
 import org.springframework.graphql.data.query.QueryByExampleDataFetcher;
 import org.springframework.graphql.execution.RuntimeWiringConfigurer;
 import org.springframework.graphql.server.WebGraphQlHandler;
@@ -185,7 +185,8 @@ private static GraphQlSetup initGraphQlSetup(@Nullable QueryByExampleExecutor<?>
 
 	private WebGraphQlRequest request(String query) {
 		return new WebGraphQlRequest(
-				URI.create("/"), new HttpHeaders(), Collections.singletonMap("query", query), "1", null);
+				URI.create("/"), new HttpHeaders(), null,
+				Collections.singletonMap("query", query), "1", null);
 	}
 
 
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/query/mongo/QueryByExampleDataFetcherReactiveMongoDbTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/query/mongo/QueryByExampleDataFetcherReactiveMongoDbTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -44,8 +44,8 @@
 import org.springframework.graphql.ResponseHelper;
 import org.springframework.graphql.data.query.QueryByExampleDataFetcher;
 import org.springframework.graphql.execution.RuntimeWiringConfigurer;
-import org.springframework.graphql.server.WebGraphQlRequest;
 import org.springframework.graphql.server.WebGraphQlHandler;
+import org.springframework.graphql.server.WebGraphQlRequest;
 import org.springframework.graphql.server.WebGraphQlResponse;
 import org.springframework.http.HttpHeaders;
 import org.springframework.lang.Nullable;
@@ -157,7 +157,8 @@ private static GraphQlSetup initGraphQlSetup(@Nullable ReactiveQueryByExampleExe
 
 	private WebGraphQlRequest request(String query) {
 		return new WebGraphQlRequest(
-				URI.create("/"), new HttpHeaders(), Collections.singletonMap("query", query), "1", null);
+				URI.create("/"), new HttpHeaders(), null,
+				Collections.singletonMap("query", query), "1", null);
 	}
 
 
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/observation/PropagationWebGraphQlInterceptorTests.java b/spring-graphql/src/test/java/org/springframework/graphql/observation/PropagationWebGraphQlInterceptorTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020-2022 the original author or authors.
+ * Copyright 2020-2023 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.
@@ -82,7 +82,9 @@ private void assertThatContextContains(GraphQLContext context, Map<String, Strin
 	WebGraphQlRequest createRequest(Map<String, String> headers) {
 		HttpHeaders httpHeaders = new HttpHeaders();
 		headers.forEach(httpHeaders::set);
-		return new WebGraphQlRequest(URI.create("https://example.org/graphql"), httpHeaders, Map.of("query", "{ notUsed }"), "1", null);
+		return new WebGraphQlRequest(
+				URI.create("https://example.org/graphql"), httpHeaders, null,
+				Map.of("query", "{ notUsed }"), "1", null);
 	}
 
 	private Mono<ExecutionGraphQlResponse> emptyExecutionResult(ExecutionGraphQlRequest request) {
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/server/WebGraphQlHandlerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/server/WebGraphQlHandlerTests.java
diff --git a/spring-graphql/src/test/java/org/springframework/graphql/server/WebGraphQlInterceptorTests.java b/spring-graphql/src/test/java/org/springframework/graphql/server/WebGraphQlInterceptorTests.java

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,5 @@`
`1`	`1`	`/*`
`2`		`- * Copyright 2002-2022 the original author or authors.`
	`2`	`+ * Copyright 2002-2023 the original author or authors.`
`3`	`3`	`*`
`4`	`4`	`* Licensed under the Apache License, Version 2.0 (the "License");`
`5`	`5`	`* you may not use this file except in compliance with the License.`
`@@ -148,7 +148,8 @@ public Mono<Void> handle(WebSocketSession session) {`
`148`	`148`	`return GraphQlStatus.close(session, GraphQlStatus.INVALID_MESSAGE_STATUS);`
`149`	`149`	`}`
`150`	`150`	`WebSocketGraphQlRequest request = new WebSocketGraphQlRequest(`
`151`		`- handshakeInfo.getUri(), handshakeInfo.getHeaders(), payload, id, null, sessionInfo);`
	`151`	`+ handshakeInfo.getUri(), handshakeInfo.getHeaders(), handshakeInfo.getCookies(),`
	`152`	`+ payload, id, null, sessionInfo);`
`152`	`153`	`if (logger.isDebugEnabled()) {`
`153`	`154`	`logger.debug("Executing: " + request);`
`154`	`155`	`}`