Merge 91c373a into c8a470e

wtforms · Oct 28, 2014 · 2f83233 · 2f83233
2 parents c8a470e + 91c373a
commit 2f83233
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 7 deletions.
diff --git a/docs/config.rst b/docs/config.rst
@@ -14,6 +14,8 @@ WTF_CSRF_ENABLED     Disable/enable CSRF protection for forms.
                      Default is True.
 WTF_I18N_ENABLED     Disable/enable I18N support. This should work
                      together with Flask-Babel. Default is True.
+WTF_CSRF_HEADERS     CSRF token HTTP headers checked. Default is
+                     **['X-CSRFToken', 'X-CSRF-Token']**
 WTF_CSRF_SECRET_KEY  A random string for generating CSRF token.
                      Default is the same as SECRET_KEY.
 WTF_CSRF_TIME_LIMIT  CSRF token expiring time. Default is **3600**

diff --git a/flask_wtf/csrf.py b/flask_wtf/csrf.py
@@ -139,6 +139,9 @@ def __init__(self, app=None):
 
     def init_app(self, app):
         app.jinja_env.globals['csrf_token'] = generate_csrf
+        app.config.setdefault(
+            'WTF_CSRF_HEADERS', ['X-CSRFToken', 'X-CSRF-Token']
+        )
         app.config.setdefault('WTF_CSRF_SSL_STRICT', True)
         app.config.setdefault('WTF_CSRF_ENABLED', True)
         app.config.setdefault('WTF_CSRF_METHODS', ['POST', 'PUT', 'PATCH'])
@@ -178,13 +181,10 @@ def _csrf_protect():
                 for key in request.form:
                     if key.endswith('csrf_token'):
                         csrf_token = request.form[key]
-            if not csrf_token:
-                # You can get csrf token from header
-                # The header name is the same as Django
-                csrf_token = request.headers.get('X-CSRFToken')
-            if not csrf_token:
-                # The header name is the same as Rails
-                csrf_token = request.headers.get('X-CSRF-Token')
+            for header_name in app.config['WTF_CSRF_HEADERS']:
+                if csrf_token is not None:
+                    break
+                csrf_token = request.headers.get(header_name)
             if not validate_csrf(csrf_token):
                 reason = 'CSRF token missing or incorrect.'
                 return self._error_response(reason)

diff --git a/tests/test_csrf.py b/tests/test_csrf.py
@@ -170,6 +170,40 @@ def test_valid_secure_csrf(self):
         )
         assert response.status_code == 200
 
+    def test_empty_csrf_headers(self):
+        response = self.client.get("/", base_url='https://localhost/')
+        csrf_token = get_csrf_token(response.data)
+        self.app.config['WTF_CSRF_HEADERS'] = list()
+        response = self.client.post(
+            "/",
+            data={"name": "danny"},
+            headers={
+                'X-CSRFToken': csrf_token,
+            },
+            environ_base={
+                'HTTP_REFERER': 'https://localhost/',
+            },
+            base_url='https://localhost/',
+        )
+        assert response.status_code == 400
+
+    def test_custom_csrf_headers(self):
+        response = self.client.get("/", base_url='https://localhost/')
+        csrf_token = get_csrf_token(response.data)
+        self.app.config['WTF_CSRF_HEADERS'] = ['X-XSRF-TOKEN']
+        response = self.client.post(
+            "/",
+            data={"name": "danny"},
+            headers={
+                'X-XSRF-TOKEN': csrf_token,
+            },
+            environ_base={
+                'HTTP_REFERER': 'https://localhost/',
+            },
+            base_url='https://localhost/',
+        )
+        assert response.status_code == 200
+
     def test_not_endpoint(self):
         response = self.client.post('/not-endpoint')
         assert response.status_code == 404