Add support for JWTs in headers and cookies in the same app

refs #26

Author: vimalloc

docs/options.rst

@@ -13,7 +13,8 @@ The available options are:
 
 ================================= =========================================
 ``JWT_TOKEN_LOCATION``            Where to look for a JWT when processing a request. The options are ``'headers'`` or
-                                  ``'cookies'``. Defaults to ``'headers'``
+                                  ``'cookies'``. You can pass in a list to check more then one location: ```['headers', 'cookies']```.
+                                  Defaults to ``'headers'``
 ``JWT_HEADER_NAME``               What header to look for the JWT in a request. Only used if we are sending
                                   the JWT in via headers. Defaults to ``'Authorization'``
 ``JWT_HEADER_TYPE``               What type of header the JWT is in. Defaults to ``'Bearer'``. This can be

examples/csrf_protection_with_cookies.py

@@ -9,7 +9,7 @@
 app.secret_key = 'super-secret'  # Change this!
 
 # Configure application to store JWTs in cookies
-app.config['JWT_TOKEN_LOCATION'] = 'cookies'
+app.config['JWT_TOKEN_LOCATION'] = ['cookies']
 
 # Only allow JWT cookies to be sent over https. In production, this
 # should likely be True

examples/jwt_in_cookie.py

@@ -16,7 +16,7 @@
 # Configure application to store JWTs in cookies. Whenever you make
 # a request to a protected endpoint, you will need to send in the
 # access or refresh JWT via a cookie.
-app.config['JWT_TOKEN_LOCATION'] = 'cookies'
+app.config['JWT_TOKEN_LOCATION'] = ['cookies']
 
 # Set the cookie paths, so that you are only sending your access token
 # cookie to the access endpoints, and only sending your refresh token

flask_jwt_extended/config.py

@@ -2,13 +2,8 @@
 from flask import current_app
 
 
-# TODO support for cookies and headers at the same time. This could be useful
-#      for using cookies in a web browser (more secure), and headers in a mobile
-#      app (don't have to worry about csrf/xss there, and headers are easier to
-#      manage in that environment)
-
 # Where to look for the JWT. Available options are cookies or headers
-TOKEN_LOCATION = 'headers'
+TOKEN_LOCATION = ['headers']
 
 # Options for JWTs when the TOKEN_LOCATION is headers
 HEADER_NAME = 'Authorization'
@@ -43,10 +38,14 @@
 
 
 def get_token_location():
-    location = current_app.config.get('JWT_TOKEN_LOCATION', TOKEN_LOCATION)
-    if location not in ['headers', 'cookies']:
-        raise RuntimeError('JWT_LOCATION_LOCATION must be "headers" or "cookies"')
-    return location
+    locations = current_app.config.get('JWT_TOKEN_LOCATION', TOKEN_LOCATION)
+    if not isinstance(locations, list):
+        locations = [locations]
+    for location in locations:
+        if location not in ('headers', 'cookies'):
+            raise RuntimeError('JWT_LOCATION_LOCATION can only contain '
+                               '"headers" and/or "cookies"')
+    return locations
 
 
 def get_jwt_header_name():

flask_jwt_extended/utils.py

@@ -87,7 +87,7 @@ def _encode_access_token(identity, secret, algorithm, token_expire_delta,
         'type': 'access',
         'user_claims': user_claims,
     }
-    if get_token_location() == 'cookies' and get_cookie_csrf_protect():
+    if 'cookies' in get_token_location() and get_cookie_csrf_protect() is True:
         token_data['csrf'] = _create_csrf_token()
     encoded_token = jwt.encode(token_data, secret, algorithm).decode('utf-8')
 
@@ -119,7 +119,7 @@ def _encode_refresh_token(identity, secret, algorithm, token_expire_delta):
         'identity': identity,
         'type': 'refresh',
     }
-    if get_token_location() == 'cookies' and get_cookie_csrf_protect():
+    if 'cookies' in get_token_location() and get_cookie_csrf_protect() is True:
         token_data['csrf'] = _create_csrf_token()
     encoded_token = jwt.encode(token_data, secret, algorithm).decode('utf-8')
 
@@ -153,9 +153,6 @@ def _decode_jwt(token, secret, algorithm):
             raise JWTDecodeError("Missing or invalid claim: fresh")
         if 'user_claims' not in data or not isinstance(data['user_claims'], dict):
             raise JWTDecodeError("Missing or invalid claim: user_claims")
