Merge pull request #23 from doordash/refresh-token

Add refresh token feature

Author: jpadilla

README.md

@@ -61,6 +61,22 @@ Now in order to access protected api urls you must include the `Authorization: J
 $ curl -H "Authorization: JWT <your_token>" http://localhost:8000/protected-url/
 ```
 
+## Refresh Token
+If `JWT_ALLOW_REFRESH` is True, issued tokens can be "refreshed" to obtain a new brand token with renewed expiration time. Add a URL pattern like this:
+```python
+    url(r'^api-token-refresh/', 'rest_framework_jwt.views.refresh_jwt_token'),
+```
+
+Pass in an existing token to the refresh endpoint as follows: `{"token": EXISTING_TOKEN}`. Note that only non-expired tokens will work. The JSON response looks the same as the normal obtain token endpoint `{"token": NEW_TOKEN}`.
+
+```bash
+$ curl -X POST -H "Content-Type: application/json" -d '{"token":"<EXISTING_TOKEN>}' http://localhost:8000/api-token-refresh/
+```
+
+Refresh with tokens can be repeated (token1 -> token2 -> token3), but this chain of token stores the time that the original token (obtained with username/password credentials), as `orig_iat`. You can only keep refreshing tokens up to `JWT_REFRESH_EXPIRATION_DELTA`.
+
+A typical use case might be a web app where you'd like to keep the user "logged in" the site without having to re-enter their password, or get kicked out by surprise before their token expired. Imagine they had a 1-hour token and are just at the last minute while they're still doing something. With mobile you could perhaps store the username/password to get a new token, but this is not a great idea in a browser. Each time the user loads the page, you can check if there is an existing non-expired token and if it's close to being expired, refresh it to extend their session. In other words, if a user is actively using your site, they can keep their "session" alive.
+
 ## Additional Settings
 There are some additional settings that you can override similar to how you'd do it with Django REST framework itself. Here are all the available defaults.
 
