server: tweak token generation to add checksum

Author: frankie567

server/polar/kit/crypto.py

@@ -1,26 +1,46 @@
 import hashlib
 import hmac
 import secrets
+import string
+import zlib
+
+
+def _crc32_to_base62(number: int) -> str:
+    characters = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+    base = len(characters)
+    encoded = ""
+    while number:
+        number, remainder = divmod(number, base)
+        encoded = characters[remainder] + encoded
+    return encoded.zfill(6)  # Ensure the checksum is 6 characters long
 
 
 def get_token_hash(token: str, *, secret: str) -> str:
     hash = hmac.new(secret.encode("ascii"), token.encode("ascii"), hashlib.sha256)
     return hash.hexdigest()
 
 
-def generate_token(*, prefix: str = "", nbytes: int | None = None) -> str:
-    return f"{prefix}{secrets.token_urlsafe(nbytes)}"
+def generate_token(*, prefix: str = "") -> str:
+    # Generate a high entropy random token
+    token = "".join(
+        secrets.choice(string.ascii_letters + string.digits) for _ in range(37)
+    )
+
+    # Calculate a 32-bit CRC checksum
+    checksum = zlib.crc32(token.encode("utf-8")) & 0xFFFFFFFF
+    checksum_base62 = _crc32_to_base62(checksum)
+
+    # Concatenate the prefix, token, and checksum
+    return f"{prefix}{token}{checksum_base62}"
 
 
-def generate_token_hash_pair(
-    *, secret: str, prefix: str = "", nbytes: int | None = None
-) -> tuple[str, str]:
+def generate_token_hash_pair(*, secret: str, prefix: str = "") -> tuple[str, str]:
     """
     Generate a token suitable for sensitive values
     like magic link tokens.
 
     Returns both the actual value and its HMAC-SHA256 hash.
     Only the latter shall be stored in database.
     """
-    token = generate_token(prefix=prefix, nbytes=nbytes)
+    token = generate_token(prefix=prefix)
     return token, get_token_hash(token, secret=secret)

server/polar/oauth2/authorization_server.py

@@ -70,7 +70,7 @@ def generate_client_registration_info(
         }
 
     def generate_client_id(self) -> str:
-        return generate_token(prefix=CLIENT_ID_PREFIX, nbytes=16)
+        return generate_token(prefix=CLIENT_ID_PREFIX)
 
     def generate_client_secret(self) -> str:
         return generate_token(prefix=CLIENT_SECRET_PREFIX)