-    if get_token_location() == 'cookies' and get_cookie_csrf_protect():
-        if 'csrf' not in data or not isinstance(data['csrf'], six.string_types):
-            raise JWTDecodeError("Missing or invalid claim: csrf")
     return data
 
 
@@ -200,17 +197,41 @@ def _decode_jwt_from_cookies(type):
 
     if get_cookie_csrf_protect():
         csrf_header_key = get_csrf_header_name()
-        csrf = request.headers.get(csrf_header_key, None)
-        if not csrf or not safe_str_cmp(csrf,  token['csrf']):
-            raise NoAuthorizationError("Missing or invalid csrf double submit header")
-
+        csrf_token_from_header = request.headers.get(csrf_header_key, None)
+        csrf_token_from_cookie = token.get('csrf', None)
+
+        # Verify the csrf tokens are present and matching
+        if csrf_token_from_cookie is None:
+            raise JWTDecodeError("Missing claim: 'csrf'")
+        if not isinstance(csrf_token_from_cookie, six.string_types):
+            raise JWTDecodeError("Invalid claim: 'csrf' (must be a string)")
+        if csrf_token_from_header is None:
+            raise NoAuthorizationError("Missing CSRF token in headers")
+        if not safe_str_cmp(csrf_token_from_header,  csrf_token_from_cookie):
+            raise NoAuthorizationError("CSRF double submit tokens do not match")
     return token
 
 
 def _decode_jwt_from_request(type):
-    token_location = get_token_location()
-    if token_location == 'headers':
+    token_locations = get_token_location()
+
+    # JWT can be in either headers or cookies
+    if 'headers' in token_locations and 'cookies' in token_locations:
+        try:
+            return _decode_jwt_from_headers()
+        except NoAuthorizationError:
+            pass
+        try:
+            return _decode_jwt_from_cookies(type)
+        except NoAuthorizationError:
+            pass
+        raise NoAuthorizationError("Missing JWT in header and cookies")
+
+    # JWT can only be in headers
+    elif 'headers' in token_locations:
         return _decode_jwt_from_headers()
+
+    # JWT can only be in cookie
     else:
         return _decode_jwt_from_cookies(type)
 

tests/test_config.py

@@ -25,7 +25,7 @@ def setUp(self):
 
     def test_default_configs(self):
         with self.app.test_request_context():
-            self.assertEqual(get_token_location(), 'headers')
+            self.assertEqual(get_token_location(), ['headers'])
             self.assertEqual(get_jwt_header_name(), 'Authorization')
             self.assertEqual(get_jwt_header_type(), 'Bearer')
 
@@ -69,7 +69,7 @@ def test_override_configs(self):
         self.app.config['JWT_BLACKLIST_TOKEN_CHECKS'] = 'all'
 
         with self.app.test_request_context():
-            self.assertEqual(get_token_location(), 'cookies')
+            self.assertEqual(get_token_location(), ['cookies'])
             self.assertEqual(get_jwt_header_name(), 'Auth')
             self.assertEqual(get_jwt_header_type(), 'JWT')
 

tests/test_jwt_encode_decode.py

@@ -302,36 +302,6 @@ def test_decode_invalid_jwt(self):
                 encoded_token = jwt.encode(token_data, 'secret', 'HS256').decode('utf-8')
                 _decode_jwt(encoded_token, 'secret', 'HS256')
 
-        # Missing and bad csrf tokens
-        self.app.config['JWT_TOKEN_LOCATION'] = 'cookies'
-        self.app.config['JWT_COOKIE_CSRF_PROTECTION'] = True
-        with self.app.test_request_context():
-            now = datetime.utcnow()
-            with self.assertRaises(JWTDecodeError):
-                token_data = {
-                    'exp': now + timedelta(minutes=5),
-                    'iat': now,
-                    'nbf': now,
-                    'jti': 'banana',
-                    'identity': 'banana',
-                    'type': 'refresh',
-                }
-                encoded_token = jwt.encode(token_data, 'secret', 'HS256').decode('utf-8')
-                _decode_jwt(encoded_token, 'secret', 'HS256')
-
-            with self.assertRaises(JWTDecodeError):
-                token_data = {
-                    'exp': now + timedelta(minutes=5),
-                    'iat': now,
-                    'nbf': now,
-                    'jti': 'banana',
-                    'identity': 'banana',
-                    'type': 'refresh',
-                    'csrf': True
-                }
-                encoded_token = jwt.encode(token_data, 'secret', 'HS256').decode('utf-8')
-                _decode_jwt(encoded_token, 'secret', 'HS256')
-
     def test_create_access_token_with_object(self):
         # Complex object to test building a JWT from. Normally if you are using
         # this functionality, this is something that would be retrieved from

