Allow different authorization header strings (closes #2)

Author: vimalloc

Diff for: README.md

@@ -359,6 +359,7 @@ The available options are:
 
 | Name   | Description | Options | Default|
 | ------ | ----------- | ------- | ------ |
+|JWT_AUTH_HEADER | What to use in the authorization header (ex: Bearer <access_token>) | Any string (empty string to have it just be the access token in the authorization header) | 'Bearer' |
 |JWT_ACCESS_TOKEN_EXPIRES | How long an access token should live | datetime.timedelta | 15 minutes|
 |JWT_REFRESH_TOKEN_EXPIRES | How long a refresh token should live | datetime.timedelta | 30 days |
 |JWT_ALGORITHM | Which algorithm to use with the JWT. | [See here] (https://pyjwt.readthedocs.io/en/latest/algorithms.html) | HS256 |

Diff for: flask_jwt_extended/config.py

@@ -3,6 +3,9 @@
 
 # Defaults
 
+# Authorize header type, what we are expecting to see in the auth header
+AUTH_HEADER = 'Bearer'
+
 # How long an access token will live before it expires.
 ACCESS_TOKEN_EXPIRES = datetime.timedelta(minutes=15)
 
@@ -28,6 +31,10 @@
 BLACKLIST_TOKEN_CHECKS = 'refresh'
 
 
+def get_auth_header():
+    return current_app.config.get('JWT_AUTH_HEADER', AUTH_HEADER)
+
+
 def get_access_expires():
     return current_app.config.get('JWT_ACCESS_TOKEN_EXPIRES', ACCESS_TOKEN_EXPIRES)
 

Diff for: flask_jwt_extended/utils.py

@@ -14,7 +14,7 @@
     from flask import _request_ctx_stack as ctx_stack
 
 from flask_jwt_extended.config import get_access_expires, get_refresh_expires, \
-    get_algorithm, get_blacklist_enabled, get_blacklist_checks
+    get_algorithm, get_blacklist_enabled, get_blacklist_checks, get_auth_header
 from flask_jwt_extended.exceptions import JWTEncodeError, JWTDecodeError, \
     InvalidHeaderError, NoAuthHeaderError, WrongTokenError, RevokedTokenError, \
     FreshTokenRequired
@@ -143,15 +143,19 @@ def _decode_jwt_from_request():
         raise NoAuthHeaderError("Missing Authorization Header")
 
     # Make sure the header is valid
+    expected_header = get_auth_header()
     parts = auth_header.split()
-    if parts[0] != 'Bearer':
-        msg = "Badly formatted authorization header. Should be 'Bearer <JWT>'"
-        raise InvalidHeaderError(msg)
-    elif len(parts) != 2:
-        msg = "Badly formatted authorization header. Should be 'Bearer <JWT>'"
-        raise InvalidHeaderError(msg)
-
-    token = parts[1]
+    if not expected_header:
+        if len(parts) != 1:
+            msg = "Badly formatted authorization header. Should be '<JWT>'"
+            raise InvalidHeaderError(msg)
+        token = parts[0]
+    else:
+        if parts[0] != expected_header or len(parts) != 2:
+            msg = "Bad authorization header. Expected '{} <JWT>'".format(expected_header)
+            raise InvalidHeaderError(msg)
+        token = parts[1]
+
     secret = _get_secret_key()
     algorithm = get_algorithm()
     return _decode_jwt(token, secret, algorithm)

Diff for: tests/test_config.py

@@ -6,7 +6,7 @@
 
 from flask_jwt_extended.config import get_access_expires, get_refresh_expires, \
     get_algorithm, get_blacklist_enabled, get_blacklist_store, \
-    get_blacklist_checks
+    get_blacklist_checks, get_auth_header
 from flask_jwt_extended import JWTManager
 
 
@@ -26,6 +26,7 @@ def test_default_configs(self):
             self.assertEqual(get_blacklist_enabled(), False)
             self.assertEqual(get_blacklist_store(), None)
             self.assertEqual(get_blacklist_checks(), 'refresh')
+            self.assertEqual(get_auth_header(), 'Bearer')
 
     def test_override_configs(self):
         self.app.config['JWT_ACCESS_TOKEN_EXPIRES'] = timedelta(minutes=5)
@@ -34,6 +35,7 @@ def test_override_configs(self):
         self.app.config['JWT_BLACKLIST_ENABLED'] = True
         self.app.config['JWT_BLACKLIST_STORE'] = simplekv.memory.DictStore()
         self.app.config['JWT_BLACKLIST_TOKEN_CHECKS'] = 'all'
+        self.app.config['JWT_AUTH_HEADER'] = 'JWT'
 
         with self.app.test_request_context():
             self.assertEqual(get_access_expires(), timedelta(minutes=5))
@@ -42,3 +44,4 @@ def test_override_configs(self):
             self.assertEqual(get_blacklist_enabled(), True)
             self.assertIsInstance(get_blacklist_store(), simplekv.memory.DictStore)
             self.assertEqual(get_blacklist_checks(), 'all')
+            self.assertEqual(get_auth_header(), 'JWT')

Diff for: tests/test_protected_endpoints.py

@@ -62,8 +62,8 @@ def _jwt_post(self, url, jwt):
         data = json.loads(response.get_data(as_text=True))
         return status_code, data
 
-    def _jwt_get(self, url, jwt):
-        auth_header = 'Bearer {}'.format(jwt)
+    def _jwt_get(self, url, jwt, auth_header='Bearer'):
+        auth_header = '{} {}'.format(auth_header, jwt).strip()
         response = self.client.get(url, headers={'Authorization': auth_header})
         status_code = response.status_code
         data = json.loads(response.get_data(as_text=True))
@@ -278,3 +278,23 @@ def claims():
         status, data = self._jwt_get('/claims', access_token)
         self.assertEqual(status, 200)
         self.assertEqual(data, {'username': 'test', 'claims': {'foo': 'bar'}})
+
+    def test_different_auth_header(self):
+        response = self.client.post('/auth/login')
+        data = json.loads(response.get_data(as_text=True))
+        access_token = data['access_token']
+
+        self.app.config['JWT_AUTH_HEADER'] = 'JWT'
+        status, data = self._jwt_get('/protected', access_token, auth_header='JWT')
+        self.assertEqual(data, {'msg': 'hello world'})
+        self.assertEqual(status, 200)
+
+        self.app.config['JWT_AUTH_HEADER'] = ''
+        status, data = self._jwt_get('/protected', access_token, auth_header='')
+        self.assertEqual(data, {'msg': 'hello world'})
+        self.assertEqual(status, 200)
+
+        self.app.config['JWT_AUTH_HEADER'] = ''
+        status, data = self._jwt_get('/protected', access_token, auth_header='Bearer')
+        self.assertIn('msg', data)
+        self.assertEqual(status, 422)