From 666f1752251dc2f00fff4a355afc073a3b9d7bab Mon Sep 17 00:00:00 2001
From: Rob Winch <rwinch@users.noreply.github.com>
Date: Thu, 7 Jul 2022 16:36:35 -0500
Subject: [PATCH] LazyCsrfTokenRepository#loadToken Supports Deferring
 Delegation

Previously LazyCsrfTokenRepository supported lazily saving the CsrfToken
which allowed for lazily saving the CsrfToken. However, it did not
support lazily reading the CsrfToken. This meant every request required
reading the CsrfToken (often the HttpSession).

This commit allows for lazily reading the CsrfToken and thus prevents
unnecessary reads to the HttpSession.

Closes gh-11700
---
 .../web/csrf/LazyCsrfTokenRepository.java     | 63 +++++++++++++++++++
 .../csrf/LazyCsrfTokenRepositoryTests.java    | 12 ++++
 2 files changed, 75 insertions(+)

diff --git a/web/src/main/java/org/springframework/security/web/csrf/LazyCsrfTokenRepository.java b/web/src/main/java/org/springframework/security/web/csrf/LazyCsrfTokenRepository.java
index b0f4263f2a..a9817cc9c7 100644
--- a/web/src/main/java/org/springframework/security/web/csrf/LazyCsrfTokenRepository.java
+++ b/web/src/main/java/org/springframework/security/web/csrf/LazyCsrfTokenRepository.java
@@ -38,6 +38,8 @@ public final class LazyCsrfTokenRepository implements CsrfTokenRepository {
 
 	private final CsrfTokenRepository delegate;
 
+	private boolean deferLoadToken;
+
 	/**
 	 * Creates a new instance
 	 * @param delegate the {@link CsrfTokenRepository} to use. Cannot be null
@@ -48,6 +50,15 @@ public final class LazyCsrfTokenRepository implements CsrfTokenRepository {
 		this.delegate = delegate;
 	}
 
+	/**
+	 * Determines if {@link #loadToken(HttpServletRequest)} should be lazily loaded.
+	 * @param deferLoadToken true if should lazily load
+	 * {@link #loadToken(HttpServletRequest)}. Default false.
+	 */
+	public void setDeferLoadToken(boolean deferLoadToken) {
+		this.deferLoadToken = deferLoadToken;
+	}
+
 	/**
 	 * Generates a new token
 	 * @param request the {@link HttpServletRequest} to use. The
@@ -77,6 +88,9 @@ public final class LazyCsrfTokenRepository implements CsrfTokenRepository {
 	 */
 	@Override
 	public CsrfToken loadToken(HttpServletRequest request) {
+		if (this.deferLoadToken) {
+			return new LazyLoadCsrfToken(request, this.delegate);
+		}
 		return this.delegate.loadToken(request);
 	}
 
@@ -92,6 +106,55 @@ public final class LazyCsrfTokenRepository implements CsrfTokenRepository {
 		return response;
 	}
 
+	private final class LazyLoadCsrfToken implements CsrfToken {
+
+		private final HttpServletRequest request;
+
+		private final CsrfTokenRepository tokenRepository;
+
+		private CsrfToken token;
+
+		private LazyLoadCsrfToken(HttpServletRequest request, CsrfTokenRepository tokenRepository) {
+			this.request = request;
+			this.tokenRepository = tokenRepository;
+		}
+
+		private CsrfToken getDelegate() {
+			if (this.token != null) {
+				return this.token;
+			}
+			// load from the delegate repository
+			this.token = LazyCsrfTokenRepository.this.delegate.loadToken(this.request);
+			if (this.token == null) {
+				// return a generated token that is lazily saved since
+				// LazyCsrfTokenRepository#loadToken always returns a value
+				this.token = generateToken(this.request);
+			}
+			return this.token;
+		}
+
+		@Override
+		public String getHeaderName() {
+			return getDelegate().getHeaderName();
+		}
+
+		@Override
+		public String getParameterName() {
+			return getDelegate().getParameterName();
+		}
+
+		@Override
+		public String getToken() {
+			return getDelegate().getToken();
+		}
+
+		@Override
+		public String toString() {
+			return "LazyLoadCsrfToken{" + "token=" + this.token + '}';
+		}
+
+	}
+
 	private static final class SaveOnAccessCsrfToken implements CsrfToken {
 
 		private transient CsrfTokenRepository tokenRepository;
diff --git a/web/src/test/java/org/springframework/security/web/csrf/LazyCsrfTokenRepositoryTests.java b/web/src/test/java/org/springframework/security/web/csrf/LazyCsrfTokenRepositoryTests.java
index 5ad35258d4..35d57634af 100644
--- a/web/src/test/java/org/springframework/security/web/csrf/LazyCsrfTokenRepositoryTests.java
+++ b/web/src/test/java/org/springframework/security/web/csrf/LazyCsrfTokenRepositoryTests.java
@@ -31,6 +31,7 @@ import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException
 import static org.mockito.BDDMockito.given;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.verifyZeroInteractions;
 
 /**
@@ -98,4 +99,15 @@ public class LazyCsrfTokenRepositoryTests {
 		verify(this.delegate).loadToken(this.request);
 	}
 
+	@Test
+	public void loadTokenWhenDeferLoadToken() {
+		given(this.delegate.loadToken(this.request)).willReturn(this.token);
+		this.repository.setDeferLoadToken(true);
+		CsrfToken loadToken = this.repository.loadToken(this.request);
+		verifyNoInteractions(this.delegate);
+		assertThat(loadToken.getToken()).isEqualTo(this.token.getToken());
+		assertThat(loadToken.getHeaderName()).isEqualTo(this.token.getHeaderName());
+		assertThat(loadToken.getParameterName()).isEqualTo(this.token.getParameterName());
+	}
+
 }