feat: Add personal API key hash mode sha256 and an upgrade path (#20327)

posthog/api/test/test_feature_flag.py

@@ -2105,17 +2105,19 @@ def test_local_evaluation_for_invalid_cohorts(self, mock_capture):
 
         self.client.logout()
 
-        with self.assertNumQueries(10):
-            # E   1. SELECT "posthog_personalapikey"."id"
-            # E   2. UPDATE "posthog_personalapikey" SET "last_used_at" = '2024-01-31T13:01:37.394080+00:00'
-            # E   3. SELECT "posthog_team"."id", "posthog_team"."uuid"
-            # E   4. SELECT "posthog_organizationmembership"."id", "posthog_organizationmembership"."organization_id"
-            # E   5. SELECT "posthog_cohort"."id"  -- all cohorts
-            # E   6. SELECT "posthog_featureflag"."id", "posthog_featureflag"."key", -- all flags
-            # E   7. SELECT "posthog_cohort". id = 99999
-            # E   8. SELECT "posthog_cohort". id = deleted cohort
-            # E   9. SELECT "posthog_cohort". id = cohort from other team
-            # E   10. SELECT "posthog_grouptypemapping"."id", -- group type mapping
+        with self.assertNumQueries(12):
+            # E   1. SAVEPOINT
+            # E   2. SELECT "posthog_personalapikey"."id"
+            # E   3. RELEASE SAVEPOINT
+            # E   4. UPDATE "posthog_personalapikey" SET "last_used_at" = '2024-01-31T13:01:37.394080+00:00'
+            # E   5. SELECT "posthog_team"."id", "posthog_team"."uuid"
+            # E   6. SELECT "posthog_organizationmembership"."id", "posthog_organizationmembership"."organization_id"
+            # E   7. SELECT "posthog_cohort"."id"  -- all cohorts
+            # E   8. SELECT "posthog_featureflag"."id", "posthog_featureflag"."key", -- all flags
+            # E   9. SELECT "posthog_cohort". id = 99999
+            # E  10. SELECT "posthog_cohort". id = deleted cohort
+            # E  11. SELECT "posthog_cohort". id = cohort from other team
+            # E  12. SELECT "posthog_grouptypemapping"."id", -- group type mapping
 
             response = self.client.get(
                 f"/api/feature_flag/local_evaluation?token={self.team.api_token}&send_cohorts",

posthog/api/test/test_personal_api_keys.py

@@ -130,7 +130,7 @@ def setUp(self):
         self.key = PersonalAPIKey.objects.create(label="Test", user=self.user, secure_value=hash_key_value(self.value))
         self.value_390000 = generate_random_token_personal()
         self.key_390000 = PersonalAPIKey.objects.create(
-            label="Test", user=self.user, secure_value=hash_key_value(self.value_390000, iterations=390000)
+            label="Test", user=self.user, secure_value=hash_key_value(self.value_390000, "pbkdf2", iterations=390000)
         )
         self.value_hardcoded = "phx_0a1b2c3d4e5f6g7h8i9j0k1l2m3n4o5p"
         self.key_hardcoded = PersonalAPIKey.objects.create(
@@ -154,19 +154,34 @@ def test_no_key(self):
         )
 
     def test_header_resilient(self):
+        key_before = PersonalAPIKey.objects.get(id=self.key.id).secure_value
+        self.assertTrue(key_before.startswith("sha256$"))
+
         response = self.client.get(
             f"/api/projects/{self.team.id}/dashboards/",
             HTTP_AUTHORIZATION=f"Bearer  {self.value}  ",
         )
         self.assertEqual(response.status_code, 200)
 
+        # Retrieve key from db to check that no update was made
+        key_after = PersonalAPIKey.objects.get(id=self.key.id).secure_value
+        self.assertEqual(key_after, key_before)
+
     def test_header_alternative_iteration_count(self):
+        key_before = PersonalAPIKey.objects.get(id=self.key_390000.id).secure_value
+        self.assertTrue(key_before.startswith("pbkdf2_sha256$390000$"))
+
         response = self.client.get(
             f"/api/projects/{self.team.id}/dashboards/",
             HTTP_AUTHORIZATION=f"Bearer {self.value_390000}",
         )
         self.assertEqual(response.status_code, 200)
 
+        # Retrieve key from db to check if hash was updated to latest mode
+        key_after = PersonalAPIKey.objects.get(id=self.key_390000.id).secure_value
+        self.assertEqual(key_after, hash_key_value(self.value_390000))
+        self.assertNotEqual(key_after, key_before)
+
     def test_header_hardcoded(self):
         response = self.client.get(
             f"/api/projects/{self.team.id}/dashboards/",

posthog/auth.py

@@ -6,6 +6,7 @@
 
 import jwt
 from django.apps import apps
+from django.db import transaction
 from django.http import HttpRequest, JsonResponse
 from django.utils import timezone
 from rest_framework import authentication
@@ -16,7 +17,7 @@
 from posthog.jwt import PosthogJwtAudience, decode_jwt
 from posthog.models.personal_api_key import (
     hash_key_value,
-    PERSONAL_API_KEY_ITERATIONS_TO_TRY,
+    PERSONAL_API_KEY_MODES_TO_TRY,
 )
 from posthog.models.sharing_configuration import SharingConfiguration
 from posthog.models.user import User
@@ -68,27 +69,37 @@ def find_key(
         return key_with_source[0] if key_with_source is not None else None
 
     @classmethod
+    @transaction.atomic
     def validate_key(cls, personal_api_key_with_source):
         from posthog.models import PersonalAPIKey
 
         personal_api_key, source = personal_api_key_with_source
         personal_api_key_object = None
+        mode_used = None
 
-        for iterations in PERSONAL_API_KEY_ITERATIONS_TO_TRY:
-            secure_value = hash_key_value(personal_api_key, iterations=iterations)
+        for mode, iterations in PERSONAL_API_KEY_MODES_TO_TRY:
+            secure_value = hash_key_value(personal_api_key, mode=mode, iterations=iterations)
             try:
                 personal_api_key_object = (
                     PersonalAPIKey.objects.select_related("user")
                     .filter(user__is_active=True)
                     .get(secure_value=secure_value)
                 )
+                mode_used = mode
                 break
             except PersonalAPIKey.DoesNotExist:
                 pass
 
         if not personal_api_key_object:
             raise AuthenticationFailed(detail=f"Personal API key found in request {source} is invalid.")
 
+        # Upgrade the key if it's not in the latest mode. We can do this since above we've already checked
+        # that the key is valid in some mode, and we do check for all modes one by one.
+        if mode_used != "sha256":
+            key_to_update = PersonalAPIKey.objects.select_for_update().get(id=personal_api_key_object.id)
+            key_to_update.secure_value = hash_key_value(personal_api_key)
+            key_to_update.save(update_fields=["secure_value"])
+
         return personal_api_key_object
 
     @classmethod

posthog/models/personal_api_key.py

@@ -1,27 +1,37 @@
+import hashlib
+from typing import Optional, Literal
+
 from django.contrib.auth.hashers import PBKDF2PasswordHasher
 from django.db import models
 from django.utils import timezone
 
 from .utils import generate_random_token
 
-# Fixed iteration count for PBKDF2PasswordHasher hasher.
-# This is the iteration count used by PostHog since the beginning of time.
-# Changing this would break all existing personal API keys.
-PERSONAL_API_KEY_ITERATIONS = 260000
-
-PERSONAL_API_KEY_ITERATIONS_TO_TRY = (
-    PERSONAL_API_KEY_ITERATIONS,
-    390000,  # This is the iteration count used briefly on some API keys.
+ModeType = Literal["sha256", "pbkdf2"]
+PERSONAL_API_KEY_MODES_TO_TRY: tuple[tuple[ModeType, Optional[int]], ...] = (
+    ("sha256", None),  # Moved to simple hashing in 2024-02
+    ("pbkdf2", 260000),  # This is the iteration count used by PostHog since the beginning of time.
+    ("pbkdf2", 390000),  # This is the iteration count used briefly on some API keys.
 )
 
-# A constant salt is not nearly as good as user-specific, but we must be able to look up a personal API key
-# by itself. Some salt is slightly better than none though.
-PERSONAL_API_KEY_SALT = "posthog_personal_api_key"
+LEGACY_PERSONAL_API_KEY_SALT = "posthog_personal_api_key"
+
+
+def hash_key_value(value: str, mode: ModeType = "sha256", iterations: Optional[int] = None) -> str:
+    if mode == "pbkdf2":
+        if not iterations:
+            raise ValueError("Iterations must be provided when using legacy PBKDF2 mode")
+
+        hasher = PBKDF2PasswordHasher()
+        return hasher.encode(value, LEGACY_PERSONAL_API_KEY_SALT, iterations=iterations)
 
+    if iterations:
+        raise ValueError("Iterations must not be provided when using simple hashing mode")
 
-def hash_key_value(value: str, iterations: int = PERSONAL_API_KEY_ITERATIONS) -> str:
-    hasher = PBKDF2PasswordHasher()
-    return hasher.encode(value, PERSONAL_API_KEY_SALT, iterations=iterations)
+    # Inspiration on why no salt:
+    # https://github.com/jazzband/django-rest-knox/issues/188
+    value = hashlib.sha256(value.encode()).hexdigest()
+    return f"sha256${value}"  # Following format from Django's PBKDF2PasswordHasher
 
 
 class PersonalAPIKey(models.Model):