Expire OAuth2AuthorizationRequest when saving to the session

oauth2/oauth2-client/src/main/java/org/springframework/security/oauth2/client/web/HttpSessionOAuth2AuthorizationRequestRepository.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 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.
@@ -16,23 +16,33 @@
 
 package org.springframework.security.oauth2.client.web;
 
+import java.io.Serializable;
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Map.Entry;
 
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import javax.servlet.http.HttpSession;
 
+import org.springframework.security.core.SpringSecurityCoreVersion;
 import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
 import org.springframework.security.oauth2.core.endpoint.OAuth2ParameterNames;
 import org.springframework.util.Assert;
 
 /**
  * An implementation of an {@link AuthorizationRequestRepository} that stores
  * {@link OAuth2AuthorizationRequest} in the {@code HttpSession}.
+ * <p>
+ * <b>NOTE:</b> {@link OAuth2AuthorizationRequest}s expire after two minutes, the default
+ * duration can be configured via {@link #setAuthorizationRequestTimeToLive(Duration)}.
  *
  * @author Joe Grandja
  * @author Rob Winch
+ * @author Craig Andrews
  * @since 5.0
  * @see AuthorizationRequestRepository
  * @see OAuth2AuthorizationRequest
@@ -45,15 +55,22 @@ public final class HttpSessionOAuth2AuthorizationRequestRepository
 
 	private final String sessionAttributeName = DEFAULT_AUTHORIZATION_REQUEST_ATTR_NAME;
 
+	private Clock clock = Clock.systemUTC();
+
+	private Duration authorizationRequestTimeToLive = Duration.ofSeconds(120);
+
+	private int maxActiveAuthorizationRequestsPerSession = 10;
+
 	@Override
 	public OAuth2AuthorizationRequest loadAuthorizationRequest(HttpServletRequest request) {
 		Assert.notNull(request, "request cannot be null");
 		String stateParameter = this.getStateParameter(request);
 		if (stateParameter == null) {
 			return null;
 		}
-		Map<String, OAuth2AuthorizationRequest> authorizationRequests = this.getAuthorizationRequests(request);
-		return authorizationRequests.get(stateParameter);
+		Map<String, OAuth2AuthorizationRequestReference> authorizationRequests = this.getAuthorizationRequests(request);
+		OAuth2AuthorizationRequestReference wrappedWithCreated = authorizationRequests.get(stateParameter);
+		return (wrappedWithCreated != null) ? wrappedWithCreated.wrapped : null;
 	}
 
 	@Override
@@ -67,8 +84,14 @@ public void saveAuthorizationRequest(OAuth2AuthorizationRequest authorizationReq
 		}
 		String state = authorizationRequest.getState();
 		Assert.hasText(state, "authorizationRequest.state cannot be empty");
-		Map<String, OAuth2AuthorizationRequest> authorizationRequests = this.getAuthorizationRequests(request);
-		authorizationRequests.put(state, authorizationRequest);
+		Map<String, OAuth2AuthorizationRequestReference> authorizationRequests = this.getAuthorizationRequests(request);
+		authorizationRequests.put(state, new OAuth2AuthorizationRequestReference(authorizationRequest,
+				this.clock.instant().plus(this.authorizationRequestTimeToLive)));
+		if (authorizationRequests.size() > this.maxActiveAuthorizationRequestsPerSession) {
+			authorizationRequests.entrySet().stream()
+					.sorted((e, f) -> e.getValue().expiresAt.compareTo(f.getValue().expiresAt)).findFirst()
+					.map(Entry::getKey).ifPresent(authorizationRequests::remove);
+		}
 		request.getSession().setAttribute(this.sessionAttributeName, authorizationRequests);
 	}
 
@@ -79,15 +102,16 @@ public OAuth2AuthorizationRequest removeAuthorizationRequest(HttpServletRequest
 		if (stateParameter == null) {
 			return null;
 		}
-		Map<String, OAuth2AuthorizationRequest> authorizationRequests = this.getAuthorizationRequests(request);
-		OAuth2AuthorizationRequest originalRequest = authorizationRequests.remove(stateParameter);
+		Map<String, OAuth2AuthorizationRequestReference> authorizationRequests = this.getAuthorizationRequests(request);
+		OAuth2AuthorizationRequestReference wrappedWithCreatedOriginalRequest = authorizationRequests
+				.remove(stateParameter);
 		if (!authorizationRequests.isEmpty()) {
 			request.getSession().setAttribute(this.sessionAttributeName, authorizationRequests);
 		}
 		else {
 			request.getSession().removeAttribute(this.sessionAttributeName);
 		}
-		return originalRequest;
+		return (wrappedWithCreatedOriginalRequest != null) ? wrappedWithCreatedOriginalRequest.wrapped : null;
 	}
 
 	@Override
@@ -113,14 +137,70 @@ private String getStateParameter(HttpServletRequest request) {
 	 * @return a non-null and mutable map of {@link OAuth2AuthorizationRequest#getState()}
 	 * to an {@link OAuth2AuthorizationRequest}.
 	 */
