Authentication in WebSocket with Django and Django Rest Framework (DRF)

Table of contents

1. What is authentication?
2. Project Setup
3. Creating necessary models for the project
4. Configuring Django Channels
5. Implementing WebSocket Authentication
6. Testing WebSocket Authentication
7. Conclusion

What is authentication?

Authentication is the process of verifying the identity of a user or system. It ensures that the entity attempting to access a resource is indeed who or what it claims to be. This process is fundamental for security in any system, as it helps to protect sensitive data and functionalities from unauthorized access.

Project Setup

1. Prerequisites : Python >=3.8

2. Installation

Create a folder, virtual environment and activate a virtual environment (Linux/MacOS/Windows)

# Create a folder for the project
mkdir auth_on_websocket

# Create virtual environment
python -m venv venv

# On Linux/MacOS:
source venv/bin/activate

# On Windows:
venv\Scripts\activate

# Install the folowing packages
 pip install django channels daphne djangorestframework djangorestframework-simplejwt

3. Folder Setup

django-admin startproject websocket_auth .

django-admin startapp accounts

django-admin startapp chats

4. Adding necessary settings in setting.py

ASGI_APPLICATION = 'websocket_auth.asgi.application'

INSTALLED_APPS = [
    "daphne", 
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',


    # THIRD_PARTY_APPS
    "rest_framework",
    "rest_framwork_simplejwt"
    "channels",

    # LOCAL_APPS
    "accounts",
    "chats"

]


REST_FRAMEWORK = {
    "DEFAULT_AUTHENTICATION_CLASSES": (
        "rest_framework_simplejwt.authentication.JWTAuthentication",
    ),
     
}


AUTH_USER_MODEL = 'accounts.CustomUser'

Creating necessary models for the project

Create a folder, utils, and create a file base_model.py, then paste the code below

from django.db import models

class BaseModel(models.Model):
    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    created_at = models.DateField(auto_now_add=True)
    updated_at = models.DateField(auto_now=True)

    class Meta:
        abstract = True

also create two apps, accounts, and chats, which can be done using the command below in your terminal, in the project’s folder, auth_over_websocket .

python manage.py <app_name>

In your file accounts/models.py

from django.db import models
from django.contrib.auth.models import AbstractUser
from utils.base_model import BaseModel





class CustomUser(AbstractUser, BaseModel):
    email = models.EmailField(unique=True, blank=False, null=False)
    first_name = models.CharField(max_length=40, blank=False, null=False)


    USERNAME_FIELD = "username"
    
    def __str__(self) -> str:
        return self.email

In your file chats/models.py

from django.db import models
from django.contrib.auth import get_user_model
from utils.base_model import BaseModel

User = get_user_model()



class WebSocketChatID(BaseModel):
    user = models.ForeignKey(User, related_name='chat_user', on_delete=models.CASCADE)
    other_user = models.ForeignKey(User, related_name='chat_other_user', on_delete=models.CASCADE)

    def __str__(self):
        return f" WebsocketChatID between {self.user.username} and {self.other_user.username}"

class Message(BaseModel):
    sender = models.ForeignKey(User, related_name='sent_messages', on_delete=models.CASCADE)
    receiver = models.ForeignKey(User, related_name='received_messages', on_delete=models.CASCADE)

Configuring Django Channels

In order to setup channels, you may follow the configuration below, or refer to the official documentation for more information.

1. Adding channels settings in websocket_auth/settings.py

# SETTINGS, NOTE: use for development
 
CHANNEL_LAYERS = {
    'default': {
        'BACKEND': 'channels.layers.InMemoryChannelLayer'
    }
}

Implementing WebSocket Authentication

Since we are using JWTAuthentication, we will write a custom middleware, JWTAuthMiddleware for authenticating users over WebSocket connection.

In your websocket_auth folder, create a folder chat_config, then add file, jwt_middleware.py, then paste the code below in it.

from channels.middleware import BaseMiddleware
from rest_framework_simplejwt.tokens import AccessToken
from channels.db import database_sync_to_async

