From cfc66b4cb7fec1e84767b817619a322047330d0f Mon Sep 17 00:00:00 2001
From: Chris Wagner <jwag.wagner@gmail.com>
Date: Tue, 20 Feb 2024 12:36:24 -0800
Subject: [PATCH] Fix 2 CSRF issues. (#928)

1) If CSRF_COOKIE_NAME wasn't set, then our CSRF_HEADER wasn't added to the list in WTforms and therefore didn't work.
2) If you set CSRF_PROTECT_MECHANISMS to an empty list (to hopefully disable CSRF) - it in fact didn't stop CSRF...

closes #870
---
 flask_security/core.py       |  8 ++++----
 flask_security/decorators.py |  4 ++--
 tests/test_csrf.py           | 31 ++++++++++++++++++++++++++++++-
 3 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/flask_security/core.py b/flask_security/core.py
index c15de9e5..fad452cb 100644
--- a/flask_security/core.py
+++ b/flask_security/core.py
@@ -1667,18 +1667,18 @@ def _csrf_init(app):
         if cv("CSRF_COOKIE_NAME", app=app) and not csrf:
             # Common use case is for cookie value to be used as contents for header
             # which is only looked at when CsrfProtect is initialized.
-            # Yes, this is opinionated - they can always get CSRF token via:
-            # 'get /login'
             raise ValueError(
                 "CSRF_COOKIE defined however CsrfProtect not part of application"
             )
 
         if csrf:
             csrf.exempt("flask_security.views.logout")
+            # Add configured header to WTF_CSRF_HEADERS
+            if ch := cv("CSRF_HEADER", app=app):
+                if ch not in app.config["WTF_CSRF_HEADERS"]:
+                    app.config["WTF_CSRF_HEADERS"].append(ch)
         if cv("CSRF_COOKIE_NAME", app=app):
             app.after_request(csrf_cookie_handler)
-            # Add configured header to WTF_CSRF_HEADERS
-            app.config["WTF_CSRF_HEADERS"].append(cv("CSRF_HEADER", app=app))
 
     def set_form_info(self, name: str, form_info: FormInfo) -> None:
         """Set form instantiation info.
diff --git a/flask_security/decorators.py b/flask_security/decorators.py
index c6db7db1..23d2e469 100644
--- a/flask_security/decorators.py
+++ b/flask_security/decorators.py
@@ -246,8 +246,8 @@ def handle_csrf(method: str, json_response: bool = False) -> ResponseValue | Non
                     payload = json_error_response(errors=e.description)
                     return _security._render_json(payload, 400, None, None)
                 raise
-        else:
-            set_request_attr("fs_ignore_csrf", True)
+            return None
+    set_request_attr("fs_ignore_csrf", True)
     return None
 
 
diff --git a/tests/test_csrf.py b/tests/test_csrf.py
index d0d0ebe0..67e47557 100644
--- a/tests/test_csrf.py
+++ b/tests/test_csrf.py
@@ -259,10 +259,13 @@ def test_cp_reset(app, client):
 
 @pytest.mark.changeable()
 @pytest.mark.csrf(csrfprotect=True)
+@pytest.mark.settings(csrf_header="X-XSRF-Token")
 def test_cp_with_token(app, client):
     # Make sure can use returned CSRF-Token in Header.
     # Since the csrf token isn't in the form - must enable app-wide CSRF
     # using CSRFProtect() - as the above mark does.
+    # Using X-XSRF-Token as header tests that we properly
+    # add that as a known header to WTFforms.
     auth_token, csrf_token = json_login(client, use_header=True)
 
     # make sure returned csrf_token works in header.
@@ -277,7 +280,7 @@ def test_cp_with_token(app, client):
             "/change",
             content_type="application/json",
             json=data,
-            headers={"X-CSRF-Token": csrf_token},
+            headers={"X-XSRF-Token": csrf_token},
         )
         assert response.status_code == 200
     assert mp.success == 1 and mp.failure == 0
@@ -417,6 +420,32 @@ def test_different_mechanisms_nc(app, client_nc):
     assert mp.success == 0 and mp.failure == 0
 
 
+@pytest.mark.changeable()
+@pytest.mark.csrf(csrfprotect=True)
+@pytest.mark.settings(csrf_protect_mechanisms=[])
+def test_cp_with_token_empty_mechanisms(app, client):
+    # If no mechanisms - shouldn't do any CSRF
+    auth_token, csrf_token = json_login(client, use_header=True)
+
+    # make sure returned csrf_token works in header.
+    data = dict(
+        password="password",
+        new_password="battery staple",
+        new_password_confirm="battery staple",
+    )
+
+    response = client.post(
+        "/change",
+        content_type="application/json",
+        json=data,
+        headers={
+            "Content-Type": "application/json",
+            "Authentication-Token": auth_token,
+        },
+    )
+    assert response.status_code == 200
+
+
 @pytest.mark.settings(csrf_ignore_unauth_endpoints=True, CSRF_COOKIE_NAME="XSRF-Token")
 def test_csrf_cookie(app, sqlalchemy_datastore):
     app.config["WTF_CSRF_ENABLED"] = True