Refactor and add tests

.github/workflows/tests.yaml

@@ -0,0 +1,25 @@
+# From https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-nodejs#using-the-nodejs-starter-workflow
+name: Node.js and Solana CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    strategy:
+      matrix:
+        node-version: [20.x]
+
+    steps:
+      - uses: actions/checkout@v3
+
+      # Install everything
+      - run: npm ci
+
+      # Run tests
+      - run: npm test

.gitignore

@@ -37,4 +37,6 @@ next-env.d.ts
 
 # database files
 walletData.json
-ipData.json
+ipData.json
+.env
+test-ledger

README.md

@@ -1 +1,41 @@
-Solana Devnet Faucet with rate limiting
+# Solana Devnet Faucet with rate limiting
+
+This is the code for the [Solana Devnet Faucet](https://faucet.solana.com/)
+
+## Run tests
+
+```
+npm run test
+```
+
+Or to run a single test, for example:
+
+```
+npx jest -t 'is a PDA'
+```
+
+## Run locally for development
+
+You'll need an `.env` file with:
+
+```
+FAUCET_KEYPAIR=[numbers...]
+POSTGRES_STRING="some string"
+```
+
+And then run:
+
+```
+npm run dev
+```
+
+## Deploy
+
+Deploying is done automatically as soon as the code is committed onto master via Vercel.
+
+Vercel also needs these details:
+
+```
+RPC_URL: "string"
+CLOUDFLARE_SECRET: "string"
+```

app/page.tsx

@@ -45,7 +45,7 @@ export default function Home() {
   });
   const [showVerifyDialog, setShowVerifyDialog] = useState<boolean>(false);
   const toaster = useToast();