class JWTAuthMiddleware(BaseMiddleware):
    
    async def __call__(self, scope, receive, send):
    
        token = self.get_token_from_scope(scope)
        
        if token != None:
            user_id = await self.get_user_from_token(token) 
            if user_id:
                scope['user_id'] = user_id

            else:
                scope['error'] = 'Invalid token'

        if token == None:
            scope['error'] = 'provide an auth token'    
    
                
        return await super().__call__(scope, receive, send)

    def get_token_from_scope(self, scope):
        headers = dict(scope.get("headers", []))
        
        auth_header = headers.get(b'authorization', b'').decode('utf-8')
        
        if auth_header.startswith('Bearer '):
            return auth_header.split(' ')[1]
        
        else:
            return None
        
    @database_sync_to_async
    def get_user_from_token(self, token):
            try:
                access_token = AccessToken(token)
                return access_token['user_id']
            except:
                return None

In your websocket_auth folder, create a folder chat_config, then add file, consumers.py, then paste the code below in it.

import json
from channels.generic.websocket import AsyncWebsocketConsumer
from chats.models import Message
from channels.db import database_sync_to_async
from asgiref.sync import async_to_sync, sync_to_async


class DirectMessageConsumer(AsyncWebsocketConsumer):

    async def connect(self):
        self.id = self.scope['url_route']['kwargs']['chat_id']
        self.chat_room_name = f'chat_{self.id}'

        await self.accept()

        if self.is_error_exists():
            error = {
                'error': str(self.scope['error'])
            }
            await self.send(text_data=json.dumps(error))
            await self.close()

        else:            
            await self.channel_layer.group_add(    
                self.chat_room_name,
                self.channel_name
            )
        
    async def disconnect(self, close_code):
        # Leave room group
        await self.channel_layer.group_discard(
            self.chat_room_name, self.channel_name
        )
    
    # Receive message from WebSocket
    async def receive(self, text_data):
        if self.scope.get('user_id') is not None:
            
            sender_id = self.scope.get('user_id')
            text_data_json = json.loads(text_data)
            message = text_data_json['message']
            receiver_id = text_data_json['receiver_id']

            message = await database_sync_to_async(Message.objects.create)(
                sender_id=sender_id,
                receiver_id=receiver_id,  
                message=message
            )    

            await self.channel_layer.group_send(
                self.chat_room_name,
            {
                'type': 'chat_message',
                'message_info': await self.message_information(message),
                }
            )
        
    async def chat_message(self, event):   
        await self.send(text_data=json.dumps(event))

    @sync_to_async
    def message_information(self, message):
        return {
            "message": message.message,
            "sender": message.sender.username       
            }



    def is_error_exists(self):
        """This checks if error exists during websockets"""

        return True if 'error' in self.scope else False

In the websocket_auth/asgi.py, paste the code below in

import os
from channels.routing import ProtocolTypeRouter, URLRouter

from django.core.asgi import get_asgi_application


os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'websocket_auth.settings')
django_asgi_application = get_asgi_application()

from websocket_auth.chat_config import routing
from websocket_auth.chat_config.jwt_middleware import JWTAuthMiddleware 



application = ProtocolTypeRouter(
    {
        'http': django_asgi_application,
        'websocket': 
        JWTAuthMiddleware(
            URLRouter(
                routing.websocket_urlpatterns
            )
        )
    }
)

Testing WebSocket Connection

1. Using Postman, create a WebSocket connection from the sidebar, click on New which pops up the image below, click on WebSocket.

Click on Headers, then add a new field, Authorization, with value as “Bearer <access_token>”, then click on connect

If setup is correct, you should receive a response indicating the connection is established or an error message if there’s an issue with the authentication.

Conclusion

This article delves into WebSocket authentication in Django, leveraging Django Channels and Django Rest Framework with JWT. It encompasses project setup, model creation, Django Channels configuration, custom middleware for authentication.

For the complete code implementation, you can access the GitHub repository here