Fix: Prevent abuse of verify-payment endpoint by adding rate-limiting & idempotency (#205) | Fix: persist selected model in localStorage (#204)

backend/routes/billing.ts

@@ -273,35 +273,37 @@ billingRouter.post("/verify-payment", authMiddleware, async (req, res) => {
 
       // Verify signature
       if (expectedSignature === signature) {
-        // Payment is authentic - update records
-        await prisma.paymentHistory.update({
-          where: { paymentId: paymentRecord.paymentId },
-          data: {
-            status: "SUCCESS",
-            cfPaymentId: razorpay_payment_id
-          }
-        });
-
-        // Update user to premium status and add 12k credits for yearly plan
-        await prisma.user.update({
-          where: { id: userId },
-          data: {
-            isPremium: true,
-            credits: { increment: 12000 } // Add 12k credits for yearly plan
+        await prisma.$transaction(async (tx) => {
+          const freshPayment = await tx.paymentHistory.findUnique({
+            where: { paymentId: paymentRecord.paymentId }
+          });
+          if (freshPayment?.status !== "PENDING") {
+            throw new Error("Payment already processed");
           }
+          await tx.paymentHistory.update({
+            where: { paymentId: paymentRecord.paymentId },
+            data: { status: "SUCCESS", cfPaymentId: razorpay_payment_id }
+          });
+          await tx.user.update({
+            where: { id: userId },
+            data: {
+              isPremium: true,
+              credits: { increment: isYearlyPlan ? 12000 : 1000 }
+            }
+          });
         });
 
         return res.json({
           success: true,
-          message: "Yearly plan payment verified successfully. You now have 12,000 credits!"
+          message: isYearlyPlan
+            ? "Yearly plan payment verified successfully. You now have 12,000 credits!"
+            : "Payment verified successfully"
         });
       } else {
-        // Invalid signature
-        await prisma.paymentHistory.update({
-          where: { paymentId: paymentRecord.paymentId },
-          data: {
-            status: "FAILED"
-          }
+        // Invalid signature: mark FAILED only if it was still PENDING
+        await prisma.paymentHistory.updateMany({
+          where: { paymentId: paymentRecord.paymentId, status: "PENDING" },
+          data: { status: "FAILED" }
         });
 
         return res.status(400).json({
@@ -335,34 +337,42 @@ billingRouter.post("/verify-payment", authMiddleware, async (req, res) => {
       // Verify signature
       if (expectedSignature === signature) {
         // Payment is authentic - update records
-        await prisma.paymentHistory.update({
-          where: { paymentId: paymentRecord.paymentId },
+        // Atomically mark the payment SUCCESS only if it's still PENDING.
+        const updated = await prisma.paymentHistory.updateMany({
+          where: {
+            paymentId: paymentRecord.paymentId,
+            status: "PENDING"
+          },
           data: {
             status: "SUCCESS",
             cfPaymentId: razorpay_payment_id
           }
         });
 
-        // Update user to premium status and add credits
+        if (updated.count === 0) {
+          return res.status(409).json({
+            success: false,
+            error: "Payment already processed or invalid payment state"
+          });
+        }
+
         await prisma.user.update({
           where: { id: userId },
           data: {
             isPremium: true,
-            credits: { increment: 1000 } // Add 1000 credits for monthly subscription
+            credits: { increment: 1000 }
           }
         });
 
         return res.json({
           success: true,
-          message: "Payment verified successfully"
+          message: "Yearly plan payment verified successfully. You now have 12,000 credits!"
         });
       } else {
-        // Invalid signature
-        await prisma.paymentHistory.update({
-          where: { paymentId: paymentRecord.paymentId },
-          data: {
-            status: "FAILED"
-          }
+        // Invalid signature: mark FAILED only if it was still PENDING
+        await prisma.paymentHistory.updateMany({
+          where: { paymentId: paymentRecord.paymentId, status: "PENDING" },
+          data: { status: "FAILED" }
         });
 
         return res.status(400).json({

backend/routes/rzpWebhookRouter.ts

@@ -60,29 +60,33 @@ rzpWebhookRouter.post("/", async (req, res) => {
 
       console.log(`Processing subscription activation for user: ${userId}`);
 
-      // Update user to premium and add 1000 credits
-      await prisma.user.update({
-        where: { id: userId },
-        data: {
-          isPremium: true,
-          credits: {
-            increment: 1000
-          }
-        }
-      });
-
-      // Update payment history status to SUCCESS
-      await prisma.paymentHistory.updateMany({
-        where: { 
+      // Atomically mark payment history SUCCESS only if it's still PENDING
+      const updated = await prisma.paymentHistory.updateMany({
+        where: {
           bankReference: subscriptionId,
           status: "PENDING"
         },
         data: {
           status: "SUCCESS",
-          updatedAt: new Date()
+          cfPaymentId: subscriptionId
         }
       });
 
+      if (updated.count === 0) {
+        // Already processed or nothing to do
+        console.info(`Webhook: subscription ${subscriptionId} already processed`);
+        return res.status(200).json({ message: "Already processed" });
+      }
+
+      // Only increment credits when we actually transitioned a PENDING payment to SUCCESS
+      await prisma.user.update({
+        where: { id: userId },
+        data: {
+          isPremium: true,
+          credits: { increment: 1000 }
+        }
+      });
+
       // Update or create subscription record
       const existingSubscription = await prisma.subscription.findFirst({
         where: { rzpSubscriptionId: subscriptionId }
@@ -113,6 +117,72 @@ rzpWebhookRouter.post("/", async (req, res) => {
       }
 
       console.log(`Successfully activated subscription for user ${userId}`);
+    } else if (event === "payment.captured") {
+      const payment = payload.payment.entity;
+      const { notes, id: paymentId } = payment;
+
+      if (notes.app_name !== "1AI") {
+        return res.status(200).json({ message: "Webhook processed successfully" });
+      }
+
+      // Extract user ID from notes
+      let userId: string | null = null;
+      if (notes && typeof notes === 'object') {
+        userId = notes.customer_id || notes.userId;
+      }
+
+      if (!userId) {
+        console.error("No user ID found in payment notes");
+        return res.status(400).json({ error: "User ID not found in notes" });
+      }
+
+      console.log(`Processing payment capture for user: ${userId}`);
+
+      // Find the corresponding payment record
+      const paymentRecord = await prisma.paymentHistory.findFirst({
+        where: {
+          paymentId: paymentId,
+          status: "PENDING"
+        }
+      });
+
+      if (!paymentRecord) {
+        console.error("No pending payment record found for payment ID:", paymentId);
+        return res.status(404).json({ error: "Payment record not found" });
+      }
+
+      // Replace the monthly-subscription success path that did unconditional updates
+      const updated = await prisma.paymentHistory.updateMany({
+        where: {
+          paymentId: paymentRecord.paymentId,
+          status: "PENDING"
+        },
+        data: {
+          status: "SUCCESS",
+          cfPaymentId: paymentId
+        }
+      });
+
+      if (updated.count === 0) {
+        return res.status(409).json({
+          success: false,
+          error: "Payment already processed or invalid payment state"
+        });
+      }
+
+      // Safe to update user once
+      await prisma.user.update({
+        where: { id: userId },
+        data: {
+          isPremium: true,
+          credits: { increment: 1000 } // monthly credits
+        }
+      });
+
+      return res.json({
+        success: true,
+        message: "Payment verified successfully"
+      });
     }
 
     res.status(200).json({ message: "Webhook processed successfully" });

frontend/components/ui/ui-input.tsx

@@ -54,7 +54,14 @@ interface UIInputProps {
 const UIInput = ({
   conversationId: initialConversationId,
 }: UIInputProps = {}) => {
-  const [model, setModel] = useState<string>(DEFAULT_MODEL_ID);
+  const [model, setModel] = useState<string>(() => {
+    if (typeof window !== "undefined") {
+      const storedModel = localStorage.getItem("selectedModel") || DEFAULT_MODEL_ID;
+      console.log("Loaded model from localStorage:", storedModel);
+      return storedModel;
+    }
+    return DEFAULT_MODEL_ID;
+  });
   const [query, setQuery] = useState<string>("");
   const [messages, setMessages] = useState<Message[]>([]);
   const [showWelcome, setShowWelcome] = useState(true);
@@ -91,6 +98,13 @@ const UIInput = ({
     }
   }, [conversation, initialConversationId]);
 
+  // Persist model selection in localStorage
+  useEffect(() => {
+    if (typeof window !== "undefined") {
+      localStorage.setItem("selectedModel", model);
+    }
+  }, [model]);
+
   useGlobalKeyPress({
     inputRef: textareaRef,
     onKeyPress: (key: string) => {