tests/test_protected_endpoints.py

@@ -358,6 +358,8 @@ def setUp(self):
         self.app.config['JWT_TOKEN_LOCATION'] = 'cookies'
         self.app.config['JWT_ACCESS_COOKIE_PATH'] = '/api/'
         self.app.config['JWT_REFRESH_COOKIE_PATH'] = '/auth/refresh'
+        self.app.config['JWT_ACCESS_COOKIE_NAME'] = 'access_token_cookie'
+        self.app.config['JWT_ALGORITHM'] = 'HS256'
         self.jwt_manager = JWTManager(self.app)
         self.client = self.app.test_client()
 
@@ -549,9 +551,17 @@ def test_access_endpoints_with_cookies_and_csrf(self):
         self.assertEqual(status_code, 401)
         self.assertIn('msg', data)
 
-        # Try with logged in and bad double submit token
+        # Try with logged in and bad header name for double submit token
+        response = self.client.get('/api/protected',
+                                   headers={'bad-header-name': 'banana'})
+        status_code = response.status_code
+        data = json.loads(response.get_data(as_text=True))
+        self.assertEqual(status_code, 401)
+        self.assertIn('msg', data)
+
+        # Try with logged in and bad header data for double submit token
         response = self.client.get('/api/protected',
-                                   headers={'X-CSRF-ACCESS-TOKEN': 'banana'})
+                                   headers={'X-CSRF-TOKEN': 'banana'})
         status_code = response.status_code
         data = json.loads(response.get_data(as_text=True))
         self.assertEqual(status_code, 401)
@@ -564,3 +574,159 @@ def test_access_endpoints_with_cookies_and_csrf(self):
         data = json.loads(response.get_data(as_text=True))
         self.assertEqual(status_code, 200)
         self.assertEqual(data, {'msg': 'hello world'})