@@ -75,12 +91,18 @@ JWT_AUTH = {
     'JWT_PAYLOAD_HANDLER':
     'rest_framework_jwt.utils.jwt_payload_handler',
 
+    'JWT_PAYLOAD_GET_USER_ID_HANDLER':
+    'rest_framework_jwt.utils.jwt_get_user_id_from_payload_handler',
+
     'JWT_SECRET_KEY': settings.SECRET_KEY,
     'JWT_ALGORITHM': 'HS256',
     'JWT_VERIFY': True,
     'JWT_VERIFY_EXPIRATION': True,
     'JWT_LEEWAY': 0,
-    'JWT_EXPIRATION_DELTA': datetime.timedelta(seconds=300)
+    'JWT_EXPIRATION_DELTA': datetime.timedelta(seconds=300),
+
+    'JWT_ALLOW_REFRESH': False,
+    'JWT_REFRESH_EXPIRATION_DELTA': datetime.timedelta(days=7),
 }
 ```
 This packages uses the JSON Web Token Python implementation, [PyJWT](https://github.com/progrium/pyjwt) and allows to modify some of it's available options.
@@ -126,8 +148,21 @@ Default is `True`.
 
 Default is `0` seconds.
 
-
 ### JWT_EXPIRATION_DELTA
 This is an instance of Python's `datetime.timedelta`. This will be added to `datetime.utcnow()` to set the expiration time.
 
 Default is `datetime.timedelta(seconds=300)`(5 minutes).
+
+### JWT_ALLOW_REFRESH
+Enable token refresh functionality. Token issued from `rest_framework_jwt.views.obtain_jwt_token` will have an `orig_iat` field. Default is `False`
+
+### JWT_REFRESH_EXPIRATION_DELTA
+Limit on token refresh, is a `datetime.timedelta` instance. This is how much time after the original token that future tokens can be refreshed from.
+
+Default is `datetime.timedelta(days=7)` (7 days).
+
+### JWT_PAYLOAD_HANDLER
+Specify a custom function to generate the token payload
+
+### JWT_PAYLOAD_GET_USER_ID_HANDLER
+If you store `user_id` differently than the default payload handler does, implement this function to fetch `user_id` from the payload.

rest_framework_jwt/authentication.py

@@ -13,6 +13,7 @@
 
 
 jwt_decode_handler = api_settings.JWT_DECODE_HANDLER
+jwt_get_user_id_from_payload = api_settings.JWT_PAYLOAD_GET_USER_ID_HANDLER
 
 
 class JSONWebTokenAuthentication(BaseAuthentication):
@@ -62,7 +63,7 @@ def authenticate_credentials(self, payload):
         Returns an active user that matches the payload's user id and email.
         """
         try:
-            user_id = payload.get('user_id')
+            user_id = jwt_get_user_id_from_payload(payload)
 
             if user_id:
                 user = User.objects.get(pk=user_id, is_active=True)

rest_framework_jwt/serializers.py

@@ -1,3 +1,7 @@
+from calendar import timegm
+from datetime import datetime, timedelta
+import jwt
+
 from django.contrib.auth import authenticate, get_user_model
 from rest_framework import serializers
 
@@ -6,6 +10,8 @@
 
 jwt_payload_handler = api_settings.JWT_PAYLOAD_HANDLER
 jwt_encode_handler = api_settings.JWT_ENCODE_HANDLER
+jwt_decode_handler = api_settings.JWT_DECODE_HANDLER
+jwt_get_user_id_from_payload = api_settings.JWT_PAYLOAD_GET_USER_ID_HANDLER
 
 
 class JSONWebTokenSerializer(serializers.Serializer):
@@ -41,6 +47,13 @@ def validate(self, attrs):
 
                 payload = jwt_payload_handler(user)
 
+                # Include original issued at time for a brand new token,
+                # to allow token refresh
+                if api_settings.JWT_ALLOW_REFRESH:
+                    payload['orig_iat'] = timegm(
+                        datetime.utcnow().utctimetuple()
+                    )
+
                 return {
                     'token': jwt_encode_handler(payload)
                 }
@@ -50,3 +63,61 @@ def validate(self, attrs):
         else:
             msg = 'Must include "username" and "password"'
             raise serializers.ValidationError(msg)
+
+
+class RefreshJSONWebTokenSerializer(serializers.Serializer):
+    """
+    Check an access token
+    """
+    token = serializers.CharField()
+
+    def validate(self, attrs):
+        token = attrs['token']
+
+        # Check payload valid (based off of JSONWebTokenAuthentication,
+        # may want to refactor)
+        try:
+            payload = jwt_decode_handler(token)
+        except jwt.ExpiredSignature:
+            msg = 'Signature has expired.'
+            raise serializers.ValidationError(msg)
+        except jwt.DecodeError:
+            msg = 'Error decoding signature.'
+            raise serializers.ValidationError(msg)
+
+        # Make sure user exists (may want to refactor this)
+        User = get_user_model()
+        try:
+            user_id = jwt_get_user_id_from_payload(payload)
+            if user_id:
+                user = User.objects.get(pk=user_id, is_active=True)
+            else:
+                msg = 'Invalid payload'
+                raise serializers.ValidationError(msg)
+        except User.DoesNotExist:
+            raise serializers.ValidationError("User doesn't exist")
+
+        # Get and check 'orig_iat'
+        orig_iat = payload.get('orig_iat')
+        if orig_iat:
+            # Verify expiration
+            refresh_limit = api_settings.JWT_REFRESH_EXPIRATION_DELTA
+            if isinstance(refresh_limit, timedelta):
+                refresh_limit = (refresh_limit.days * 24 * 3600 +
+                                 refresh_limit.seconds)
+            expiration_timestamp = (
+                orig_iat +
+                int(refresh_limit)
+            )
+            now_timestamp = timegm(datetime.utcnow().utctimetuple())
+            if now_timestamp > expiration_timestamp:
+                raise serializers.ValidationError("Refresh has expired")
+        else:
+            raise serializers.ValidationError("orig_iat field is required")
+
+        new_payload = jwt_payload_handler(user)
+        new_payload['orig_iat'] = orig_iat
+
+        return {
+            'token': jwt_encode_handler(new_payload)
+        }

rest_framework_jwt/settings.py

@@ -16,19 +16,26 @@
     'JWT_PAYLOAD_HANDLER':
     'rest_framework_jwt.utils.jwt_payload_handler',
 
+    'JWT_PAYLOAD_GET_USER_ID_HANDLER':
+    'rest_framework_jwt.utils.jwt_get_user_id_from_payload_handler',
+
     'JWT_SECRET_KEY': settings.SECRET_KEY,
     'JWT_ALGORITHM': 'HS256',
     'JWT_VERIFY': True,
     'JWT_VERIFY_EXPIRATION': True,
     'JWT_LEEWAY': 0,
-    'JWT_EXPIRATION_DELTA': datetime.timedelta(seconds=300)
+    'JWT_EXPIRATION_DELTA': datetime.timedelta(seconds=300),
+
+    'JWT_ALLOW_REFRESH': False,
+    'JWT_REFRESH_EXPIRATION_DELTA': datetime.timedelta(days=7),
 }
 
 # List of settings that may be in string import notation.
 IMPORT_STRINGS = (
     'JWT_ENCODE_HANDLER',
     'JWT_DECODE_HANDLER',
     'JWT_PAYLOAD_HANDLER',
+    'JWT_PAYLOAD_GET_USER_ID_HANDLER',
 )
 
 api_settings = APISettings(USER_SETTINGS, DEFAULTS, IMPORT_STRINGS)

rest_framework_jwt/tests/test_views.py

@@ -1,20 +1,29 @@
+from calendar import timegm
+from datetime import datetime, timedelta
+import time
+
 from django.test import TestCase
 from django.conf import settings
 from django.contrib.auth.models import User
+
 from rest_framework import status
 from rest_framework.compat import patterns
 from rest_framework.test import APIClient
 
 from rest_framework_jwt import utils
 from rest_framework_jwt.runtests.models import CustomUser
+from rest_framework_jwt.settings import api_settings, DEFAULTS
 
 urlpatterns = patterns(
     '',
     (r'^auth-token/$', 'rest_framework_jwt.views.obtain_jwt_token'),
+    (r'^auth-token-refresh/$', 'rest_framework_jwt.views.refresh_jwt_token'),
 )
 
+orig_datetime = datetime
+
 
-class ObtainJSONWebTokenTests(TestCase):
+class BaseTestCase(TestCase):
     urls = 'rest_framework_jwt.tests.test_views'
 
     def setUp(self):
@@ -29,6 +38,9 @@ def setUp(self):
             'password': self.password
         }
 
+
+class ObtainJSONWebTokenTests(BaseTestCase):
+
     def test_jwt_login_json(self):
         """
         Ensure JWT login view using JSON POST works.
@@ -145,3 +157,101 @@ def test_jwt_login_json_bad_creds(self):
 
     def tearDown(self):
         settings.AUTH_USER_MODEL = self.ORIG_AUTH_USER_MODEL
+
+
+class RefreshJSONWebTokenTests(BaseTestCase):
+    urls = 'rest_framework_jwt.tests.test_views'
+
+    def setUp(self):
+        super(RefreshJSONWebTokenTests, self).setUp()
+        api_settings.JWT_ALLOW_REFRESH = True
+
+    def get_token(self):
+        client = APIClient(enforce_csrf_checks=True)
+        response = client.post('/auth-token/', self.data, format='json')
+        return response.data['token']
+
+    def create_token(self, user, exp=None, orig_iat=None):
+        payload = utils.jwt_payload_handler(self.user)
+        if exp:
+            payload['exp'] = exp
+
+        if orig_iat:
+            payload['orig_iat'] = timegm(orig_iat.utctimetuple())
+
+        token = utils.jwt_encode_handler(payload)
+        return token
+
+    def test_refresh_jwt(self):
+        """
+        Test getting a refreshed token from original token works
+        """
+        client = APIClient(enforce_csrf_checks=True)
+
+        orig_token = self.get_token()
+        orig_token_decoded = utils.jwt_decode_handler(orig_token)
+
+        expected_orig_iat = timegm(datetime.utcnow().utctimetuple())
+
+        # Make sure 'orig_iat' exists and is the current time (give some slack)
+        orig_iat = orig_token_decoded['orig_iat']
+        self.assertLessEqual(orig_iat - expected_orig_iat, 1)
+
+        # wait a few seconds, so new token will have different exp
+        time.sleep(2)
+
+        # Now try to get a refreshed token
+        response = client.post('/auth-token-refresh/', {'token': orig_token},
+                               format='json')
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+        new_token = response.data['token']
+        new_token_decoded = utils.jwt_decode_handler(new_token)
+
+        # Make sure 'orig_iat' on the new token is same as original
+        self.assertEquals(new_token_decoded['orig_iat'], orig_iat)
+        self.assertGreater(new_token_decoded['exp'], orig_token_decoded['exp'])
+
+    def test_refresh_jwt_fails_with_expired_token(self):
+        """
+        Test that using an expired token to refresh won't work
+        """
+        client = APIClient(enforce_csrf_checks=True)
+
+        # Make an expired token..
+        token = self.create_token(
+            self.user,
+            exp=datetime.utcnow() - timedelta(seconds=5),
+            orig_iat=datetime.utcnow() - timedelta(hours=1)
+        )
+
+        response = client.post('/auth-token-refresh/', {'token': token},
+                               format='json')
+        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+        self.assertRegexpMatches(response.data['non_field_errors'][0],
+                                 'Signature has expired')
+
+    def test_refresh_jwt_after_refresh_expiration(self):
+        """
+        Test that token can't be refreshed after token refresh limit
+        """
+        client = APIClient(enforce_csrf_checks=True)
+
+        orig_iat = (datetime.utcnow() - api_settings.JWT_REFRESH_EXPIRATION_DELTA -
+                    timedelta(seconds=5))
+        token = self.create_token(
+            self.user,
+            exp=datetime.utcnow() + timedelta(hours=1),
+            orig_iat=orig_iat
+        )
+
+        response = client.post('/auth-token-refresh/', {'token': token},
+                               format='json')
+        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+        self.assertEqual(response.data['non_field_errors'][0],
+                         'Refresh has expired')
+
+    def tearDown(self):
+        # Restore original settings
+        api_settings.JWT_ALLOW_REFRESH = \
+            DEFAULTS['JWT_ALLOW_REFRESH']

rest_framework_jwt/utils.py

@@ -1,4 +1,4 @@
-import datetime
+from datetime import datetime
 import jwt
 
 from rest_framework_jwt.settings import api_settings
@@ -9,10 +9,18 @@ def jwt_payload_handler(user):
         'user_id': user.pk,
         'email': user.email,
         'username': user.get_username(),
-        'exp': datetime.datetime.utcnow() + api_settings.JWT_EXPIRATION_DELTA
+        'exp': datetime.utcnow() + api_settings.JWT_EXPIRATION_DELTA
     }
 
 
+def jwt_get_user_id_from_payload_handler(payload):
+    """
+    Override this function if user_id is formatted differently in payload
+    """
+    user_id = payload.get('user_id')
+    return user_id
+
+
 def jwt_encode_handler(payload):
     return jwt.encode(
         payload,