From e5a032cfa8081b295e4784c7dcfda666547f164e Mon Sep 17 00:00:00 2001
From: Prathamesh Musale <prathamesh.musale0@gmail.com>
Date: Wed, 24 Jul 2024 18:33:33 +0530
Subject: [PATCH] Use a configurable period limit

---
 environments/local.toml |  1 +
 src/index.ts            | 42 ++++++++++++++++++++++++++++++++---------
 2 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/environments/local.toml b/environments/local.toml
index 2eca458..f9fbfe3 100644
--- a/environments/local.toml
+++ b/environments/local.toml
@@ -8,6 +8,7 @@
 
 [server]
   port = 3000
+  periodInSecs = 60
   transferAmount = 1000000
   periodTransferLimit = 3000000
   dbDir = "db"
diff --git a/src/index.ts b/src/index.ts
index bdcdf78..f5b313d 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -12,7 +12,7 @@ import KeyvSqlite from '@keyv/sqlite';
 
 const CONFIG_PATH = 'environments/local.toml';
 const FAUCET_DATA_FILE = 'faucet_data.sqlite';
-const FAUCET_DATA_TTL = 86400000; // 24 hrs
+const TTL_MULTIPLIER = 1.5; // 1.5 times the configured period
 
 interface Config {
   registry: {
@@ -25,12 +25,18 @@ interface Config {
   },
   server: {
     port: number
+    periodInSecs: number
     transferAmount: string
     periodTransferLimit: string
     dbDir: string
   }
 }
 
+interface FaucetStoreValue {
+  periodStart: string
+  amountSent: string
+}
+
 async function main (): Promise<void> {
   // Read and parse the configuration
   const configFile = fs.readFileSync(CONFIG_PATH, 'utf-8');
@@ -54,14 +60,28 @@ async function main (): Promise<void> {
       return res.status(400).json({ error: 'invalid address' });
     }
 
-    // Check rate limit
-    const now = Date.now();
-    const today = new Date(now).toISOString().split('T')[0];
-    const faucetStoreKey = `${accountAddress}:${today}`;
-    const amountSentToAddress = await faucetDataStore.get(faucetStoreKey) || '0';
+    const currentTimeInSeconds = Math.floor(Date.now() / 1000);
+    let amountSentToAccount = BigInt(0);
+    let accountPeriodStart = currentTimeInSeconds;
 
-    if (BigInt(amountSentToAddress) + BigInt(config.server.transferAmount) > BigInt(config.server.periodTransferLimit)) {
-      return res.status(429).json({ error: 'Limit exceeded' });
+    // Get accounts data from faucet data store
+    const accountFaucetData: FaucetStoreValue = await faucetDataStore.get(accountAddress);
+
+    // Check rate limit
+    if (accountFaucetData !== undefined) {
+      const elapsedTime = currentTimeInSeconds - Number(accountFaucetData.periodStart);
+
+      // Check if saved entry is still within configured period
+      // Else, create a fresh entry
+      if (elapsedTime < config.server.periodInSecs) {
+        amountSentToAccount = BigInt(accountFaucetData.amountSent);
+        accountPeriodStart = Number(accountFaucetData.periodStart);
+
+        // Check if limit has been reached
+        if (amountSentToAccount + BigInt(config.server.transferAmount) > BigInt(config.server.periodTransferLimit)) {
+          return res.status(429).json({ error: 'Limit exceeded' });
+        }
+      }
     }
 
     try {
@@ -69,7 +89,11 @@ async function main (): Promise<void> {
       console.log(`Sent tokens to address: ${accountAddress}, txHash: ${txHash}`);
 
       // Update rate limit
-      await faucetDataStore.set(faucetStoreKey, (BigInt(amountSentToAddress) + BigInt(config.server.transferAmount)).toString(), FAUCET_DATA_TTL);
+      const updatedAccountFaucetData: FaucetStoreValue = {
+        periodStart: accountPeriodStart.toString(),
+        amountSent: amountSentToAccount.toString()
+      };
+      await faucetDataStore.set(accountAddress, updatedAccountFaucetData, config.server.periodInSecs * TTL_MULTIPLIER * 1000);
 
       res.json({ success: true, txHash });
     } catch (error) {