Add getter & builder methods for OAuth2AuthorizationRequest and OAuth2AuthorizationConsent inside OAuth2AuthenticationContext

Author: MrJovanovic13

oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/authentication/OAuth2AuthorizationCodeRequestAuthenticationContext.java

@@ -21,6 +21,8 @@
 import java.util.function.Consumer;
 
 import org.springframework.lang.Nullable;
+import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
+import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationConsent;
 import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
 import org.springframework.util.Assert;
 
@@ -63,6 +65,27 @@ public RegisteredClient getRegisteredClient() {
 		return get(RegisteredClient.class);
 	}
 
+	/**
+	 * Returns the {@link OAuth2AuthorizationRequest oauth2 authorization request}.
+	 *
+	 * @return the {@link OAuth2AuthorizationRequest}
+	 */
+	@Nullable
+	public OAuth2AuthorizationRequest getOAuth2AuthorizationRequest() {
+		return get(OAuth2AuthorizationRequest.class);
+	}
+
+	/**
+	 * Returns the {@link OAuth2AuthorizationConsent oauth2 authorization consent}.
+	 *
+	 * @return the {@link OAuth2AuthorizationConsent}
+	 */
+	@Nullable
+	public OAuth2AuthorizationConsent getOAuth2AuthorizationConsent() {
+		return get(OAuth2AuthorizationConsent.class);
+	}
+
+
 	/**
 	 * Constructs a new {@link Builder} with the provided {@link OAuth2AuthorizationCodeRequestAuthenticationToken}.
 	 *
@@ -78,6 +101,20 @@ public static Builder with(OAuth2AuthorizationCodeRequestAuthenticationToken aut
 	 */
 	public static final class Builder extends AbstractBuilder<OAuth2AuthorizationCodeRequestAuthenticationContext, Builder> {
 
+		/**
+		 * Associates an attribute.
+		 *
+		 * @param key the key for the attribute
+		 * @param value the value of the attribute
+		 * @return the {@link Builder} for further configuration
+		 */
+		@Override
+		public Builder put(Object key, Object value) {
+			Assert.notNull(key, "key cannot be null");
+			getContext().put(key, value);
+			return getThis();
+		}
+
 		private Builder(OAuth2AuthorizationCodeRequestAuthenticationToken authentication) {
 			super(authentication);
 		}
@@ -92,6 +129,26 @@ public Builder registeredClient(RegisteredClient registeredClient) {
 			return put(RegisteredClient.class, registeredClient);
 		}
 
+		/**
+		 * Sets the {@link OAuth2AuthorizationRequest oauth2 authorization request}.
+		 *
+		 * @param oAuth2AuthorizationRequest the {@link OAuth2AuthorizationRequest}
+		 * @return the {@link Builder} for further configuration
+		 */
+		public Builder oAuth2AuthorizationRequest(OAuth2AuthorizationRequest oAuth2AuthorizationRequest) {
+			return put(OAuth2AuthorizationRequest.class, oAuth2AuthorizationRequest);
+		}
+
+		/**
+		 * Sets the {@link OAuth2AuthorizationConsent oauth2 authorization consent}.
+		 *
+		 * @param oAuth2AuthorizationConsent the {@link OAuth2AuthorizationConsent}
+		 * @return the {@link Builder} for further configuration
+		 */
+		public Builder oAuth2AuthorizationConsent(OAuth2AuthorizationConsent oAuth2AuthorizationConsent) {
+			return put(OAuth2AuthorizationConsent.class, oAuth2AuthorizationConsent);
+		}
+
 		/**
 		 * Builds a new {@link OAuth2AuthorizationCodeRequestAuthenticationContext}.
 		 *

oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/authentication/OAuth2AuthorizationCodeRequestAuthenticationProvider.java

@@ -177,8 +177,8 @@ public Authentication authenticate(Authentication authentication) throws Authent
 		OAuth2AuthorizationCodeRequestAuthenticationContext contextWithAuthorizationRequestAndAuthorizationConsent =
 				OAuth2AuthorizationCodeRequestAuthenticationContext.with(authorizationCodeRequestAuthentication)
 						.registeredClient(registeredClient)
-						.context(context -> context.put(OAuth2AuthorizationRequest.class, authorizationRequest))
-						.context(context -> context.put(OAuth2AuthorizationConsent.class, currentAuthorizationConsent))
+						.oAuth2AuthorizationConsent(currentAuthorizationConsent)
+						.oAuth2AuthorizationRequest(authorizationRequest)
 						.build();
 
 		if (requiresAuthorizationConsent.test(contextWithAuthorizationRequestAndAuthorizationConsent)) {
@@ -296,14 +296,14 @@ private boolean requireAuthorizationConsent(OAuth2AuthorizationCodeRequestAuthen
 			return false;
 		}
 
-		OAuth2AuthorizationRequest authorizationRequest = context.get(OAuth2AuthorizationRequest.class);
+		OAuth2AuthorizationRequest authorizationRequest = context.getOAuth2AuthorizationRequest();
 		// 'openid' scope does not require consent
 		if (authorizationRequest.getScopes().contains(OidcScopes.OPENID) &&
 				authorizationRequest.getScopes().size() == 1) {
 			return false;
 		}
 
-		OAuth2AuthorizationConsent authorizationConsent = context.get(OAuth2AuthorizationConsent.class);
+		OAuth2AuthorizationConsent authorizationConsent = context.getOAuth2AuthorizationConsent();
 		if (authorizationConsent != null &&
 				authorizationConsent.getScopes().containsAll(authorizationRequest.getScopes())) {
 			return false;