-	private Map<String, OAuth2AuthorizationRequest> getAuthorizationRequests(HttpServletRequest request) {
+	private Map<String, OAuth2AuthorizationRequestReference> getAuthorizationRequests(HttpServletRequest request) {
 		HttpSession session = request.getSession(false);
-		Map<String, OAuth2AuthorizationRequest> authorizationRequests = (session != null)
-				? (Map<String, OAuth2AuthorizationRequest>) session.getAttribute(this.sessionAttributeName) : null;
+		Map<String, OAuth2AuthorizationRequestReference> authorizationRequests = (session != null)
+				? (Map<String, OAuth2AuthorizationRequestReference>) session.getAttribute(this.sessionAttributeName)
+				: null;
 		if (authorizationRequests == null) {
 			return new HashMap<>();
 		}
+		// remove expired entries
+		authorizationRequests.entrySet().removeIf((entry) -> entry.getValue().expiresAt.isBefore(this.clock.instant()));
 		return authorizationRequests;
 	}
 
+	/**
+	 * Sets the {@link Clock} used in {@link Instant#now(Clock)} when setting the instant
+	 * created for {@link OAuth2AuthorizationRequest}.
+	 * @param clock the clock
+	 * @since 5.5
+	 */
+	void setClock(Clock clock) {
+		Assert.notNull(clock, "clock cannot be null");
+		this.clock = clock;
+	}
+
+	/**
+	 * Sets the {@link Duration} for which {@link OAuth2AuthorizationRequest} should
+	 * expire.
+	 * @param authorizationRequestTimeToLive the {@link Duration} a
+	 * {@link OAuth2AuthorizationRequest} is considered not expired. Must not be negative.
+	 * @since 5.5
+	 */
+	public void setAuthorizationRequestTimeToLive(Duration authorizationRequestTimeToLive) {
+		Assert.notNull(authorizationRequestTimeToLive, "oAuth2AuthorizationRequestExpiresIn cannot be null");
+		Assert.state(!authorizationRequestTimeToLive.isNegative(),
+				"oAuth2AuthorizationRequestExpiresIn cannot be negative");
+		this.authorizationRequestTimeToLive = authorizationRequestTimeToLive;
+	}
+
+	/**
+	 * Sets the maximum number of {@link OAuth2AuthorizationRequest} that can be
+	 * stored/active for a session. If the maximum number are present in a session when an
+	 * attempt is made to save another one, then the oldest will be removed.
+	 * @param maxActiveAuthorizationRequests must not be negative.
+	 */
+	public void setMaxActiveAuthorizationRequestsPerSession(int maxActiveAuthorizationRequestsPerSession) {
+		Assert.state(maxActiveAuthorizationRequestsPerSession > 0,
+				"maxActiveAuthorizationRequestsPerSession must be greater than zero");
+		this.maxActiveAuthorizationRequestsPerSession = maxActiveAuthorizationRequestsPerSession;
+	}
+
+	private static final class OAuth2AuthorizationRequestReference implements Serializable {
+
+		private static final long serialVersionUID = SpringSecurityCoreVersion.SERIAL_VERSION_UID;
+
+		private final Instant expiresAt;
+
+		private final OAuth2AuthorizationRequest wrapped;
+
+		private OAuth2AuthorizationRequestReference(OAuth2AuthorizationRequest wrapped, Instant created) {
+			Assert.notNull(wrapped, "wrapped cannot be null");
+			this.expiresAt = created;
+			this.wrapped = wrapped;
+		}
+
+	}
+
 }

oauth2/oauth2-client/src/test/java/org/springframework/security/oauth2/client/web/HttpSessionOAuth2AuthorizationRequestRepositoryTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 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.
@@ -16,6 +16,10 @@
 
 package org.springframework.security.oauth2.client.web;
 
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+import java.time.ZoneId;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -36,6 +40,7 @@
  * Tests for {@link HttpSessionOAuth2AuthorizationRequestRepository}.
  *
  * @author Joe Grandja
