From e22f4453659ca8da9770d3dc1cc8dfdc3a6488d8 Mon Sep 17 00:00:00 2001
From: esfomeado <jonymilk@live.com.pt>
Date: Thu, 29 Jul 2021 23:20:03 +0100
Subject: [PATCH] Get CSRF token from local storage

---
 .../java/org/springdoc/core/Constants.java    |  9 +++-
 .../core/SwaggerUiConfigProperties.java       | 46 +++++++++++++++++++
 .../ui/AbstractSwaggerIndexTransformer.java   | 26 ++++++++++-
 3 files changed, 78 insertions(+), 3 deletions(-)

diff --git a/springdoc-openapi-common/src/main/java/org/springdoc/core/Constants.java b/springdoc-openapi-common/src/main/java/org/springdoc/core/Constants.java
index 2db6c2b4f..4b2631e5b 100644
--- a/springdoc-openapi-common/src/main/java/org/springdoc/core/Constants.java
+++ b/springdoc-openapi-common/src/main/java/org/springdoc/core/Constants.java
@@ -298,12 +298,17 @@ public final class Constants {
 	/**
 	 * The constant CSRF_DEFAULT_COOKIE_NAME.
 	 */
-	public static final String CSRF_DEFAULT_COOKIE_NAME= "XSRF-TOKEN";
+	public static final String CSRF_DEFAULT_COOKIE_NAME = "XSRF-TOKEN";
+
+	/**
+	 * The constant CSRF_DEFAULT_LOCAL_STORAGE_KEY
+	 */
+	public static final String CSRF_DEFAULT_LOCAL_STORAGE_KEY = "XSRF-TOKEN";
 
 	/**
 	 * The constant CSRF_DEFAULT_HEADER_NAME.
 	 */
-	public static final String CSRF_DEFAULT_HEADER_NAME= "X-XSRF-TOKEN";
+	public static final String CSRF_DEFAULT_HEADER_NAME = "X-XSRF-TOKEN";
 
 	/**
 	 * The constant OPERATION_ATTRIBUTE.
diff --git a/springdoc-openapi-common/src/main/java/org/springdoc/core/SwaggerUiConfigProperties.java b/springdoc-openapi-common/src/main/java/org/springdoc/core/SwaggerUiConfigProperties.java
index 5c8298503..77c59e036 100644
--- a/springdoc-openapi-common/src/main/java/org/springdoc/core/SwaggerUiConfigProperties.java
+++ b/springdoc-openapi-common/src/main/java/org/springdoc/core/SwaggerUiConfigProperties.java
@@ -117,11 +117,21 @@ public static class Csrf {
 		 */
 		private boolean enabled;
 
+		/**
+		 * Use Local storage.
+		 */
+		private boolean useLocalStorage;
+
 		/**
 		 * The Cookie name.
 		 */
 		private String cookieName = Constants.CSRF_DEFAULT_COOKIE_NAME;
 
+		/**
+		 * The Local storage key.
+		 */
+		private String localStorageKey = Constants.CSRF_DEFAULT_LOCAL_STORAGE_KEY;
+
 		/**
 		 * The Header name.
 		 */
@@ -145,6 +155,24 @@ public void setEnabled(boolean enabled) {
 			this.enabled = enabled;
 		}
 
+		/**
+		 * Use Local storage boolean.
+		 *
+		 * @return the boolean
+		 */
+		public boolean isUseLocalStorage() {
+			return useLocalStorage;
+		}
+
+		/**
+		 * Sets useLocalStorage.
+		 *
+		 * @param useLocalStorage the use local storage
+		 */
+		public void setUseLocalStorage(boolean useLocalStorage) {
+			this.useLocalStorage = useLocalStorage;
+		}
+
 		/**
 		 * Gets cookie name.
 		 *
@@ -163,6 +191,24 @@ public void setCookieName(String cookieName) {
 			this.cookieName = cookieName;
 		}
 
+		/**
+		 * Gets local storage key.
+		 *
+		 * @return the cookie name
+		 */
+		public String getLocalStorageKey() {
+			return localStorageKey;
+		}
+
+		/**
+		 * Sets local storage key.
+		 *
+		 * @param localStorageKey the local storage key
+		 */
+		public void setLocalStorageKey(String localStorageKey) {
+			this.localStorageKey = localStorageKey;
+		}
+
 		/**
 		 * Gets header name.
 		 *
diff --git a/springdoc-openapi-common/src/main/java/org/springdoc/ui/AbstractSwaggerIndexTransformer.java b/springdoc-openapi-common/src/main/java/org/springdoc/ui/AbstractSwaggerIndexTransformer.java
index a63d358f4..5c1bce941 100644
--- a/springdoc-openapi-common/src/main/java/org/springdoc/ui/AbstractSwaggerIndexTransformer.java
+++ b/springdoc-openapi-common/src/main/java/org/springdoc/ui/AbstractSwaggerIndexTransformer.java
@@ -143,7 +143,11 @@ protected String defaultTransformations(InputStream inputStream) throws IOExcept
 			html = overwriteSwaggerDefaultUrl(html);
 		}
 		if (swaggerUiConfig.isCsrfEnabled()) {
-			html = addCSRF(html);
+			if (swaggerUiConfig.getCsrf().isUseLocalStorage()) {
+				html = addCSRFLocalStorage(html);
+			} else {
+				html = addCSRF(html);
+			}
 		}
 		if (swaggerUiConfig.getSyntaxHighlight() != null) {
 			html = addSyntaxHighlight(html);
@@ -174,6 +178,26 @@ protected String addCSRF(String html) {
 		return html.replace(PRESETS, stringBuilder.toString());
 	}
 
+	/**
+	 * Add csrf string.
+	 *
+	 * @param html the html
+	 * @return the string
+	 */
+	protected String addCSRFLocalStorage(String html) {
+		StringBuilder stringBuilder = new StringBuilder();
+		stringBuilder.append("requestInterceptor: (request) => {\n");
+		stringBuilder.append("const value = window.localStorage.getItem('");
+		stringBuilder.append(swaggerUiConfig.getCsrf().getLocalStorageKey() + "');\n");
+		stringBuilder.append("request.headers['");
+		stringBuilder.append(swaggerUiConfig.getCsrf().getHeaderName());
+		stringBuilder.append("'] = value;\n");
+		stringBuilder.append("return request;\n");
+		stringBuilder.append("},\n");
+		stringBuilder.append(PRESETS);
+		return html.replace(PRESETS, stringBuilder.toString());
+	}
+
 	/**
 	 * Add syntax highlight string.
 	 *