How to implement Auto Expiring Token in Django Rest Framework

Introduction

Today I will be talking about one of the crucial features you might have found missing in the Django Rest Framework, and that is the ability to create authentication tokens that expire after a set period, let's say after 20-25 days, every user that had authenticated previously gets prompted to re-authenticate again, which means the client will have to use the REST API to obtain the auth token once again.

One way you might have tried to solve this problem is by reaching out for some JWT token package, but I thought there must be a way to do this without having to switch to an entirely different type of authentication token system. That's when I found the solution that I am going to share with you all today.

Background

First of all, when I utilize token authentication I usually make use of a package called dj-rest-auth[^1] which provides me with endpoints for password management (forgot password, reset and change) and authentication (login, logout and register). So part of my solution will include changes to the configuration of that package.

Let's start by describing our solution, by default the tokens generated by Django Rest Framework are valid for a lifetime and a user can have only one authentication token at a time, now upon logging out you can either choose to just clear the cookies and still maintain that old authentication token in the database, which is a security flaw in itself, or implement a logout endpoint that will delete that token. If you were to choose the latter path, you will soon face another issue, which is, if the user decided to log into their account on multiple devices, then they will be using the same authentication token, and if you delete that token, their sessions on both the devices will get logged out, which is something undesirable.

To be able to resolve the above-mentioned problem, we will be modifying the relation between the Token and User to be a ForeignKey instead, that way, a user can have multiple tokens (equivalent to having multiple sessions on different devices).

So let's begin.

Custom Token Model

The following will be the definition of our Token model:

from django.conf import settings
from django.db import models
from rest_framework.authtoken.models import Token as AuthToken


class Token(AuthToken):
    key = models.CharField("Key", max_length=40, db_index=True, unique=True)
    user = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        related_name="auth_token",
        on_delete=models.CASCADE,
        verbose_name="User",
    )

For your reference, I defined the above model in an app called accounts, which is where I usually define my custom user model as well.

Now since I use dj-rest-auth, inside by project settings, I will define the following:

REST_AUTH_TOKEN_MODEL = "accounts.models.Token"

This setting was described in the package's configuration docs here..

Next up we will be creating a utility function used by dj-rest-auth for creating tokens:

import datetime

import pytz
from django.utils import timezone


def custom_create_token(token_model, user, serializer):
    token = token_model.objects.create(user=user)
    utc_now = timezone.now()
    utc_now = utc_now.replace(tzinfo=pytz.utc)
    token.created = utc_now
    token.save()
    return token

I saved the above function inside accounts/utils.py.

After this, define the following inside your project settings:

REST_AUTH_TOKEN_CREATOR = "accounts.utils.custom_create_token"

The above setting is also a configuration required by the dj-rest-auth package.

Now let's pause and explain what we have achieved so far. We have implemented the functionality to be able to create multiple tokens for a single user, meaning they can log onto multiple devices with each session being unique.

Migrations

If you were using the Token model provided by Django Rest Framework earlier, then you will need to make sure you remove the installed app for the token provided by it, to do that, go inside your project settings, and remove the following line from the INSTALLED_APPS:

INSTALLED_APPS = [
    ...
    'rest_framework.authtoken' # <== This line should be removed
]

When you're done with that, let's create the migrations for the custom Token model we defined earlier, to do that run the following command in the terminal:

python manage.py makemigrations accounts

We used accounts here because our model resides in accounts app, change it accordingly as per your implementation. This will create the migrations, now migrate using:

python manage.py migrate accounts

NOTE: If you had any tokens previously created using the default Token model, then those tokens will become invalid later as we will be switching out the Authentication Backend as well.

Custom Authentication Backend

Now we are going to be defining a custom Authentication Backend[^2] which will be derived from the default TokenAuthentication available in Django Rest Framework, and have some extended functionality.

Following is the code for our authentication backend:

import datetime

import pytz
from django.conf import settings
from django.utils import timezone
from rest_framework.authentication import TokenAuthentication
from rest_framework.exceptions import AuthenticationFailed

from accounts.models import Token


class ExpiringTokenAuthentication(TokenAuthentication):
    """
    Expiring token for mobile and desktop clients.
    It expires every {n} hrs requiring client to supply valid username 
    and password for new one to be created.
    """

    model = Token

    def authenticate_credentials(self, key, request=None):
        models = self.get_model()

        try:
            token = models.objects.select_related("user").get(key=key)
        except models.DoesNotExist:
            raise AuthenticationFailed(
                {"error": "Invalid or Inactive Token", "is_authenticated": False}
            )

        if not token.user.is_active:
            raise AuthenticationFailed(
                {"error": "Invalid user", "is_authenticated": False}
            )

        utc_now = timezone.now()
        utc_now = utc_now.replace(tzinfo=pytz.utc)

        if token.created < utc_now - settings.TOKEN_TTL:
            raise AuthenticationFailed(
                {"error": "Token has expired", "is_authenticated": False}
            )
        return token.user, token

I wrote the above code inside accounts/authentication.py.

To be able to use the above-mentioned authentication backend, you will need to make changes to your REST_FRAMEWORK config in your project settings as follows:

REST_FRAMEWORK = {
    "DEFAULT_AUTHENTICATION_CLASSES": (
        "rest_framework.authentication.TokenAuthentication", # Replace this line
        "accounts.authentication.ExpiringTokenAuthentication", # <-- with this line 
        "rest_framework.authentication.SessionAuthentication",
    ),
 # ...
}

If you noticed carefully, in the code for our custom authentication backend, I mentioned settings.TOKEN_TTL, which is another functionality that I have added to allow configuring the Token's Lifespan or Token - Time to Live through project settings instead of setting it to a fixed value.

Now under your project settings, define the following:

import datetime

TOKEN_TTL = datetime.timedelta(days=15)

I have set the token lifetime to 15 days, but you are free to configure it as per your choice.

That's it for this article, please let me know if you found this useful in the comments below.

References

1. Custom Authentication in DRF

2. Dj-Rest-Auth