+ * @author Craig Andrews
  */
 @RunWith(MockitoJUnitRunner.class)
 public class HttpSessionOAuth2AuthorizationRequestRepositoryTests {
@@ -237,6 +242,62 @@ public void removeAuthorizationRequestWhenNotSavedThenNotRemoved() {
 		assertThat(removedAuthorizationRequest).isNull();
 	}
 
+	@Test
+	public void removeAuthorizationRequestWhenExpired() {
+		final Duration expiresIn = Duration.ofMinutes(2);
+		this.authorizationRequestRepository.setAuthorizationRequestTimeToLive(expiresIn);
+		this.authorizationRequestRepository.setClock(Clock.fixed(Instant.ofEpochMilli(0), ZoneId.systemDefault()));
+		MockHttpServletRequest request = new MockHttpServletRequest();
+		MockHttpServletResponse response = new MockHttpServletResponse();
+		String state1 = "state-1122";
+		OAuth2AuthorizationRequest authorizationRequest1 = createAuthorizationRequest().state(state1).build();
+		this.authorizationRequestRepository.saveAuthorizationRequest(authorizationRequest1, request, response);
+		String state2 = "state-3344";
+		this.authorizationRequestRepository
+				.setClock(Clock.fixed(Instant.ofEpochMilli(1).plus(expiresIn), ZoneId.systemDefault()));
+		OAuth2AuthorizationRequest authorizationRequest2 = createAuthorizationRequest().state(state2).build();
+		this.authorizationRequestRepository.saveAuthorizationRequest(authorizationRequest2, request, response);
+		request.addParameter(OAuth2ParameterNames.STATE, state1);
+		OAuth2AuthorizationRequest loadedAuthorizationRequest1 = this.authorizationRequestRepository
+				.loadAuthorizationRequest(request);
+		assertThat(loadedAuthorizationRequest1).isNull();
+		request.removeParameter(OAuth2ParameterNames.STATE);
+		request.addParameter(OAuth2ParameterNames.STATE, state2);
+		OAuth2AuthorizationRequest loadedAuthorizationRequest2 = this.authorizationRequestRepository
+				.loadAuthorizationRequest(request);
+		assertThat(loadedAuthorizationRequest2).isEqualTo(authorizationRequest2);
+	}
+
+	@Test
+	public void removeOldestAuthorizationRequestWhenMoreThanMax() {
+		this.authorizationRequestRepository.setMaxActiveAuthorizationRequestsPerSession(2);
+		MockHttpServletRequest request = new MockHttpServletRequest();
+		MockHttpServletResponse response = new MockHttpServletResponse();
+		String state1 = "state-1122";
+		OAuth2AuthorizationRequest authorizationRequest1 = createAuthorizationRequest().state(state1).build();
+		this.authorizationRequestRepository.saveAuthorizationRequest(authorizationRequest1, request, response);
+		String state2 = "state-3344";
+		OAuth2AuthorizationRequest authorizationRequest2 = createAuthorizationRequest().state(state2).build();
+		this.authorizationRequestRepository.saveAuthorizationRequest(authorizationRequest2, request, response);
+		String state3 = "state-4455";
+		OAuth2AuthorizationRequest authorizationRequest3 = createAuthorizationRequest().state(state3).build();
+		this.authorizationRequestRepository.saveAuthorizationRequest(authorizationRequest3, request, response);
+		request.addParameter(OAuth2ParameterNames.STATE, state1);
+		OAuth2AuthorizationRequest loadedAuthorizationRequest1 = this.authorizationRequestRepository
+				.loadAuthorizationRequest(request);
+		assertThat(loadedAuthorizationRequest1).isNull();
+		request.removeParameter(OAuth2ParameterNames.STATE);
+		request.addParameter(OAuth2ParameterNames.STATE, state2);
+		OAuth2AuthorizationRequest loadedAuthorizationRequest2 = this.authorizationRequestRepository
+				.loadAuthorizationRequest(request);
+		assertThat(loadedAuthorizationRequest2).isEqualTo(authorizationRequest2);
+		request.removeParameter(OAuth2ParameterNames.STATE);
+		request.addParameter(OAuth2ParameterNames.STATE, state3);
+		OAuth2AuthorizationRequest loadedAuthorizationRequest3 = this.authorizationRequestRepository
+				.loadAuthorizationRequest(request);
+		assertThat(loadedAuthorizationRequest3).isEqualTo(authorizationRequest3);
+	}
+
 	private OAuth2AuthorizationRequest.Builder createAuthorizationRequest() {
 		return OAuth2AuthorizationRequest.authorizationCode().authorizationUri("https://example.com/oauth2/authorize")
 				.clientId("client-id-1234").state("state-1234");