Merge pull request #62 from erichonkanen/master

Added jwt_response_payload_handler which allows returning additional response data

Author: jpadilla

README.md

@@ -28,11 +28,11 @@ In your `settings.py`, add `JSONWebTokenAuthentication` to Django REST framework
 ```python
 REST_FRAMEWORK = {
     'DEFAULT_PERMISSION_CLASSES': (
-    	'rest_framework.permissions.IsAuthenticated',
+        'rest_framework.permissions.IsAuthenticated',
     ),
     'DEFAULT_AUTHENTICATION_CLASSES': (
-    	'rest_framework.authentication.SessionAuthentication',
-	    'rest_framework.authentication.BasicAuthentication',
+        'rest_framework.authentication.SessionAuthentication',
+        'rest_framework.authentication.BasicAuthentication',
         'rest_framework_jwt.authentication.JSONWebTokenAuthentication',
     ),
 }
@@ -100,6 +100,9 @@ JWT_AUTH = {
     'JWT_PAYLOAD_GET_USER_ID_HANDLER':
     'rest_framework_jwt.utils.jwt_get_user_id_from_payload_handler',
 
+    'JWT_RESPONSE_PAYLOAD_HANDLER':
+    'rest_framework_jwt.utils.jwt_response_payload_handler',
+
     'JWT_SECRET_KEY': settings.SECRET_KEY,
     'JWT_ALGORITHM': 'HS256',
     'JWT_VERIFY': True,
@@ -175,6 +178,22 @@ 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.
 
+### JWT_RESPONSE_PAYLOAD_HANDLER
+Responsible for controlling the response data returned after login or refresh. Override to return a custom response such as including the serialized representation of the User.
+
+Defaults to return the JWT token.
+
+Example:
+```
+def jwt_response_payload_handler(token, user=None):
+    return {
+        'token': token,
+        'user': UserSerializer(user).data
+    }
+```
+
+Default is `{'token': token}`
+
 ### JWT_AUTH_HEADER_PREFIX
 You can modify the Authorization header value prefix that is required to be sent together with the token. The default value is `JWT`. This decision was introduced in PR [#4](https://github.com/GetBlimp/django-rest-framework-jwt/pull/4) to allow using both this package and OAuth2 in DRF.
 

rest_framework_jwt/serializers.py

@@ -1,6 +1,7 @@
+import jwt
+
 from calendar import timegm
 from datetime import datetime, timedelta
-import jwt
 
 from django.contrib.auth import authenticate
 from django.utils.translation import ugettext as _
@@ -63,7 +64,8 @@ def validate(self, attrs):
                     )
 
                 return {
-                    'token': jwt_encode_handler(payload)
+                    'token': jwt_encode_handler(payload),
+                    'user': user
                 }
             else:
                 msg = _('Unable to login with provided credentials.')
@@ -133,5 +135,6 @@ def validate(self, attrs):
         new_payload['orig_iat'] = orig_iat
 
         return {
-            'token': jwt_encode_handler(new_payload)
+            'token': jwt_encode_handler(new_payload),
+            'user': user
         }

rest_framework_jwt/settings.py

@@ -19,6 +19,9 @@
     'JWT_PAYLOAD_GET_USER_ID_HANDLER':
     'rest_framework_jwt.utils.jwt_get_user_id_from_payload_handler',
 
+    'JWT_RESPONSE_PAYLOAD_HANDLER':
+    'rest_framework_jwt.utils.jwt_response_payload_handler',
+
     'JWT_SECRET_KEY': settings.SECRET_KEY,
     'JWT_ALGORITHM': 'HS256',
     'JWT_VERIFY': True,
@@ -38,6 +41,7 @@
     'JWT_DECODE_HANDLER',
     'JWT_PAYLOAD_HANDLER',
     'JWT_PAYLOAD_GET_USER_ID_HANDLER',
+    'JWT_RESPONSE_PAYLOAD_HANDLER',
 )
 
 api_settings = APISettings(USER_SETTINGS, DEFAULTS, IMPORT_STRINGS)

rest_framework_jwt/utils.py

@@ -1,6 +1,7 @@
-from datetime import datetime
 import jwt
 
+from datetime import datetime
+
 from rest_framework_jwt.settings import api_settings
 
 
@@ -53,3 +54,23 @@ def jwt_decode_handler(token):
         verify_expiration=api_settings.JWT_VERIFY_EXPIRATION,
         leeway=api_settings.JWT_LEEWAY
     )
+
+
+def jwt_response_payload_handler(token, user=None):
+    """
+    Returns the response data for both the login and refresh views.
+    Override to return a custom response such as including the
+    serialized representation of the User.
+
+    Example:
+
+    def jwt_response_payload_handler(token, user=None):
+        return {
+            'token': token,
+            'user': UserSerializer(user).data
+        }
+
+    """
+    return {
+        'token': token
+    }

rest_framework_jwt/views.py

@@ -4,8 +4,12 @@
 from rest_framework import renderers
 from rest_framework.response import Response
 
+from rest_framework_jwt.settings import api_settings
+
 from .serializers import JSONWebTokenSerializer, RefreshJSONWebTokenSerializer
 
+jwt_response_payload_handler = api_settings.JWT_RESPONSE_PAYLOAD_HANDLER
+
 
 class ObtainJSONWebToken(APIView):
     """
@@ -23,7 +27,10 @@ class ObtainJSONWebToken(APIView):
     def post(self, request):
         serializer = self.serializer_class(data=request.DATA)
         if serializer.is_valid():
-            return Response({'token': serializer.object['token']})
+            user = serializer.object.get('user') or request.user
+            token = serializer.object.get('token')
+            response_data = jwt_response_payload_handler(token, user)
+            return Response(response_data)
         return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
 
 
@@ -45,7 +52,10 @@ class RefreshJSONWebToken(APIView):
     def post(self, request):
         serializer = self.serializer_class(data=request.DATA)
         if serializer.is_valid():
-            return Response({'token': serializer.object['token']})
+            user = serializer.object.get('user') or request.user
+            token = serializer.object.get('token')
+            response_data = jwt_response_payload_handler(token, user)
+            return Response(response_data)
         return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
 
 

tests/test_authentication.py

@@ -1,8 +1,8 @@
-from django.contrib.auth.models import User
 from django.http import HttpResponse
 from django.test import TestCase
 from django.utils import unittest
 from django.conf.urls import patterns
+from django.contrib.auth import get_user_model
 
 from rest_framework import permissions, status
 from rest_framework.authentication import OAuth2Authentication
@@ -14,6 +14,7 @@
 from rest_framework_jwt.settings import api_settings, DEFAULTS
 from rest_framework_jwt.authentication import JSONWebTokenAuthentication
 
+User = get_user_model()
 
 DJANGO_OAUTH2_PROVIDER_NOT_INSTALLED = 'django-oauth2-provider not installed'
 

tests/test_serializers.py

@@ -3,11 +3,12 @@
 import rest_framework
 from django.test import TestCase
 from django.utils import unittest
-from django.contrib.auth.models import User
+from django.contrib.auth import get_user_model
 
 from rest_framework_jwt.serializers import JSONWebTokenSerializer
 from rest_framework_jwt import utils
 
+User = get_user_model()
 
 drf2 = rest_framework.VERSION < StrictVersion('3.0.0')
 drf3 = rest_framework.VERSION >= StrictVersion('3.0.0')

tests/test_utils.py

@@ -1,10 +1,12 @@
 import json
 from jwt import base64url_decode
 
+from django.contrib.auth import get_user_model
 from django.test import TestCase
-from django.contrib.auth.models import User
 from rest_framework_jwt import utils
 
+User = get_user_model()
+
 
 class UtilsTests(TestCase):
     def setUp(self):
@@ -36,3 +38,10 @@ def test_jwt_decode(self):
         decoded_payload = utils.jwt_decode_handler(token)
 
         self.assertEqual(decoded_payload, payload)
+
+    def test_jwt_response_payload(self):
+        payload = utils.jwt_payload_handler(self.user)
+        token = utils.jwt_encode_handler(payload)
+        response_data = utils.jwt_response_payload_handler(token, self.user)
+
+        self.assertEqual(response_data, dict(token=token))

tests/test_views.py

@@ -6,15 +6,18 @@
 from django.test import TestCase
 from django.test.utils import override_settings
 from django.utils import unittest
-from django.contrib.auth.models import User
 from django.conf.urls import patterns
+from django.contrib.auth import get_user_model
 
 from rest_framework import status
 from rest_framework.test import APIClient
 
 from rest_framework_jwt import utils
 from rest_framework_jwt.settings import api_settings, DEFAULTS
 
+from . import utils as test_utils
+
+User = get_user_model()
 
 NO_CUSTOM_USER_MODEL = 'Custom User Model only supported after Django 1.5'
 
@@ -43,6 +46,32 @@ def setUp(self):
         }
 
 
+class TestCustomResponsePayload(BaseTestCase):
+
+    def setUp(self):
+        api_settings.JWT_RESPONSE_PAYLOAD_HANDLER = test_utils\
+            .jwt_response_payload_handler
+        return super(TestCustomResponsePayload, self).setUp()
+
+    def test_jwt_login_custom_response_json(self):
+        """
+        Ensure JWT login view using JSON POST works.
+        """
+        client = APIClient(enforce_csrf_checks=True)
+
+        response = client.post('/auth-token/', self.data, format='json')
+
+        decoded_payload = utils.jwt_decode_handler(response.data['token'])
+
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertEqual(decoded_payload['username'], self.username)
+        self.assertEqual(response.data['user'], self.username)
+
+    def tearDown(self):
+        api_settings.JWT_RESPONSE_PAYLOAD_HANDLER =\
+            DEFAULTS['JWT_RESPONSE_PAYLOAD_HANDLER']
+
+
 class ObtainJSONWebTokenTests(BaseTestCase):
 
     def test_jwt_login_json(self):

tests/utils.py

@@ -0,0 +1,19 @@
+def jwt_response_payload_handler(token, user=None):
+    """
+    Returns the response data for both the login and refresh views.
+    Override to return a custom response such as including the
+    serialized representation of the User.
+
+    Example:
+
+    def jwt_response_payload_handler(token, user=None):
+        return {
+            'token': token,
+            'user': UserSerializer(user).data
+        }
+
+    """
+    return {
+        'user': user.get_username(),
+        'token': token
+    }