From ab921ccb3115f3914db297bd565c6b4b546a4700 Mon Sep 17 00:00:00 2001
From: Oliver <oliver.henry.walters@gmail.com>
Date: Mon, 22 Jan 2024 23:07:35 +1100
Subject: [PATCH] [PUI] Logout Fixes (#6318)

* Refactor method to extract token from request

* Reimplement error-report API endpoint

- Removed in previous commit - b8b3dfc90e0bca42df4f936303f3f314b2111205
- Adds unit tests to ensure it doesn't happen again

* Adds custom logout view for API

- Ensure correct token gets deleted
- Our new custom token setup is incompatible with default dj-rest-auth
---
 InvenTree/InvenTree/middleware.py   | 38 ++++++++++++++---------------
 InvenTree/InvenTree/urls.py         |  1 +
 InvenTree/users/api.py              | 24 ++++++++++++++++++
 src/frontend/src/functions/auth.tsx | 11 +++++----
 4 files changed, 49 insertions(+), 25 deletions(-)

diff --git a/InvenTree/InvenTree/middleware.py b/InvenTree/InvenTree/middleware.py
index 1691b2225b..86b254ac46 100644
--- a/InvenTree/InvenTree/middleware.py
+++ b/InvenTree/InvenTree/middleware.py
@@ -18,6 +18,23 @@ from users.models import ApiToken
 logger = logging.getLogger('inventree')
 
 
+def get_token_from_request(request):
+    """Extract token information from a request object."""
+    auth_keys = ['Authorization', 'authorization']
+
+    token = None
+
+    for k in auth_keys:
+        if auth_header := request.headers.get(k, None):
+            auth_header = auth_header.strip().lower().split()
+
+            if len(auth_header) > 1 and auth_header[0].startswith('token'):
+                token = auth_header[1]
+                break
+
+    return token
+
+
 class AuthRequiredMiddleware(object):
     """Check for user to be authenticated."""
 
@@ -25,28 +42,9 @@ class AuthRequiredMiddleware(object):
         """Save response object."""
         self.get_response = get_response
 
-    def get_auth_headers(self, request):
-        """Extract authorization headers from request."""
-        keys = ['Authorization', 'authorization']
-
-        for k in keys:
-            if k in request.headers.keys():
-                return request.headers[k]
-
-        return None
-
     def check_token(self, request) -> bool:
         """Check if the user is authenticated via token."""
-        auth = self.get_auth_headers(request)
-
-        if not auth:
-            return False
-
-        auth = auth.strip().lower().split()
-
-        if len(auth) > 1 and auth[0].startswith('token'):
-            token = auth[1]
-
+        if token := get_token_from_request(request):
             # Does the provided token match a valid user?
             try:
                 token = ApiToken.objects.get(key=token)
diff --git a/InvenTree/InvenTree/urls.py b/InvenTree/InvenTree/urls.py
index 547764cf5d..038b75ace2 100644
--- a/InvenTree/InvenTree/urls.py
+++ b/InvenTree/InvenTree/urls.py
@@ -148,6 +148,7 @@ apipatterns = [
                 SocialAccountDisconnectView.as_view(),
                 name='social_account_disconnect',
             ),
+            path('logout/', users.api.Logout.as_view(), name='api-logout'),
             path('', include('dj_rest_auth.urls')),
         ]),
     ),
diff --git a/InvenTree/users/api.py b/InvenTree/users/api.py
index 8343962fdd..241b69815a 100644
--- a/InvenTree/users/api.py
+++ b/InvenTree/users/api.py
@@ -7,6 +7,7 @@ from django.contrib.auth import get_user, login
 from django.contrib.auth.models import Group, User
 from django.urls import include, path, re_path
 
+from dj_rest_auth.views import LogoutView
 from rest_framework import exceptions, permissions
 from rest_framework.response import Response
 from rest_framework.views import APIView
@@ -200,6 +201,29 @@ class GroupList(ListCreateAPI):
     ordering_fields = ['name']
 
 
+class Logout(LogoutView):
+    """API view for logging out via API."""
+
+    def logout(self, request):
+        """Logout the current user.
+
+        Deletes user token associated with request.
+        """
+        from InvenTree.middleware import get_token_from_request
+
+        if request.user:
+            token_key = get_token_from_request(request)
+
+            if token_key:
+                try:
+                    token = ApiToken.objects.get(key=token_key, user=request.user)
+                    token.delete()
+                except ApiToken.DoesNotExist:
+                    pass
+
+        return super().logout(request)
+
+
 class GetAuthToken(APIView):
     """Return authentication token for an authenticated user."""
 
diff --git a/src/frontend/src/functions/auth.tsx b/src/frontend/src/functions/auth.tsx
index 7039ea06d6..d93f172266 100644
--- a/src/frontend/src/functions/auth.tsx
+++ b/src/frontend/src/functions/auth.tsx
@@ -48,16 +48,17 @@ export const doClassicLogin = async (username: string, password: string) => {
  * Logout the user (invalidate auth token)
  */
 export const doClassicLogout = async () => {
+  // Set token in context
+  const { setToken } = useSessionState.getState();
+
+  setToken(undefined);
+
   // Logout from the server session
   await api.post(apiUrl(ApiPaths.user_logout));
 
-  // Set token in context
-  const { setToken } = useSessionState.getState();
-  setToken(undefined);
-
   notifications.show({
     title: t`Logout successful`,
-    message: t`See you soon.`,
+    message: t`You have been logged out`,
     color: 'green',
     icon: <IconCheck size="1rem" />
   });