EgCoder | Django Rest Framework custom JWT authentication

This article is not a tutorial or a guide, it is more like a request for code review and validate the implementation from more experienced Django developers, so please don't use this code unless you are able to review and validate yourself

Introduction

Web application security is critical and doing it on your own in production is a huge responsibility, that is why I love Django because it takes care of most of the security vulnerabilities out there, but my problem shows up when I wanted to build my back-end as a restful API so I can connect all type of clients (SPA, mobile, desktop, ... etc) Django Rest Framework (DRF) comes with different builtin authentication classes, token authentication or JWT are the way to go for my use case but I still have that worry of how to save the tokens in client-side everybody says don't save the token in localstorage because of XSS attacks and better to save your token in httponly cookie, but cookies are open to CSRF attack too and DRF disable CSRF protection for all the APIView so what is the best practice to do this. for a while I kept using it as everybody especially in mobile (i do react-native) we have secure storages and also if the device is not (rooted android or jailbroken ios) then each app is sandboxed and the tokens are safe in most cases The problem still exists in web clients (SPAs), so I came into an implementation that might be useful and I wanted to document it here so I get feedback from more experienced developers, I can summarize my implementation into the following steps:

the user sends a POST request with username and password to log in, then the server will do 3 things
- generate an access_token which is a short life jwt (maybe 5 mins) and send it in the response body
- generate a refresh_token which is a long life jwt (days) and send it in an httponly cookie, so it won't be accessible from the client javascript
- send a normal cookie that contains a CSRF token
The developer needs to be sure that all unsafe views (POST, UPDATE, PUT, Delete) are protected by the builtin Django CSRF protection because as I mentioned above DRF disables them by default.
in the client-side the developer should take care of:
- In the client-side, every request will contain refresh token in the cookies automatically (be sure that your client domain is whitelisted in the server cors headers settings)
- Send the access_token in the Authorization header.
- Send the CSRF token in the X-CSRFTOKEN header if he is doing POST request
- when he needs a new access_token, he need to send a POST request to refresh token endpoint

enough talking, let us see some code

Project Setup

python3 -m venv .venv
source .venv/bin/activate
pip install django django-cors-headers djangorestframework PyJWT

# create Django project with the venv activates
django-admin startproject project

# create an app
./manage.py start app accounts

in the project settings enable the apps and add some settings

INSTALLED_APPS = [
    ...
    # 3rd party apps
    'corsheaders',
    'rest_framework',

    # project apps
    'accounts',
]

CORS_ALLOW_CREDENTIALS = True # to accept cookies via ajax request
CORS_ORIGIN_WHITELIST = [
    'http://localhost:3000' # the domain for front-end app(you can add more than 1)
]

REST_FRAMEWORK = {
    'DEFAULT_PERMISSION_CLASSES': (
        'rest_framework.permissions.IsAuthenticated', # make all endpoints private
    )
}

before applying the first migration and as recommended in the official Django docs I like to start my project with a custom user model even if I won't use it now in accounts.models define the user model

from django.contrib.auth.models import AbstractUser

class User(AbstractUser):
    pass

in the project.settings.py define the user model

...
AUTH_USER_MODEL = 'accounts.User'
...

later you can refer to the user model by one of the following

from django.conf import settings
User = settings.AUTH_USER_MODEL

# OR

from django.contrib.auth import get_user_model
User = get_user_model()

and in accounts.admin register the new user model to the admin site

from django.contrib import admin
from django.contrib.auth.admin import UserAdmin
from accounts.models import User

admin.site.register(User, UserAdmin)

create a super user

./manage.py createsuperuser

and finally, an endpoint to test, as we declare in the above settings all the endpoint will require authentication by default, we can override this in certain views as we will going to do in the login later

we will create the user profile endpoint that will return current authenticated user object as JSON, for that we will need to create a user serializer

# accounts.serializers
from rest_framework import serializers
from accounts.models import User


class UserSerializer(serializers.ModelSerializer):
    class Meta:
        model = User
        fields = ['id', 'username', 'email',
                  'first_name', 'last_name', 'is_active']

