diff --git a/whitelist/assembly/__tests__/whitelist.spec.ts b/whitelist/assembly/__tests__/whitelist.spec.ts new file mode 100644 index 0000000..11378f1 --- /dev/null +++ b/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(); + }); + }); +}); diff --git a/whitelist/assembly/index.ts b/whitelist/assembly/index.ts index 848b176..57549e0 100644 --- a/whitelist/assembly/index.ts +++ b/whitelist/assembly/index.ts @@ -1,3 +1,130 @@ -if (!ASC_NO_ASSERT) { - assert(false, "noAssert should be true"); -} \ No newline at end of file +import { context, Context, PersistentMap, storage, logging } from "near-sdk-as"; + +const FOUNDATION_KEY = "f"; + +// Use PersistentMap as a lookup set for whitelisted accounts +const whitelist = new PersistentMap("w"); +const factoryWhitelist = new PersistentMap("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); +}