fix: [Bounty] Whitelist

whitelist/assembly/__tests__/whitelist.spec.ts

@@ -0,0 +1,182 @@
+import { VMContext, Context } from "near-sdk-as";
+import {
+  new_default,
+  add_staking_pool,
+  remove_staking_pool,
+  is_whitelisted,
+  add_factory,
+  remove_factory,
+  is_factory_whitelisted,
+} from "../index";
+
+const FOUNDATION = "foundation.near";
+const POOL1 = "pool1.near";
+const POOL2 = "pool2.near";
+const FACTORY1 = "factory1.near";
+const FACTORY2 = "factory2.near";
+const NON_FOUNDATION = "alice.near";
+
+function setFoundation(): void {
+  VMContext.setPredecessor_account_id(FOUNDATION);
+}
+
+function init(): void {
+  setFoundation();
+  new_default(FOUNDATION);
+}
+
+describe("Whitelist Contract", () => {
+  beforeEach(() => {
+    init();
+  });
+
+  describe("Initialization", () => {
+    it("should not allow double initialization", () => {
+      expect(() => {
+        new_default(FOUNDATION);
+      }).toThrow();
+    });
+  });
+
+  describe("Staking Pool Whitelist (Foundation)", () => {
+    it("should add a staking pool to whitelist", () => {
+      const result = add_staking_pool(POOL1);
+      expect(result).toBeTruthy();
+      expect(is_whitelisted(POOL1)).toBeTruthy();
+    });
+
+    it("should return false when adding already whitelisted pool", () => {
+      add_staking_pool(POOL1);
+      const result = add_staking_pool(POOL1);
+      expect(result).toBeFalsy();
+    });
+
+    it("should remove a staking pool from whitelist", () => {
+      add_staking_pool(POOL1);
+      const result = remove_staking_pool(POOL1);
+      expect(result).toBeTruthy();
+      expect(is_whitelisted(POOL1)).toBeFalsy();
+    });
+
+    it("should return false when removing non-whitelisted pool", () => {
+      const result = remove_staking_pool(POOL1);
+      expect(result).toBeFalsy();
+    });
+
+    it("should handle multiple pools", () => {
+      add_staking_pool(POOL1);
+      add_staking_pool(POOL2);
+      expect(is_whitelisted(POOL1)).toBeTruthy();
+      expect(is_whitelisted(POOL2)).toBeTruthy();
+    });
+
+    it("should not whitelist a pool by default", () => {
+      expect(is_whitelisted(POOL1)).toBeFalsy();
+    });
+  });
+
+  describe("Access Control", () => {
+    it("should not allow non-foundation to remove staking pool", () => {
+      add_staking_pool(POOL1);
+      VMContext.setPredecessor_account_id(NON_FOUNDATION);
+      expect(() => {
+        remove_staking_pool(POOL1);
+      }).toThrow();
+    });
+
+    it("should not allow non-foundation/non-factory to add staking pool", () => {
+      VMContext.setPredecessor_account_id(NON_FOUNDATION);
+      expect(() => {
+        add_staking_pool(POOL1);
+      }).toThrow();
+    });
+
+    it("should not allow non-foundation to add factory", () => {
+      VMContext.setPredecessor_account_id(NON_FOUNDATION);
+      expect(() => {
+        add_factory(FACTORY1);
+      }).toThrow();
+    });
+
+    it("should not allow non-foundation to remove factory", () => {
+      add_factory(FACTORY1);
+      VMContext.setPredecessor_account_id(NON_FOUNDATION);
+      expect(() => {
+        remove_factory(FACTORY1);
+      }).toThrow();
+    });
+  });
+
+  describe("Factory Whitelist", () => {
+    it("should add a factory to whitelist", () => {
+      const result = add_factory(FACTORY1);
+      expect(result).toBeTruthy();
+      expect(is_factory_whitelisted(FACTORY1)).toBeTruthy();
+    });
+
+    it("should return false when adding already whitelisted factory", () => {
+      add_factory(FACTORY1);
+      const result = add_factory(FACTORY1);
+      expect(result).toBeFalsy();
+    });
+
+    it("should remove a factory from whitelist", () => {
+      add_factory(FACTORY1);
+      const result = remove_factory(FACTORY1);
+      expect(result).toBeTruthy();
+      expect(is_factory_whitelisted(FACTORY1)).toBeFalsy();
+    });
+
+    it("should return false when removing non-whitelisted factory", () => {
+      const result = remove_factory(FACTORY1);
+      expect(result).toBeFalsy();
+    });
+
+    it("should not whitelist a factory by default", () => {
+      expect(is_factory_whitelisted(FACTORY1)).toBeFalsy();
+    });
+  });
+
+  describe("Factory Adding Staking Pool", () => {
+    it("should allow whitelisted factory to add staking pool", () => {
+      add_factory(FACTORY1);
+      VMContext.setPredecessor_account_id(FACTORY1);
+      const result = add_staking_pool(POOL1);
+      expect(result).toBeTruthy();
+      expect(is_whitelisted(POOL1)).toBeTruthy();
+    });
+
+    it("should not allow non-whitelisted factory to add staking pool", () => {
+      VMContext.setPredecessor_account_id(FACTORY1);
+      expect(() => {
+        add_staking_pool(POOL1);
+      }).toThrow();
+    });
+  });
+
+  describe("Validation", () => {
+    it("should reject empty staking pool account ID", () => {
+      expect(() => {
+        add_staking_pool("");
+      }).toThrow();
+    });
+
+    it("should reject empty factory account ID", () => {
+      expect(() => {
+        add_factory("");
+      }).toThrow();
+    });
+
+    it("should reject empty account ID for is_whitelisted", () => {
+      expect(() => {
+        is_whitelisted("");
+      }).toThrow();
+    });
+
+    it("should reject empty account ID for is_factory_whitelisted", () => {
+      expect(() => {
+        is_factory_whitelisted("");
+      }).toThrow();
+    });
+  });
+});