# project.urls
from accounts import urls as accounts_urls

urlpatterns = [
    path('accounts/', include(accounts_urls)),
]

# accounts.urls
urlpatterns = [
    path('profile', profile, name='profile'),
]

# accounts.views
from rest_framework.decorators import api_view
from rest_framework.response import Response
from .serializers import UserSerializer


@api_view(['GET'])
def profile(request):
    user = request.user
    serialized_user = UserSerializer(user).data
    return Response({'user': serialized_user })

now if you tried to acceess this endpoint you will get 403 error, as described in the introduction we need to login then send the access_token in the request header

The login endpoint will be a post request with username and password in the request body. we will make the login view public using the permission class decorator AllowAny, also we will decorate the view with @ensure_csrf_cookie forcing Django to send the CSRF cookie in the response if the login success

if the login success, we will have:

an access_token in the response body.
a refreshtoken in an httponly cookie.
a csrftoken in a normal cookie so we can read it from javascript and resend it when needed

# accounts.views
from django.contrib.auth import get_user_model
from rest_framework.response import Response
from rest_framework import exceptions
from rest_framework.permissions import AllowAny
from rest_framework.decorators import api_view, permission_classes
from django.views.decorators.csrf import ensure_csrf_cookie
from accounts.serializers import UserSerializer
from accounts.utils import generate_access_token, generate_refresh_token


@api_view(['POST'])
@permission_classes([AllowAny])
@ensure_csrf_cookie
def login_view(request):
    User = get_user_model()
    username = request.data.get('username')
    password = request.data.get('password')
    response = Response()
    if (username is None) or (password is None):
        raise exceptions.AuthenticationFailed(
            'username and password required')

    user = User.objects.filter(username=username).first()
    if(user is None):
        raise exceptions.AuthenticationFailed('user not found')
    if (not user.check_password(password)):
        raise exceptions.AuthenticationFailed('wrong password')

    serialized_user = UserSerializer(user).data

    access_token = generate_access_token(user)
    refresh_token = generate_refresh_token(user)

    response.set_cookie(key='refreshtoken', value=refresh_token, httponly=True)
    response.data = {
        'access_token': access_token,
        'user': serialized_user,
    }

    return response

and here are the functions that generate the tokens, notice that I use a different secret to sign refresh token for more security

# accounts.utils
import datetime
import jwt
from django.conf import settings


def generate_access_token(user):

    access_token_payload = {
        'user_id': user.id,
        'exp': datetime.datetime.utcnow() + datetime.timedelta(days=0, minutes=5),
        'iat': datetime.datetime.utcnow(),
    }
    access_token = jwt.encode(access_token_payload,
                              settings.SECRET_KEY, algorithm='HS256').decode('utf-8')
    return access_token


def generate_refresh_token(user):
    refresh_token_payload = {
        'user_id': user.id,
        'exp': datetime.datetime.utcnow() + datetime.timedelta(days=7),
        'iat': datetime.datetime.utcnow()
    }
    refresh_token = jwt.encode(
        refresh_token_payload, settings.REFRESH_TOKEN_SECRET, algorithm='HS256').decode('utf-8')

    return refresh_token

Custom Authentication Class for DRF

Django Rest Framework makes it easy to create a custom authentication scheme, it described in details in the official docs The following code is originally taken from DRF source code then I add my changes as required. notice that DRF enforce CSRF only in the session authentication rest_framework/authentication.py

# accounts.authentication

import jwt
from rest_framework.authentication import BaseAuthentication
from django.middleware.csrf import CsrfViewMiddleware
from rest_framework import exceptions
from django.conf import settings
from django.contrib.auth import get_user_model


class CSRFCheck(CsrfViewMiddleware):
    def _reject(self, request, reason):
        # Return the failure reason instead of an HttpResponse
        return reason