+
+    def test_access_endpoints_with_cookie_missing_csrf_field(self):
+        # Test accessing a csrf protected endpoint with a cookie that does not
+        # have a csrf token in it
+        self.app.config['JWT_COOKIE_CSRF_PROTECT'] = False
+        self._login()
+        self.app.config['JWT_COOKIE_CSRF_PROTECT'] = True
+
+        response = self.client.get('/api/protected')
+        status_code = response.status_code
+        data = json.loads(response.get_data(as_text=True))
+        self.assertEqual(status_code, 422)
+        self.assertIn('msg', data)
+
+    def test_access_endpoints_with_cookie_csrf_claim_not_string(self):
+        now = datetime.utcnow()
+        token_data = {
+            'exp': now + timedelta(minutes=5),
+            'iat': now,
+            'nbf': now,
+            'jti': 'banana',
+            'identity': 'banana',
+            'type': 'refresh',
+            'csrf': 404
+        }
+        secret = self.app.secret_key
+        algorithm = self.app.config['JWT_ALGORITHM']
+        encoded_token = jwt.encode(token_data, secret, algorithm).decode('utf-8')
+        access_cookie_key = self.app.config['JWT_ACCESS_COOKIE_NAME']
+        self.client.set_cookie('localhost', access_cookie_key, encoded_token)
+
+        self.app.config['JWT_COOKIE_CSRF_PROTECT'] = True
+        response = self.client.get('/api/protected')
+        status_code = response.status_code
+        data = json.loads(response.get_data(as_text=True))
+        self.assertEqual(status_code, 422)
+        self.assertIn('msg', data)
+
+
+class TestEndpointsWithHeadersAndCookies(unittest.TestCase):
+
+    def setUp(self):
+        self.app = Flask(__name__)
+        self.app.secret_key = 'super=secret'
+        self.app.config['JWT_TOKEN_LOCATION'] = ['cookies', 'headers']
+        self.app.config['JWT_COOKIE_CSRF_PROTECT'] = True
+        self.app.config['JWT_ACCESS_COOKIE_PATH'] = '/api/'
+        self.app.config['JWT_REFRESH_COOKIE_PATH'] = '/auth/refresh'
+        self.jwt_manager = JWTManager(self.app)
+        self.client = self.app.test_client()
+
+        @self.app.route('/auth/login_cookies', methods=['POST'])
+        def login_cookies():
+            # Create the tokens we will be sending back to the user
+            access_token = create_access_token(identity='test')
+            refresh_token = create_refresh_token(identity='test')
+
+            # Set the JWTs and the CSRF double submit protection cookies in this response
+            resp = jsonify({'login': True})
+            set_access_cookies(resp, access_token)
+            set_refresh_cookies(resp, refresh_token)
+            return resp, 200
+
+        @self.app.route('/auth/login_headers', methods=['POST'])
+        def login_headers():
+            ret = {
+                'access_token': create_access_token('test', fresh=True),
+                'refresh_token': create_refresh_token('test')
+            }
+            return jsonify(ret), 200
+
+        @self.app.route('/api/protected')
+        @jwt_required
+        def protected():
+            return jsonify({'msg': "hello world"})
+
+    def _jwt_post(self, url, jwt):
+        response = self.client.post(url, content_type='application/json',
+                                    headers={'Authorization': 'Bearer {}'.format(jwt)})
+        status_code = response.status_code
+        data = json.loads(response.get_data(as_text=True))
+        return status_code, data
+
+    def _jwt_get(self, url, jwt, header_name='Authorization', header_type='Bearer'):
+        header_type = '{} {}'.format(header_type, jwt).strip()
+        response = self.client.get(url, headers={header_name: header_type})
+        status_code = response.status_code
+        data = json.loads(response.get_data(as_text=True))
+        return status_code, data
+
+    def _login_cookies(self):
+        resp = self.client.post('/auth/login_cookies')
+        index = 1
+
+        access_cookie_str = resp.headers[index][1]
+        access_cookie_key = access_cookie_str.split('=')[0]
+        access_cookie_value = "".join(access_cookie_str.split('=')[1:])
+        self.client.set_cookie('localhost', access_cookie_key, access_cookie_value)
+        index += 1
+
+        if self.app.config['JWT_COOKIE_CSRF_PROTECT']:
+            access_csrf_str = resp.headers[index][1]
+            access_csrf_key = access_csrf_str.split('=')[0]
+            access_csrf_value = "".join(access_csrf_str.split('=')[1:])
+            self.client.set_cookie('localhost', access_csrf_key, access_csrf_value)
+            index += 1
+            access_csrf = access_csrf_value.split(';')[0]
+        else:
+            access_csrf = ""
+
+        refresh_cookie_str = resp.headers[index][1]
+        refresh_cookie_key = refresh_cookie_str.split('=')[0]
+        refresh_cookie_value = "".join(refresh_cookie_str.split('=')[1:])
+        self.client.set_cookie('localhost', refresh_cookie_key, refresh_cookie_value)
+        index += 1
+
+        if self.app.config['JWT_COOKIE_CSRF_PROTECT']:
+            refresh_csrf_str = resp.headers[index][1]
+            refresh_csrf_key = refresh_csrf_str.split('=')[0]
+            refresh_csrf_value = "".join(refresh_csrf_str.split('=')[1:])
+            self.client.set_cookie('localhost', refresh_csrf_key, refresh_csrf_value)
+            refresh_csrf = refresh_csrf_value.split(';')[0]
+        else:
+            refresh_csrf = ""
+
+        return access_csrf, refresh_csrf
+
+    def _login_headers(self):
+        resp = self.client.post('/auth/login_headers')
+        data = json.loads(resp.get_data(as_text=True))
+        return data['access_token'], data['refresh_token']
+
+    def test_accessing_endpoint_with_headers(self):
+        access_token, _ = self._login_headers()
+        header_type = '{} {}'.format('Bearer', access_token).strip()
+        response = self.client.get('/api/protected', headers={'Authorization': header_type})
+        status_code = response.status_code
+        data = json.loads(response.get_data(as_text=True))
+        self.assertEqual(status_code, 200)
+        self.assertEqual(data, {'msg': 'hello world'})
+
+    def test_accessing_endpoint_with_cookies(self):
+        access_csrf, _ = self._login_cookies()
+        response = self.client.get('/api/protected',
+                                   headers={'X-CSRF-TOKEN': access_csrf})
+        status_code = response.status_code
+        data = json.loads(response.get_data(as_text=True))
+        self.assertEqual(status_code, 200)
+        self.assertEqual(data, {'msg': 'hello world'})
+
+    def test_accessing_endpoint_without_jwt(self):
+        response = self.client.get('/api/protected')
+        status_code = response.status_code
+        data = json.loads(response.get_data(as_text=True))
+        self.assertEqual(status_code, 401)
+        self.assertIn('msg', data)