Build KarlDNS self-hosted DynDNS service

router_dyndns/store/rate_limit.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import sqlite3
+import time
+from datetime import UTC, datetime, timedelta
+
+from .types import CleanupRunRow
+
+RATE_LIMIT_FLUSH_SECONDS = 10.0
+
+
+class RateLimitCleanupMixin:
+    def list_cleanup_runs(self, limit: int = 20, offset: int = 0) -> list[CleanupRunRow]:
+        with self.connect() as conn:
+            conn.row_factory = sqlite3.Row
+            rows = conn.execute(
+                """
+                SELECT created_at, domain_challenges_deleted, unused_accounts_deleted
+                FROM cleanup_runs
+                ORDER BY created_at DESC
+                LIMIT ? OFFSET ?
+                """,
+                (limit, offset),
+            ).fetchall()
+        return [CleanupRunRow(**dict(row)) for row in rows]
+
+    def count_cleanup_runs(self) -> int:
+        with self.connect() as conn:
+            row = conn.execute("SELECT COUNT(*) FROM cleanup_runs").fetchone()
+        return int(row[0])
+
+    def allow_rate_limit(self, key: str, limit: int) -> bool:
+        if limit <= 0:
+            return False
+
+        now = time.time()
+        window = int(now) // 60
+        refill_rate = limit / 60.0
+        with self._rate_limit_lock:
+            bucket = self._rate_limit_buckets.get(key)
+            if bucket is None:
+                bucket = {"tokens": float(limit), "updated_at": now, "window": window}
+            else:
+                elapsed_seconds = max(0.0, now - float(bucket["updated_at"]))
+                bucket["tokens"] = min(float(limit), float(bucket["tokens"]) + elapsed_seconds * refill_rate)
+                bucket["updated_at"] = now
+                bucket["window"] = window
+
+            if float(bucket["tokens"]) < 1.0:
+                self._rate_limit_buckets[key] = bucket
+                return False
+
+            bucket["tokens"] = float(bucket["tokens"]) - 1.0
+            self._rate_limit_buckets[key] = bucket
+            dirty_key = (key, window)
+            self._rate_limit_dirty_counts[dirty_key] = self._rate_limit_dirty_counts.get(dirty_key, 0) + 1
+            self._prune_memory_rate_limits(now)
+
+        self._flush_rate_limits_if_due()
+        return True
+
+    def flush_rate_limits(self) -> None:
+        with self._rate_limit_lock:
+            dirty_counts = self._rate_limit_dirty_counts
+            self._rate_limit_dirty_counts = {}
+            self._last_rate_limit_flush = time.monotonic()
+
+        if not dirty_counts:
+            return
+
+        try:
+            with self.connect() as conn:
+                conn.executemany(
+                    """
+                    INSERT INTO rate_limits(key, window, count)
+                    VALUES (?, ?, ?)
+                    ON CONFLICT(key, window) DO UPDATE SET count=count + excluded.count
+                    """,
+                    [(key, window, count) for (key, window), count in dirty_counts.items()],
+                )
+        except Exception:
+            with self._rate_limit_lock:
+                for dirty_key, count in dirty_counts.items():
+                    self._rate_limit_dirty_counts[dirty_key] = self._rate_limit_dirty_counts.get(dirty_key, 0) + count
+            raise
+
+    def _flush_rate_limits_if_due(self) -> None:
+        if time.monotonic() - self._last_rate_limit_flush < RATE_LIMIT_FLUSH_SECONDS:
+            return
+        self.flush_rate_limits()
+
+    def _prune_memory_rate_limits(self, now: float) -> None:
+        stale_before = now - 120
+        stale_keys = [
+            key
+            for key, bucket in self._rate_limit_buckets.items()
+            if float(bucket["updated_at"]) < stale_before and float(bucket["tokens"]) >= 1.0
+        ]
+        for key in stale_keys:
+            del self._rate_limit_buckets[key]
+
+    def prune_rate_limit_storage(self) -> None:
+        current_window = int(time.time()) // 60
+        should_prune = self._should_prune_rate_limits(current_window)
+        if not should_prune:
+            return
+        with self.connect() as conn:
+            conn.execute("DELETE FROM rate_limits WHERE window < ?", (current_window - 2,))
+
+    def _should_prune_rate_limits(self, window: int) -> bool:
+        if self._last_rate_limit_prune_window == window:
+            return False
+        with self._rate_limit_prune_lock:
+            if self._last_rate_limit_prune_window == window:
+                return False
+            self._last_rate_limit_prune_window = window
+            return True
+
+    def cleanup(
+        self,
+        challenge_hours: int,
+        unused_account_hours: int,
+        event_history_hours: int = 24 * 180,
+    ) -> dict[str, int]:
+        self.flush_rate_limits()
+        challenge_cutoff = (datetime.now(UTC) - timedelta(hours=challenge_hours)).isoformat()
+        unused_account_cutoff = (datetime.now(UTC) - timedelta(hours=unused_account_hours)).isoformat()
+        event_history_cutoff = (datetime.now(UTC) - timedelta(hours=event_history_hours)).isoformat()
+        with self.connect() as conn:
+            challenge_cursor = conn.execute(
+                "DELETE FROM domain_challenges WHERE verified_at IS NULL AND created_at < ?",
+                (challenge_cutoff,),
+            )
+            account_cursor = conn.execute(
+                """
+                DELETE FROM accounts
+                WHERE created_at < ?
+                  AND NOT EXISTS (
+                    SELECT 1 FROM records WHERE records.hostname = accounts.hostname
+                  )
+                """,
+                (unused_account_cutoff,),
+            )
+            rollup_cursor = conn.execute(
+                "DELETE FROM update_event_rollups WHERE last_seen_at < ?",
+                (unused_account_cutoff,),
+            )
+            event_cursor = conn.execute(
+                "DELETE FROM update_events WHERE created_at < ?",
+                (event_history_cutoff,),
+            )
+            current_window = int(time.time()) // 60
+            conn.execute("DELETE FROM rate_limits WHERE window < ?", (current_window - 60,))
+            conn.execute(
+                """
+                INSERT INTO cleanup_runs(created_at, domain_challenges_deleted, unused_accounts_deleted)
+                VALUES (?, ?, ?)
+                """,
+                (datetime.now(UTC).isoformat(), challenge_cursor.rowcount, account_cursor.rowcount),
+            )
+        return {
+            "domain_challenges": challenge_cursor.rowcount,
+            "unused_accounts": account_cursor.rowcount,
+            "event_rollups": rollup_cursor.rowcount,
+            "update_events": event_cursor.rowcount,
+        }