whitelist/assembly/index.ts

@@ -1,3 +1,130 @@
-if (!ASC_NO_ASSERT) {
-  assert(false, "noAssert should be true");
-}
+import { context, Context, PersistentMap, storage, logging } from "near-sdk-as";
+
+const FOUNDATION_KEY = "f";
+
+// Use PersistentMap<string, bool> as a lookup set for whitelisted accounts
+const whitelist = new PersistentMap<string, bool>("w");
+const factoryWhitelist = new PersistentMap<string, bool>("fw");
+
+/**
+ * Returns the foundation account ID stored during initialization.
+ */
+function getFoundationAccountId(): string {
+  return storage.getString(FOUNDATION_KEY)!;
+}
+
+/**
+ * Asserts that the predecessor (caller) is the foundation account.
+ */
+function assertCalledByFoundation(): void {
+  assert(
+    Context.predecessor == getFoundationAccountId(),
+    "Can only be called by the foundation"
+  );
+}
+
+/**
+ * Validates that the given account ID is valid (non-empty).
+ */
+function assertValidAccountId(accountId: string): void {
+  assert(accountId.length > 0, "The account ID is invalid");
+}
+
+// ─── Public contract methods ───
+
+/**
+ * Initializes the contract with the given foundation account ID.
+ * Can only be called once.
+ */
+export function new_default(foundation_account_id: string): void {
+  assert(!storage.contains(FOUNDATION_KEY), "The contract is already initialized");
+  assertValidAccountId(foundation_account_id);
+  storage.setString(FOUNDATION_KEY, foundation_account_id);
+}
+
+/**
+ * Adds a staking pool account to the whitelist.
+ * Can be called by the foundation or a whitelisted factory.
+ * Returns true if the pool was not already whitelisted.
+ */
+export function add_staking_pool(staking_pool_account_id: string): bool {
+  assertValidAccountId(staking_pool_account_id);
+  const predecessor = Context.predecessor;
+  const foundationAccountId = getFoundationAccountId();
+  if (predecessor != foundationAccountId) {
+    assert(
+      factoryWhitelist.contains(predecessor) && factoryWhitelist.getSome(predecessor),
+      "Can only be called by the foundation or a whitelisted factory"
+    );
+  }
+  const alreadyWhitelisted = whitelist.contains(staking_pool_account_id) && whitelist.getSome(staking_pool_account_id);
+  if (!alreadyWhitelisted) {
+    whitelist.set(staking_pool_account_id, true);
+    logging.log("Added staking pool account " + staking_pool_account_id + " to the whitelist");
+  }
+  return !alreadyWhitelisted;
+}
+
+/**
+ * Removes a staking pool account from the whitelist.
+ * Can only be called by the foundation.
+ * Returns true if the pool was previously whitelisted.
+ */
+export function remove_staking_pool(staking_pool_account_id: string): bool {
+  assertCalledByFoundation();
+  assertValidAccountId(staking_pool_account_id);
+  const wasWhitelisted = whitelist.contains(staking_pool_account_id) && whitelist.getSome(staking_pool_account_id);
+  if (wasWhitelisted) {
+    whitelist.set(staking_pool_account_id, false);
+    logging.log("Removed staking pool account " + staking_pool_account_id + " from the whitelist");
+  }
+  return wasWhitelisted;
+}
+
+/**
+ * Returns whether the given staking pool account is whitelisted.
+ */
+export function is_whitelisted(staking_pool_account_id: string): bool {
+  assertValidAccountId(staking_pool_account_id);
+  return whitelist.contains(staking_pool_account_id) && whitelist.getSome(staking_pool_account_id);
+}
+
+/**
+ * Adds a factory account to the factory whitelist.
+ * Can only be called by the foundation.
+ * Returns true if the factory was not already whitelisted.
+ */
+export function add_factory(factory_account_id: string): bool {
+  assertCalledByFoundation();
+  assertValidAccountId(factory_account_id);
+  const alreadyWhitelisted = factoryWhitelist.contains(factory_account_id) && factoryWhitelist.getSome(factory_account_id);
+  if (!alreadyWhitelisted) {
+    factoryWhitelist.set(factory_account_id, true);
+    logging.log("Added factory account " + factory_account_id + " to the factory whitelist");
+  }
+  return !alreadyWhitelisted;
+}
+
+/**
+ * Removes a factory account from the factory whitelist.
+ * Can only be called by the foundation.
+ * Returns true if the factory was previously whitelisted.
+ */
+export function remove_factory(factory_account_id: string): bool {
+  assertCalledByFoundation();
+  assertValidAccountId(factory_account_id);
+  const wasWhitelisted = factoryWhitelist.contains(factory_account_id) && factoryWhitelist.getSome(factory_account_id);
+  if (wasWhitelisted) {
+    factoryWhitelist.set(factory_account_id, false);
+    logging.log("Removed factory account " + factory_account_id + " from the factory whitelist");
+  }
+  return wasWhitelisted;
+}
+
+/**
+ * Returns whether the given factory account is whitelisted.
+ */
+export function is_factory_whitelisted(factory_account_id: string): bool {
+  assertValidAccountId(factory_account_id);
+  return factoryWhitelist.contains(factory_account_id) && factoryWhitelist.getSome(factory_account_id);
+}