From 7b2ece248c43fca0b4fab1ae86c756af255985af Mon Sep 17 00:00:00 2001
From: Diyahir <32445955+diyahir@users.noreply.github.com>
Date: Sat, 30 Mar 2024 10:59:29 +0100
Subject: [PATCH] Fixing bug

---
 packages/nextjs/hooks/useWebSocket.tsx | 13 +++--
 packages/nextjs/types/utils.ts         |  4 ++
 packages/server/src/app.ts             | 70 ++++++++++++++++----------
 3 files changed, 56 insertions(+), 31 deletions(-)

diff --git a/packages/nextjs/hooks/useWebSocket.tsx b/packages/nextjs/hooks/useWebSocket.tsx
index b3a49e9..160cf08 100644
--- a/packages/nextjs/hooks/useWebSocket.tsx
+++ b/packages/nextjs/hooks/useWebSocket.tsx
@@ -1,5 +1,5 @@
 import { useCallback, useEffect, useRef, useState } from "react";
-import { InvoiceRequest, InvoiceResponse, ServerStatus } from "~~/types/utils";
+import { ConnectionResponse, InvoiceRequest, InvoiceResponse, ServerStatus } from "~~/types/utils";
 
 export const useWebSocket = (url: string) => {
   const socket = useRef<WebSocket | null>(null);
@@ -41,12 +41,15 @@ export const useWebSocket = (url: string) => {
     socket.current.onerror = event => setError(event);
     socket.current.onmessage = event => {
       try {
-        const responseData: InvoiceResponse = JSON.parse(event.data);
-        if (responseData.status) {
-          setStatus(responseData.status as ServerStatus);
+        const responseData: ConnectionResponse | InvoiceResponse = JSON.parse(event.data);
+        if (responseData && "serverStatus" in responseData) {
+          setStatus(responseData.serverStatus as ServerStatus);
+          return;
+        }
+        if (responseData && "status" in responseData) {
+          setData(responseData);
           return;
         }
-        setData(responseData);
       } catch (err) {
         console.error("Failed to parse message", err);
       }
diff --git a/packages/nextjs/types/utils.ts b/packages/nextjs/types/utils.ts
index 762a29c..d8ed58f 100644
--- a/packages/nextjs/types/utils.ts
+++ b/packages/nextjs/types/utils.ts
@@ -19,6 +19,10 @@ export interface InvoiceResponse {
   message: string;
 }
 
+export interface ConnectionResponse {
+  serverStatus: ServerStatus;
+}
+
 export enum ServerStatus {
   ACTIVE = "ACTIVE",
   INACTIVE = "INACTIVE",
diff --git a/packages/server/src/app.ts b/packages/server/src/app.ts
index 9c8e39c..a3326b4 100644
--- a/packages/server/src/app.ts
+++ b/packages/server/src/app.ts
@@ -13,7 +13,7 @@ dotenv.config();
 // Verify environment variables
 const { PORT, LND_MACAROON, LND_SOCKET, RPC_URL, LSP_PRIVATE_KEY, CHAIN_ID } =
   process.env;
-if ( !RPC_URL || !LSP_PRIVATE_KEY || !CHAIN_ID) {
+if (!RPC_URL || !LSP_PRIVATE_KEY || !CHAIN_ID) {
   console.error("Missing environment variables");
   process.exit(1);
 }
@@ -39,7 +39,9 @@ export type CachedPayment = {
   contractId: string;
   secret: string;
 };
+
 let cachedPayments: CachedPayment[] = [];
+let pendingContracts: string[] = [];
 // ideally this should be stored in a database, but for the sake of simplicity we are using an in-memory cache
 
 console.log(`RPC Provider is running on ${RPC_URL}`);
@@ -49,17 +51,33 @@ console.log(`LSP Address: ${signer.address}`);
 wss.on("connection", (ws: WebSocket) => {
   const serverStatus = process.env.LND_MACAROON ? "ACTIVE" : "MOCK";
   console.log("Client connected");
-  ws.send(JSON.stringify({ status: serverStatus, message: "Connected to server" }));
+  ws.send(
+    JSON.stringify({
+      serverStatus: serverStatus,
+      message: "Connected to server",
+    })
+  );
 
   ws.on("message", async (message: string) => {
     console.log("Received message:", message);
+    const request: InvoiceRequest = JSON.parse(message);
+    if (pendingContracts.includes(request.contractId)) {
+      ws.send(
+        JSON.stringify({
+          status: "error",
+          message: "Contract is already being processed.",
+        })
+      );
+      return;
+    }
+    pendingContracts.push(request.contractId);
     try {
-      const request: InvoiceRequest = JSON.parse(message);
       await processInvoiceRequest(request, ws);
     } catch (error) {
       console.error("Error processing message:", error);
       ws.send(JSON.stringify({ status: "error", message: "Invalid request" }));
     }
+    pendingContracts = pendingContracts.filter((c) => c !== request.contractId);
   });
 
   ws.on("close", () => console.log("Client disconnected"));
@@ -75,22 +93,24 @@ async function processInvoiceRequest(request: InvoiceRequest, ws: WebSocket) {
 
   console.log("Invoice Request Received:", request);
 
-    // Check if LND_MACAROON and LND_SOCKET are empty to simulate mock mode
-    if (!process.env.LND_MACAROON && !process.env.LND_SOCKET) {
-      console.log("Mock Server Mode: Simulating payment success");
-      
-      // Simulate processing delay
-      await new Promise(resolve => setTimeout(resolve, 1000)); // 1 second delay for realism
-  
-      // Directly respond with a simulated success message
-      ws.send(JSON.stringify({
+  // Check if LND_MACAROON and LND_SOCKET are empty to simulate mock mode
+  if (!process.env.LND_MACAROON && !process.env.LND_SOCKET) {
+    console.log("Mock Server Mode: Simulating payment success");
+
+    // Simulate processing delay
+    await new Promise((resolve) => setTimeout(resolve, 1000)); // 1 second delay for realism
+
+    // Directly respond with a simulated success message
+    ws.send(
+      JSON.stringify({
         status: "success",
         message: "Invoice paid successfully in mock mode.",
-      }));
-  
-      // Exit early since we're in mock mode
-      return;
-    }
+      })
+    );
+
+    // Exit early since we're in mock mode
+    return;
+  }
 
   try {
     const options = { gasPrice: ethers.parseUnits("0.001", "gwei") };
@@ -135,10 +155,14 @@ async function processInvoiceRequest(request: InvoiceRequest, ws: WebSocket) {
       max_fee: providerConfig.maxLNFee,
     });
     console.log("Payment Response:", paymentResponse);
-
+    ws.send(
+      JSON.stringify({
+        status: "success",
+        message: "Invoice paid successfully.",
+      })
+    );
     // Critical point, if this withdraw fails, the LSP will lose funds
     // We should cache the paymentResponse.secret and request.contractId and retry the withdrawal if it fails
-
     await htlcContract
       .withdraw(request.contractId, "0x" + paymentResponse.secret, options)
       .then((tx: any) => {
@@ -151,13 +175,7 @@ async function processInvoiceRequest(request: InvoiceRequest, ws: WebSocket) {
           secret: paymentResponse.secret,
         });
       });
-
-    ws.send(
-      JSON.stringify({
-        status: "success",
-        message: "Invoice paid successfully.",
-      })
-    );
+    console.log("Payment processed successfully");
   } catch (error) {
     console.error("Error during invoice processing:", error);
     ws.send(