Prevent small risk of Token overwrite (#9754)

mahdirahimi1999 · browniebroke · web-flow · commit c0166d95bb64 · 2025-08-10T16:52:32.000+06:00
* Fix #9250: Prevent token overwrite and improve security - Fix key collision issue that could overwrite existing tokens - Use force_insert=True only for new token instances - Replace os.urandom with secrets.token_hex for better security - Add comprehensive test suite to verify fix and backward compatibility - Ensure existing tokens can still be updated without breaking changes * Fix code style: remove trailing whitespace and unused imports * Fix #9250: Prevent token overwrite with minimal changes - Add force_insert=True to Token.save() for new objects to prevent overwriting existing tokens - Revert generate_key method to original implementation (os.urandom + binascii) - Update tests to work with original setUp() approach - Remove verbose comments and unrelated changes per reviewer feedback * Fix flake8 violations: remove extra blank lines and trailing whitespace * Update tests/test_authtoken.py Co-authored-by: Bruno Alla <browniebroke@users.noreply.github.com> * Update tests/test_authtoken.py Co-authored-by: Bruno Alla <browniebroke@users.noreply.github.com> * Update tests/test_authtoken.py Co-authored-by: Bruno Alla <browniebroke@users.noreply.github.com> * Fix token key regeneration behavior and add test * Update tests/test_authtoken.py Co-authored-by: Bruno Alla <browniebroke@users.noreply.github.com> --------- Co-authored-by: Bruno Alla <browniebroke@users.noreply.github.com>
diff --git a/rest_framework/authtoken/models.py b/rest_framework/authtoken/models.py
@@ -27,8 +27,17 @@ class Meta:
         verbose_name_plural = _("Tokens")
 
     def save(self, *args, **kwargs):
+        """
+        Save the token instance.
+
+        If no key is provided, generates a cryptographically secure key.
+        For new tokens, ensures they are inserted as new (not updated).
+        """
         if not self.key:
             self.key = self.generate_key()
+            # For new objects, force INSERT to prevent overwriting existing tokens
+            if self._state.adding:
+                kwargs['force_insert'] = True
         return super().save(*args, **kwargs)
 
     @classmethod
diff --git a/tests/test_authtoken.py b/tests/test_authtoken.py
@@ -5,6 +5,7 @@
 from django.contrib.admin import site
 from django.contrib.auth.models import User
 from django.core.management import CommandError, call_command
+from django.db import IntegrityError
 from django.test import TestCase, modify_settings
 
 from rest_framework.authtoken.admin import TokenAdmin
@@ -48,6 +49,45 @@ def test_whitespace_in_password(self):
         self.user.save()
         assert AuthTokenSerializer(data=data).is_valid()
 
+    def test_token_creation_collision_raises_integrity_error(self):
+        user2 = User.objects.create_user('user2', 'user2@example.com', 'p')
+        existing_token = Token.objects.create(user=user2)
+
+        # Try to create another token with the same key
+        with self.assertRaises(IntegrityError):
+            Token.objects.create(key=existing_token.key, user=self.user)
+
+    def test_key_generated_on_save_when_cleared(self):
+        # Create a new user for this test to avoid conflicts with setUp token
+        user2 = User.objects.create_user('test_user2', 'test2@example.com', 'password')
+
+        # Create a token without a key - it should generate one automatically
+        token = Token(user=user2)
+        token.key = ""  # Explicitly clear the key
+        token.save()
+
+        # Verify the key was generated
+        self.assertEqual(len(token.key), 40)
+        self.assertEqual(token.user, user2)
+
+    def test_clearing_key_on_existing_token_raises_integrity_error(self):
+        """Test that clearing the key on an existing token raises IntegrityError."""
+        user = User.objects.create_user('test_user3', 'test3@example.com', 'password')
+        token = Token.objects.create(user=user)
+        token.key = ""
+
+        # This should raise IntegrityError because:
+        # 1. We're trying to update a record with an empty primary key
+        # 2. The OneToOneField constraint would be violated
+        with self.assertRaises(Exception):  # Could be IntegrityError or DatabaseError
+            token.save()
+
+    def test_saving_existing_token_without_changes_does_not_alter_key(self):
+        original_key = self.token.key
+
+        self.token.save()
+        self.assertEqual(self.token.key, original_key)
+
 
 class AuthTokenCommandTests(TestCase):
 