-  const [network, setSelectedNetwork] = useState('devnet');
+  const [network, setSelectedNetwork] = useState("devnet");
 
   const validateWallet = (address: string): boolean => {
     try {
@@ -91,7 +91,12 @@ export default function Home() {
         headers: {
           "Content-Type": "application/json",
         },
-        body: JSON.stringify({ walletAddress, amount, cloudflareCallback, network: network }),
+        body: JSON.stringify({
+          walletAddress,
+          amount,
+          cloudflareCallback,
+          network: network,
+        }),
       });
 
       if (res.ok) {

jest.config.js

@@ -0,0 +1,5 @@
+/** @type {import('ts-jest').JestConfigWithTsJest} */
+module.exports = {
+  preset: 'ts-jest',
+  testEnvironment: 'node',
+};

lib/db.test.ts

@@ -0,0 +1,55 @@
+import { describe, expect, test } from "@jest/globals";
+import { Row, checkLimits } from "./db";
+import { Pool } from "pg";
+import { MINUTES } from "./constants";
+const log = console.log;
+
+let mockRows: Array<Row> = [];
+
+// Make a mock for the pg Pool constructor
+// https://jestjs.io/docs/mock-functions#mocking-modules
+jest.mock("pg", () => {
+  return {
+    Pool: jest.fn(() => ({
+      query: jest.fn(() => {
+        return {
+          rows: mockRows,
+        };
+      }),
+    })),
+  };
+});
+
+describe("checkLimits", () => {
+  // TODO: ideally I'd like to use mockValueOnce() instead of the
+  // mockRows variable, but I couldn't get it to work.
+  test("is fine when there's no previous usage", async () => {
+    mockRows = [];
+    await checkLimits("1.1.1.1");
+  });
+
+  test("allows reasonable usage", async () => {
+    mockRows = [
+      {
+        timestamps: [Date.now() - 10 * MINUTES],
+      },
+    ];
+    await checkLimits("1.1.1.1");
+  });
+
+  test("blocks unreasonable usage", async () => {
+    mockRows = [
+      {
+        timestamps: [
+          Date.now() - 10 * MINUTES,
+          Date.now() - 10 * MINUTES,
+          Date.now() - 10 * MINUTES,
+          Date.now() - 10 * MINUTES,
+        ],
+      },
+    ];
+    await expect(checkLimits("1.1.1.1")).rejects.toThrow(
+      "You have exceeded the 2 airdrops limit in the past 1 hour(s)"
+    );
+  });
+});

lib/db.ts

@@ -0,0 +1,62 @@
+import { Pool } from "pg";
+import { HOURS } from "./constants";
+const log = console.log;
+
+export interface Row {
+  timestamps: Array<number>;
+}
+
+const pgClient = new Pool({
+  connectionString: process.env.POSTGRES_STRING as string,
+});
+
+// Eg if AIRDROPS_LIMIT_TOTAL is 2, and AIRDROPS_LIMIT_HOURS is 1,
+// then a user can only get 2 airdrops per 1 hour.
+const AIRDROPS_LIMIT_TOTAL = 2;
+const AIRDROPS_LIMIT_HOURS = 1;
+
+// Formerly called 'getOrCreateAndVerifyDatabaseEntry'
+export const checkLimits = async (
+  ipAddressWithoutDotsOrWalletAddress: string
+): Promise<void> => {
+  // Remove the . (IPV4) and : (IPV6) from the IP address
+  let databaseKey = ipAddressWithoutDotsOrWalletAddress.replace(/[\.,:]/g, "");
+
+  const entryQuery = "SELECT * FROM rate_limits WHERE key = $1;";
+  const insertQuery =
+    "INSERT INTO rate_limits (key, timestamps) VALUES ($1, $2);";
+  const updateQuery = "UPDATE rate_limits SET timestamps = $2 WHERE key = $1;";
+
+  const timeAgo = Date.now() - AIRDROPS_LIMIT_HOURS * HOURS;
+
+  const queryResult = await pgClient.query(entryQuery, [databaseKey]);
+
+  const rows = queryResult.rows as Array<Row>;
+  const entry = rows[0];
+
+  if (entry) {
+    const timestamps = entry.timestamps;
+
+    const isExcessiveUsage =
+      timestamps.filter((timestamp: number) => timestamp > timeAgo).length >=
+      AIRDROPS_LIMIT_TOTAL;
+
+    if (isExcessiveUsage) {
+      throw new Error(
+        `You have exceeded the ${AIRDROPS_LIMIT_TOTAL} airdrops limit in the past ${AIRDROPS_LIMIT_HOURS} hour(s)`
+      );
+    }
+
+    timestamps.push(Date.now());
+
+    await pgClient.query(updateQuery, [
+      ipAddressWithoutDotsOrWalletAddress,
+      timestamps,
+    ]);
+  } else {
+    await pgClient.query(insertQuery, [
+      ipAddressWithoutDotsOrWalletAddress,
+      [Date.now()],
+    ]);
+  }
+};

lib/utils.test.ts

@@ -0,0 +1,35 @@
+import { NextApiRequest } from "next";
+import { getHeaderValues } from "./utils";
+
+describe("getHeaderValues", () => {
+  it("returns an array of strings when x-forwarded-for is a string", () => {
+    const req = {
+      headers: {
+        "x-forwarded-for": "10.0.0.0",
+      },
+    } as unknown as NextApiRequest;
+    const headerName = "x-forwarded-for";
+    const result = getHeaderValues(req, headerName);
+    expect(result).toEqual(["10.0.0.0"]);
+  });
+
+  it("returns an array of strings when x-forwarded-for is an array", () => {
+    const req = {
+      headers: {
+        "x-forwarded-for": ["10.0.0.0"],
+      },
+    } as unknown as NextApiRequest;
+    const headerName = "x-forwarded-for";
+    const result = getHeaderValues(req, headerName);
+    expect(result).toEqual(["10.0.0.0"]);
+  });
+
+  it("returns an array of strings when x-forwarded-for does not exist", () => {
+    const req = {
+      headers: {},
+    } as unknown as NextApiRequest;
+    const headerName = "x-forwarded-for";
+    const result = getHeaderValues(req, headerName);
+    expect(result).toEqual([]);
+  });
+});

lib/utils.ts

@@ -1,6 +1,17 @@
 import { type ClassValue, clsx } from "clsx";
+import { NextApiRequest } from "next";
 import { twMerge } from "tailwind-merge";
 
 export function cn(...inputs: ClassValue[]) {
   return twMerge(clsx(inputs));
 }
+
+export const getHeaderValues = (req: NextApiRequest, headerName: string) => {
+  // Annoyingly, req.headers["x-forwarded-for"] can be a string or an array of strings
+  // Let's just make it an array of strings
+  let valueOrValues = req.headers[headerName] || [];
+  if (Array.isArray(valueOrValues)) {
+    return valueOrValues;
+  }
+  return [valueOrValues];
+};

lib/validate.test.ts

@@ -0,0 +1,47 @@
+import { validate } from "./validate";
+
+const WALLET_ADDRESS_FOR_TESTS = "dDCQNnDmNbFVi8cQhKAgXhyhXeJ625tvwsunRyRc7c8";
+const PDA_ADDRESS_FOR_TESTS = "4MD31b2GFAWVDYQT8KG7E5GcZiFyy4MpDUt4BcyEdJRP";
+
+// Write some tests for the validate function
+describe("validate", () => {
+  test("allows reasonable usage", () => {
+    validate("DXJfhtWicZwBpHGiBepWwwnJK7jJYNYguGDUgNYbMCCi", 1);
+  });
+
+  test("throws when wallet address is a PDA", () => {
+    expect(() => {
+      validate(PDA_ADDRESS_FOR_TESTS, 1);
+    }).toThrow("Please enter valid wallet address.");
+  });
+
+  test("throws when wallet address is empty string", () => {
+    expect(() => {
+      validate("", 1);
+    }).toThrow("Missing wallet address.");
+  });
+
+  test("throws when amount is 0", () => {
+    expect(() => {
+      validate(WALLET_ADDRESS_FOR_TESTS, 0);
+    }).toThrow("Missing SOL amount.");
+  });
+
+  test("throws when amount is negative", () => {
+    expect(() => {
+      validate(WALLET_ADDRESS_FOR_TESTS, -3);
+    }).toThrow("Requested SOL amount cannot be negative.");
+  });
+
+  test("throws when amount is too large", () => {
+    expect(() => {
+      validate(WALLET_ADDRESS_FOR_TESTS, 6);
+    }).toThrow("Requested SOL amount too large.");
+  });
+
+  test("throws when wallet address is invalid", () => {
+    expect(() => {
+      validate("invalidWalletAddress", 1);
+    }).toThrow("Please enter valid wallet address.");
+  });
+});

lib/validate.ts

@@ -0,0 +1,31 @@
+import { PublicKey } from "@solana/web3.js";
+
+const MAX_SOL_AMOUNT = 5;
+
+export const validate = (walletAddress: string, amount: number): void => {
+  if (!walletAddress) {
+    throw new Error("Missing wallet address.");
+  }
+
+  if (!amount) {
+    throw new Error("Missing SOL amount.");
+  }
+
+  if (amount < 0) {
+    throw new Error("Requested SOL amount cannot be negative.");
+  }
+
+  if (amount > MAX_SOL_AMOUNT) {
+    throw new Error("Requested SOL amount too large.");
+  }
+
+  try {
+    let pubkey = new PublicKey(walletAddress);
+    let isOnCurve = PublicKey.isOnCurve(pubkey.toBuffer());
+    if (!isOnCurve) {
+      throw new Error("Address can't be a PDA.");
+    }
+  } catch (error) {
+    throw new Error("Please enter valid wallet address.");
+  }
+};