Add email based rate limiting to email login API endpoint

Server:
 - Rate limit based on unverified email before creating user
 - Check email address for deliverability before creating user
 - Track rate limit for unverified email in new non-user keyed table

Web app:
 - Show error in login popup to user on failure/throttling
 - Simplify login popup logic by moving magic link handling logic
   into EmailSigninContext instead of passing require props via parent

Author: debanjum

src/interface/web/app/components/loginPrompt/loginPrompt.tsx

@@ -52,10 +52,6 @@ export default function LoginPrompt(props: LoginPromptProps) {
 
     const [useEmailSignIn, setUseEmailSignIn] = useState(false);
 
-    const [email, setEmail] = useState("");
-    const [checkEmail, setCheckEmail] = useState(false);
-    const [recheckEmail, setRecheckEmail] = useState(false);
-
     useEffect(() => {
         const google = (window as any).google;
 
@@ -118,49 +114,13 @@ export default function LoginPrompt(props: LoginPromptProps) {
         });
     };
 
-    function handleMagicLinkSignIn() {
-        fetch("/auth/magic", {
-            method: "POST",
-            headers: {
-                "Content-Type": "application/json",
-            },
-            body: JSON.stringify({ email: email }),
-        })
-            .then((res) => {
-                if (res.ok) {
-                    setCheckEmail(true);
-                    if (checkEmail) {
-                        setRecheckEmail(true);
-                    }
-                    return res.json();
-                } else {
-                    throw new Error("Failed to send magic link");
-                }
-            })
-            .then((data) => {
-                console.log(data);
-            })
-            .catch((err) => {
-                console.error(err);
-            });
-    }
-
     if (props.isMobileWidth) {
         return (
             <Drawer open={true} onOpenChange={props.onOpenChange}>
                 <DrawerContent className={`flex flex-col gap-4 w-full mb-4`}>
                     <div>
                         {useEmailSignIn ? (
-                            <EmailSignInContext
-                                email={email}
-                                setEmail={setEmail}
-                                checkEmail={checkEmail}
-                                setCheckEmail={setCheckEmail}
-                                setUseEmailSignIn={setUseEmailSignIn}
-                                recheckEmail={recheckEmail}
-                                setRecheckEmail={setRecheckEmail}
-                                handleMagicLinkSignIn={handleMagicLinkSignIn}
-                            />
+                            <EmailSignInContext setUseEmailSignIn={setUseEmailSignIn} />
                         ) : (
                             <MainSignInContext
                                 handleGoogleScriptLoad={handleGoogleScriptLoad}
@@ -187,16 +147,7 @@ export default function LoginPrompt(props: LoginPromptProps) {
                 </VisuallyHidden.Root>
                 <div>
                     {useEmailSignIn ? (
-                        <EmailSignInContext
-                            email={email}
-                            setEmail={setEmail}
-                            checkEmail={checkEmail}
-                            setCheckEmail={setCheckEmail}
-                            setUseEmailSignIn={setUseEmailSignIn}
-                            recheckEmail={recheckEmail}
-                            setRecheckEmail={setRecheckEmail}
-                            handleMagicLinkSignIn={handleMagicLinkSignIn}
-                        />
+                        <EmailSignInContext setUseEmailSignIn={setUseEmailSignIn} />
                     ) : (
                         <MainSignInContext
                             handleGoogleScriptLoad={handleGoogleScriptLoad}
@@ -214,26 +165,17 @@ export default function LoginPrompt(props: LoginPromptProps) {
 }
 
 function EmailSignInContext({
-    email,
-    setEmail,
-    checkEmail,
-    setCheckEmail,
     setUseEmailSignIn,
-    recheckEmail,
-    handleMagicLinkSignIn,
 }: {
-    email: string;
-    setEmail: (email: string) => void;
-    checkEmail: boolean;
-    setCheckEmail: (checkEmail: boolean) => void;
     setUseEmailSignIn: (useEmailSignIn: boolean) => void;
-    recheckEmail: boolean;
-    setRecheckEmail: (recheckEmail: boolean) => void;
-    handleMagicLinkSignIn: () => void;
 }) {
     const [otp, setOTP] = useState("");
     const [otpError, setOTPError] = useState("");
     const [numFailures, setNumFailures] = useState(0);
+    const [email, setEmail] = useState("");
+    const [checkEmail, setCheckEmail] = useState(false);
+    const [recheckEmail, setRecheckEmail] = useState(false);
+    const [sendEmailError, setSendEmailError] = useState("");
 
     function checkOTPAndRedirect() {
         const verifyUrl = `/auth/magic?code=${encodeURIComponent(otp)}&email=${encodeURIComponent(email)}`;
@@ -275,6 +217,39 @@ function EmailSignInContext({
             });
     }
 
+    function handleMagicLinkSignIn() {
+        fetch("/auth/magic", {
+            method: "POST",
+            headers: {
+                "Content-Type": "application/json",
+            },
+            body: JSON.stringify({ email: email }),
+        })
+            .then((res) => {
+                if (res.ok) {
+                    setCheckEmail(true);
+                    if (checkEmail) {
+                        setRecheckEmail(true);
+                    }
+                    return res.json();
+                } else if (res.status === 429 || res.status === 404) {
+                    res.json().then((data) => {
+                        setSendEmailError(data.detail);
+                        throw new Error(data.detail);
+                    });
+                } else {
+                    setSendEmailError("Failed to send email. Contact developers for assistance.");
+                    throw new Error("Failed to send magic link via email.");
+                }
+            })
+            .then((data) => {
+                console.log(data);
+            })
+            .catch((err) => {
+                console.error(err);
+            });
+    }
+
     return (
         <div className="flex flex-col gap-4 p-4">
             <Button
@@ -297,7 +272,7 @@ function EmailSignInContext({
                     : "You will receive a sign-in code on the email address you provide below"}
             </div>
             {!checkEmail && (
-                <>
+                <div className="flex items-center justify-center gap-4 text-muted-foreground flex-col">
                     <Input
                         placeholder="Email"
                         className="p-6 w-[300px] mx-auto rounded-lg"
@@ -320,7 +295,8 @@ function EmailSignInContext({
                         <PaperPlaneTilt className="h-6 w-6 mr-2 font-bold" />
                         {checkEmail ? "Check your email" : "Send sign in code"}
                     </Button>
-                </>
+                    {sendEmailError && <div className="text-red-500 text-sm">{sendEmailError}</div>}
+                </div>
             )}
             {checkEmail && (
                 <div className="flex items-center justify-center gap-4 text-muted-foreground flex-col">
@@ -359,9 +335,7 @@ function EmailSignInContext({
                         variant="ghost"
                         className="p-0 text-orange-500"
                         disabled={recheckEmail}
-                        onClick={() => {
-                            handleMagicLinkSignIn();
-                        }}
+                        onClick={handleMagicLinkSignIn}
                     >
                         <ArrowsClockwise className="h-6 w-6 mr-2 text-gray-500" />
                         Resend email

src/khoj/configure.py

@@ -37,6 +37,7 @@
     aget_or_create_user_by_phone_number,
     aget_user_by_phone_number,
     ais_user_subscribed,
+    delete_ratelimit_records,
     delete_user_requests,
     get_all_users,
     get_or_create_search_models,
@@ -428,8 +429,10 @@ def upload_telemetry():
 @schedule.repeat(schedule.every(31).minutes)
 @clean_connections
 def delete_old_user_requests():
-    num_deleted = delete_user_requests()
-    logger.debug(f"🗑️ Deleted {num_deleted[0]} day-old user requests")
+    num_user_ratelimit_requests = delete_user_requests()
+    num_ratelimit_requests = delete_ratelimit_records()
+    if state.verbose > 2:
+        logger.debug(f"🗑️ Deleted {num_user_ratelimit_requests + num_ratelimit_requests} stale rate limit requests")
 
 
 @schedule.repeat(schedule.every(17).minutes)

src/khoj/database/adapters/__init__.py

@@ -27,6 +27,7 @@
 from django.db.models import Prefetch, Q
 from django.db.models.manager import BaseManager
 from django.db.utils import IntegrityError
+from django.utils import timezone as django_timezone
 from django_apscheduler import util
 from django_apscheduler.models import DjangoJob, DjangoJobExecution
 from fastapi import HTTPException
@@ -49,6 +50,7 @@
     NotionConfig,
     ProcessLock,
     PublicConversation,
+    RateLimitRecord,
     ReflectiveQuestion,
     SearchModelConfig,
     ServerChatSettings,
@@ -233,20 +235,21 @@ async def acreate_user_by_phone_number(phone_number: str) -> KhojUser:
     return user
 
 
-async def aget_or_create_user_by_email(input_email: str) -> tuple[KhojUser, bool]:
-    email, is_valid_email = normalize_email(input_email)
+async def aget_or_create_user_by_email(input_email: str, check_deliverability=False) -> tuple[KhojUser, bool]:
+    # Validate deliverability to email address of new user
+    email, is_valid_email = normalize_email(input_email, check_deliverability=check_deliverability)
     is_existing_user = await KhojUser.objects.filter(email=email).aexists()
-    # Validate email address of new users
     if not is_existing_user and not is_valid_email:
         logger.error(f"Account creation failed. Invalid email address: {email}")
         return None, False
 
+    # Get/create user based on email address
     user, is_new = await KhojUser.objects.filter(email=email).aupdate_or_create(
         defaults={"username": email, "email": email}
     )
 
     # Generate a secure 6-digit numeric code
-    user.email_verification_code = f"{secrets.randbelow(1000000):06}"
+    user.email_verification_code = f"{secrets.randbelow(int(1e6)):06}"
     user.email_verification_code_expiry = datetime.now(tz=timezone.utc) + timedelta(minutes=5)
     await user.asave()
 
@@ -516,8 +519,18 @@ def get_user_notion_config(user: KhojUser):
     return config
 
 
-def delete_user_requests(window: timedelta = timedelta(days=1)):
-    return UserRequests.objects.filter(created_at__lte=datetime.now(tz=timezone.utc) - window).delete()
+def delete_user_requests(max_age: timedelta = timedelta(days=1)):
+    """Deletes UserRequests entries older than the specified max_age."""
+    cutoff = django_timezone.now() - max_age
+    deleted_count, _ = UserRequests.objects.filter(created_at__lte=cutoff).delete()
+    return deleted_count
+
+
+def delete_ratelimit_records(max_age: timedelta = timedelta(days=1)):
+    """Deletes RateLimitRecord entries older than the specified max_age."""
+    cutoff = django_timezone.now() - max_age
+    deleted_count, _ = RateLimitRecord.objects.filter(created_at__lt=cutoff).delete()
+    return deleted_count
 
 
 @arequire_valid_user

src/khoj/database/admin.py

@@ -25,6 +25,7 @@
     KhojUser,
     NotionConfig,
     ProcessLock,
+    RateLimitRecord,
     ReflectiveQuestion,
     SearchModelConfig,
     ServerChatSettings,
@@ -179,6 +180,7 @@ def get_email_login_url(self, request, queryset):
 admin.site.register(UserVoiceModelConfig, unfold_admin.ModelAdmin)
 admin.site.register(VoiceModelOption, unfold_admin.ModelAdmin)
 admin.site.register(UserRequests, unfold_admin.ModelAdmin)
+admin.site.register(RateLimitRecord, unfold_admin.ModelAdmin)
 
 
 @admin.register(Agent)

src/khoj/database/migrations/0088_ratelimitrecord.py

@@ -0,0 +1,28 @@
+# Generated by Django 5.0.13 on 2025-04-07 07:10
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("database", "0087_alter_aimodelapi_api_key"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="RateLimitRecord",
+            fields=[
+                ("id", models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")),
+                ("created_at", models.DateTimeField(auto_now_add=True)),
+                ("updated_at", models.DateTimeField(auto_now=True)),
+                ("identifier", models.CharField(db_index=True, max_length=255)),
+                ("slug", models.CharField(db_index=True, max_length=255)),
+            ],
+            options={
+                "ordering": ["-created_at"],
+                "indexes": [
+                    models.Index(fields=["identifier", "slug", "created_at"], name="database_ra_identif_031adf_idx")
+                ],
+            },
+        ),
+    ]

src/khoj/database/models/__init__.py

@@ -730,10 +730,28 @@ class Meta:
 
 
 class UserRequests(DbBaseModel):
+    """Stores user requests to the server for rate limiting."""
+
     user = models.ForeignKey(KhojUser, on_delete=models.CASCADE)
     slug = models.CharField(max_length=200)
 
 
+class RateLimitRecord(DbBaseModel):
+    """Stores individual request timestamps for rate limiting."""
+
+    identifier = models.CharField(max_length=255, db_index=True)  # IP address or email
+    slug = models.CharField(max_length=255, db_index=True)  # Differentiates limit types
+
+    class Meta:
+        indexes = [
+            models.Index(fields=["identifier", "slug", "created_at"]),
+        ]
+        ordering = ["-created_at"]
+
+    def __str__(self):
+        return f"{self.slug} - {self.identifier} at {self.created_at}"
+
+
 class DataStore(DbBaseModel):
     key = models.CharField(max_length=200, unique=True)
     value = models.JSONField(default=dict)