class SafeJWTAuthentication(BaseAuthentication):
    '''
        custom authentication class for DRF and JWT
        https://github.com/encode/django-rest-framework/blob/master/rest_framework/authentication.py
    '''

    def authenticate(self, request):

        User = get_user_model()
        authorization_heaader = request.headers.get('Authorization')

        if not authorization_heaader:
            return None
        try:
            # header = 'Token xxxxxxxxxxxxxxxxxxxxxxxx'
            access_token = authorization_heaader.split(' ')[1]
            payload = jwt.decode(
                access_token, settings.SECRET_KEY, algorithms=['HS256'])

        except jwt.ExpiredSignatureError:
            raise exceptions.AuthenticationFailed('access_token expired')
        except IndexError:
            raise exceptions.AuthenticationFailed('Token prefix missing')

        user = User.objects.filter(id=payload['user_id']).first()
        if user is None:
            raise exceptions.AuthenticationFailed('User not found')

        if not user.is_active:
            raise exceptions.AuthenticationFailed('user is inactive')

        self.enforce_csrf(request)
        return (user, None)

    def enforce_csrf(self, request):
        """
        Enforce CSRF validation
        """
        check = CSRFCheck()
        # populates request.META['CSRF_COOKIE'], which is used in process_view()
        check.process_request(request)
        reason = check.process_view(request, None, (), {})
        print(reason)
        if reason:
            # CSRF failed, bail with explicit error message
            raise exceptions.PermissionDenied('CSRF Failed: %s' % reason)

once you create that class, go to the project.settings and activate it in the REST_FRAMEWORK section like this appname.filename.classname


REST_FRAMEWORK = {
    'DEFAULT_AUTHENTICATION_CLASSES': (
        'accounts.authentication.SafeJWTAuthentication',
    ),
    'DEFAULT_PERMISSION_CLASSES': (
        'rest_framework.permissions.IsAuthenticated',
    )
}

Now as we have access_token and set the authentication method we can revisit the profile endpoint but this time we will set the Authorization header

Refresh Token View

Whenever the token is expired or you need a new token for any reason we need a refresh_token endpoint. this view needs the permission of AlloAny because we don't have an access_token but it will be protected by 2 other things

a valid refresh_token that sent in the httoponly cookie.
the CSRF token so we are sure the above cookie not compromised if the 2 conditions are satisfied then the server will generate a new valid access_token and send it back.

if the refresh_token is invalid or expired, the user will need to re-logion

import jwt
from django.conf import settings
from django.contrib.auth import get_user_model
from django.views.decorators.csrf import csrf_protect
from rest_framework import exceptions
from rest_framework.response import Response
from rest_framework.permissions import AllowAny
from rest_framework.decorators import api_view, permission_classes
from accounts.utils import generate_access_token


@api_view(['POST'])
@permission_classes([AllowAny])
@csrf_protect
def refresh_token_view(request):
    '''
    To obtain a new access_token this view expects 2 important things:
        1. a cookie that contains a valid refresh_token
        2. a header 'X-CSRFTOKEN' with a valid csrf token, client app can get it from cookies "csrftoken"
    '''
    User = get_user_model()
    refresh_token = request.COOKIES.get('refreshtoken')
    if refresh_token is None:
        raise exceptions.AuthenticationFailed(
            'Authentication credentials were not provided.')
    try:
        payload = jwt.decode(
            refresh_token, settings.REFRESH_TOKEN_SECRET, algorithms=['HS256'])
    except jwt.ExpiredSignatureError:
        raise exceptions.AuthenticationFailed(
            'expired refresh token, please login again.')

    user = User.objects.filter(id=payload.get('user_id')).first()
    if user is None:
        raise exceptions.AuthenticationFailed('User not found')

    if not user.is_active:
        raise exceptions.AuthenticationFailed('user is inactive')


    access_token = generate_access_token(user)
    return Response({'access_token': access_token})

Token Revoke

The last piece of the puzzle is a way to revoke the refresh token which has a long lifetime, you might blacklist the token or assign a uuid for the token and put it in the payload then link it to the user and save it in the database, when revoking or logout you just change that uuid in the database to not match the value in the payload, you can pick what suits your application need

Django Rest Framework custom JWT authentication

Introduction

Project Setup

Login View

Custom Authentication Class for DRF

Refresh Token View

Token Revoke