snapshot of https://github.com/zulip/zulip.git @ 48c69911c7dd4e03b20552eb3bd77ad8151c1357 (2026-09-15)

zerver/lib/redis_utils.py

@@ -0,0 +1,96 @@
+import os
+import re
+import secrets
+from collections.abc import Mapping
+from typing import Any
+
+import orjson
+import redis
+from django.conf import settings
+
+# Redis accepts keys up to 512MB in size, but there's no reason for us to use such size,
+# so we want to stay limited to 1024 characters.
+MAX_KEY_LENGTH = 1024
+
+
+class ZulipRedisError(Exception):
+    pass
+
+
+class ZulipRedisKeyTooLongError(ZulipRedisError):
+    pass
+
+
+class ZulipRedisKeyOfWrongFormatError(ZulipRedisError):
+    pass
+
+
+def get_redis_client() -> "redis.StrictRedis[bytes]":
+    return redis.StrictRedis(
+        host=settings.REDIS_HOST,
+        port=settings.REDIS_PORT,
+        password=settings.REDIS_PASSWORD,
+        db=0,
+        decode_responses=False,
+    )
+
+
+def put_dict_in_redis(
+    redis_client: "redis.StrictRedis[bytes]",
+    key_format: str,
+    data_to_store: Mapping[str, Any],
+    expiration_seconds: int,
+    token_length: int = 64,
+    token: str | None = None,
+) -> str:
+    key_length = len(key_format) - len("{token}") + token_length
+    if key_length > MAX_KEY_LENGTH:
+        raise ZulipRedisKeyTooLongError(
+            f"Requested key too long in put_dict_in_redis. Key format: {key_format}, token length: {token_length}",
+        )
+    if token is None:
+        token = secrets.token_hex(token_length // 2)
+    key = key_format.format(token=token)
+    with redis_client.pipeline() as pipeline:
+        pipeline.set(key, orjson.dumps(data_to_store))
+        pipeline.expire(key, expiration_seconds)
+        pipeline.execute()
+
+    return key
+
+
+def get_dict_from_redis(
+    redis_client: "redis.StrictRedis[bytes]",
+    key_format: str,
+    key: str,
+) -> dict[str, Any] | None:
+    # This function requires inputting the intended key_format to validate
+    # that the key fits it, as an additionally security measure. This protects
+    # against bugs where a caller requests a key based on user input and doesn't
+    # validate it - which could potentially allow users to poke around arbitrary Redis keys.
+    if len(key) > MAX_KEY_LENGTH:
+        raise ZulipRedisKeyTooLongError(
+            f"Requested key too long in get_dict_from_redis: {key}",
+        )
+    validate_key_fits_format(key, key_format)
+
+    data = redis_client.get(key)
+    if data is None:
+        return None
+    return orjson.loads(data)
+
+
+def validate_key_fits_format(key: str, key_format: str) -> None:
+    assert "{token}" in key_format
+    regex = key_format.format(token=r"[a-zA-Z0-9]+")
+
+    if not re.fullmatch(regex, key):
+        raise ZulipRedisKeyOfWrongFormatError(f"{key} does not match format {key_format}")
+
+
+REDIS_KEY_PREFIX = ""
+
+
+def bounce_redis_key_prefix_for_testing(test_name: str) -> None:
+    global REDIS_KEY_PREFIX
+    REDIS_KEY_PREFIX = test_name + ":" + str(os.getpid()) + ":"