From bdc883dcb38e77990a73ee3ad3d36c8925322d6c Mon Sep 17 00:00:00 2001
From: Joel Whitaker <joel@wolfbyte.co.uk>
Date: Thu, 29 Oct 2020 13:49:37 +0000
Subject: [PATCH] Add the ability to return the refresh token in a cookie

---
 dj_rest_auth/app_settings.py |  1 +
 dj_rest_auth/views.py        | 17 ++++++++++++++++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/dj_rest_auth/app_settings.py b/dj_rest_auth/app_settings.py
index fe6e6a2..0b28085 100644
--- a/dj_rest_auth/app_settings.py
+++ b/dj_rest_auth/app_settings.py
@@ -38,3 +38,4 @@ PasswordChangeSerializer = import_callable(
 )
 
 JWT_AUTH_COOKIE = getattr(settings, 'JWT_AUTH_COOKIE', None)
+JWT_AUTH_REFRESH_COOKIE = getattr(settings, 'JWT_AUTH_REFRESH_COOKIE', None)
diff --git a/dj_rest_auth/views.py b/dj_rest_auth/views.py
index 28b472d..970d7ec 100644
--- a/dj_rest_auth/views.py
+++ b/dj_rest_auth/views.py
@@ -86,12 +86,13 @@ class LoginView(GenericAPIView):
         response = Response(serializer.data, status=status.HTTP_200_OK)
         if getattr(settings, 'REST_USE_JWT', False):
             cookie_name = getattr(settings, 'JWT_AUTH_COOKIE', None)
+            refresh_cookie_name = getattr(settings, 'JWT_AUTH_REFRESH_COOKIE', None)
             cookie_secure = getattr(settings, 'JWT_AUTH_SECURE', False)
             cookie_httponly = getattr(settings, 'JWT_AUTH_HTTPONLY', True)
             cookie_samesite = getattr(settings, 'JWT_AUTH_SAMESITE', 'Lax')
             from rest_framework_simplejwt.settings import api_settings as jwt_settings
+            from datetime import datetime
             if cookie_name:
-                from datetime import datetime
                 expiration = (datetime.utcnow() + jwt_settings.ACCESS_TOKEN_LIFETIME)
                 response.set_cookie(
                     cookie_name,
@@ -101,6 +102,16 @@ class LoginView(GenericAPIView):
                     httponly=cookie_httponly,
                     samesite=cookie_samesite
                 )
+            if refresh_cookie_name:
+                expiration = (datetime.utcnow() + jwt_settings.REFRESH_TOKEN_LIFETIME)
+                response.set_cookie(
+                    refresh_cookie_name,
+                    self.refresh_token,
+                    expires=expiration,
+                    secure=cookie_secure,
+                    httponly=cookie_httponly,
+                    samesite=cookie_samesite
+                )
         return response
 
     def post(self, request, *args, **kwargs):
@@ -156,6 +167,10 @@ class LogoutView(APIView):
             if cookie_name:
                 response.delete_cookie(cookie_name)
 
+            refresh_cookie_name = getattr(settings, 'JWT_AUTH_REFRESH_COOKIE', None)
+            if refresh_cookie_name:
+                response.delete_cookie(refresh_cookie_name)
+
             elif 'rest_framework_simplejwt.token_blacklist' in settings.INSTALLED_APPS:
                 # add refresh token to blacklist
                 try: