From 23a7c3010ce69081acdf04f775adeae2c6592b56 Mon Sep 17 00:00:00 2001
From: sandmannn <bohdanpukalskyi@gmail.com>
Date: Tue, 4 Jun 2019 23:52:03 +0200
Subject: [PATCH] Added jwt injection for reactive test mocks

Added new implementation of jwt() method that
makes it possible to directly provide a previously
prepared JWT token to WebTestClient mutator.

Fixes: spring-projectsgh-6896
---
 .../server/SecurityMockServerConfigurers.java | 15 +++++++++++
 ...SecurityMockServerConfigurersJwtTests.java | 27 +++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/test/src/main/java/org/springframework/security/test/web/reactive/server/SecurityMockServerConfigurers.java b/test/src/main/java/org/springframework/security/test/web/reactive/server/SecurityMockServerConfigurers.java
index 89f339f1e4..41763e5356 100644
--- a/test/src/main/java/org/springframework/security/test/web/reactive/server/SecurityMockServerConfigurers.java
+++ b/test/src/main/java/org/springframework/security/test/web/reactive/server/SecurityMockServerConfigurers.java
@@ -152,6 +152,21 @@ public class SecurityMockServerConfigurers {
 		return new JwtMutator(jwtBuilder.build());
 	}
 
+	/**
+	 * Updates the ServerWebExchange to establish a {@link SecurityContext} that has a
+	 * {@link JwtAuthenticationToken} for the
+	 * {@link Authentication} and a {@link Jwt} for the
+	 * {@link Authentication#getPrincipal()}. All details are
+	 * declarative and do not require the JWT to be valid.
+	 *
+	 * @param jwt The preliminary constructed {@link Jwt}
+	 * @return the {@link JwtMutator} to further configure or use
+	 * @since 5.2
+	 */
+	public static JwtMutator mockJwt(Jwt jwt) {
+		return new JwtMutator(jwt);
+	}
+
 	public static CsrfMutator csrf() {
 		return new CsrfMutator();
 	}
diff --git a/test/src/test/java/org/springframework/security/test/web/reactive/server/SecurityMockServerConfigurersJwtTests.java b/test/src/test/java/org/springframework/security/test/web/reactive/server/SecurityMockServerConfigurersJwtTests.java
index adcf9a922f..8e763af4cb 100644
--- a/test/src/test/java/org/springframework/security/test/web/reactive/server/SecurityMockServerConfigurersJwtTests.java
+++ b/test/src/test/java/org/springframework/security/test/web/reactive/server/SecurityMockServerConfigurersJwtTests.java
@@ -15,8 +15,12 @@
  */
 package org.springframework.security.test.web.reactive.server;
 
+import java.time.Instant;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
+import java.util.HashMap;
+import java.util.Collections;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -29,6 +33,8 @@ import org.springframework.http.MediaType;
 import org.springframework.security.core.GrantedAuthority;
 import org.springframework.security.core.authority.SimpleGrantedAuthority;
 import org.springframework.security.core.context.SecurityContext;
+import org.springframework.security.oauth2.core.oidc.IdTokenClaimNames;
+import org.springframework.security.oauth2.jwt.Jwt;
 import org.springframework.security.oauth2.server.resource.authentication.JwtAuthenticationToken;
 import org.springframework.security.web.reactive.result.method.annotation.CurrentSecurityContextArgumentResolver;
 import org.springframework.security.web.server.context.SecurityContextServerWebExchangeWebFilter;
@@ -136,4 +142,25 @@ public class SecurityMockServerConfigurersJwtTests extends AbstractMockServerCon
 		assertThat((List<GrantedAuthority>) context.getAuthentication().getAuthorities())
 				.containsOnly(this.authority1);
 	}
+
+	@Test
+	public void mockJwtWhenProvidingPreparedJwtThenProducesJwtAuthentication() {
+		Map<String, Object> claims = new HashMap<>();
+		claims.put(IdTokenClaimNames.SUB, "some_user");
+		Jwt originalToken = new Jwt("token123", Instant.now(), Instant.now().plusSeconds(3600),
+				Collections.singletonMap("header1", "value1"), claims);
+		client
+				.mutateWith(mockJwt(originalToken))
+				.get()
+				.exchange()
+				.expectStatus().isOk();
+
+		SecurityContext context = securityContextController.removeSecurityContext();
+		assertThat(context.getAuthentication()).isInstanceOf(
+				JwtAuthenticationToken.class);
+		JwtAuthenticationToken retrievedToken = (JwtAuthenticationToken) context.getAuthentication();
+		assertThat(retrievedToken.getToken().getSubject()).isEqualTo("some_user");
+		assertThat(retrievedToken.getToken().getTokenValue()).isEqualTo("token123");
+		assertThat(retrievedToken.getToken().getHeaders().get("header1")).isEqualTo("value1");
+	}
 }