feat(payments): add address whitelist

Author: MohammadNassar1

src/errors.cairo

@@ -1,5 +1,7 @@
 use starknet::ContractAddress;
 
+pub const ADDRESS_ALREADY_WHITELISTED: felt252 = 'ADDRESS_ALREADY_WHITELISTED';
+pub const ADDRESS_NOT_WHITELISTED: felt252 = 'ADDRESS_NOT_WHITELISTED';
 pub const INVALID_AMOUNT_RATIO: felt252 = 'INVALID_AMOUNT_RATIO';
 pub const INVALID_AMOUNT_TOO_LARGE: felt252 = 'INVALID_AMOUNT_TOO_LARGE';
 pub const INVALID_DOWNCAST_AFTER_DIVISION: felt252 = 'INVALID_DOWNCAST_AFTER_DIVISION';

src/events.cairo

@@ -25,6 +25,18 @@ pub struct TokenRemoved {
     pub token: ContractAddress,
 }
 
+#[derive(Debug, Drop, PartialEq, starknet::Event)]
+pub struct AddressWhitelisted {
+    #[key]
+    pub address: ContractAddress,
+}
+
+#[derive(Debug, Drop, PartialEq, starknet::Event)]
+pub struct AddressRemovedFromWhitelist {
+    #[key]
+    pub address: ContractAddress,
+}
+
 #[derive(Debug, Drop, PartialEq, starknet::Event)]
 pub struct TradeExecuted {
     #[key]

src/interface.cairo

@@ -18,6 +18,10 @@ pub trait IPayments<TContractState> {
     fn remove_token(ref self: TContractState, token: ContractAddress);
     fn is_token_registered(self: @TContractState, token: ContractAddress) -> bool;
 
+    fn whitelist_address(ref self: TContractState, address: ContractAddress);
+    fn remove_from_whitelist(ref self: TContractState, address: ContractAddress);
+    fn is_whitelisted(self: @TContractState, address: ContractAddress) -> bool;
+
     fn cancel_orders(ref self: TContractState, orders: Span<Order>);
 
     // Setters:

src/payments.cairo

@@ -21,13 +21,15 @@ pub mod payments {
     use starkware_utils::signature::stark::{HashType, Signature};
     use starkware_utils::time::time::Time;
     use crate::errors::{
-        INVALID_AMOUNT_RATIO, INVALID_AMOUNT_TOO_LARGE, INVALID_DOWNCAST_AFTER_DIVISION,
-        INVALID_HIGH_FEE, INVALID_HIGH_FEE_LIMIT, INVALID_TOKEN_PAIR, INVALID_TRADE_SAME_USER,
-        INVALID_ZERO_ADDRESS, INVALID_ZERO_AMOUNT, INVALID_ZERO_TOKEN, ORDER_EXPIRED,
-        TOKEN_ALREADY_REGISTERED, TOKEN_NOT_REGISTERED, UNALLOWED_ADDRESS, transfer_failed_error,
+        ADDRESS_ALREADY_WHITELISTED, ADDRESS_NOT_WHITELISTED, INVALID_AMOUNT_RATIO,
+        INVALID_AMOUNT_TOO_LARGE, INVALID_DOWNCAST_AFTER_DIVISION, INVALID_HIGH_FEE,
+        INVALID_HIGH_FEE_LIMIT, INVALID_TOKEN_PAIR, INVALID_TRADE_SAME_USER, INVALID_ZERO_ADDRESS,
+        INVALID_ZERO_AMOUNT, INVALID_ZERO_TOKEN, ORDER_EXPIRED, TOKEN_ALREADY_REGISTERED,
+        TOKEN_NOT_REGISTERED, UNALLOWED_ADDRESS, transfer_failed_error,
     };
     use crate::events::{
-        FeeRecipientSet, FeeSet, OrderCanceled, TokenRegistered, TokenRemoved, TradeExecuted,
+        AddressRemovedFromWhitelist, AddressWhitelisted, FeeRecipientSet, FeeSet, OrderCanceled,
+        TokenRegistered, TokenRemoved, TradeExecuted,
     };
     use crate::interface::IPayments;
     use crate::order::Order;
@@ -85,6 +87,8 @@ pub mod payments {
         fee: u128,
         // Whitelisted tokens.
         tokens: Map<ContractAddress, bool>,
+        // Whitelisted addresses.
+        whitelist: Map<ContractAddress, bool>,
         // Order hash to fulfilled sell amount.
         fulfillment: Map<HashType, u128>,
     }
@@ -106,6 +110,8 @@ pub mod payments {
         FeeRecipientSet: FeeRecipientSet,
         TokenRegistered: TokenRegistered,
         TokenRemoved: TokenRemoved,
+        AddressWhitelisted: AddressWhitelisted,
+        AddressRemovedFromWhitelist: AddressRemovedFromWhitelist,
         TradeExecuted: TradeExecuted,
         OrderCanceled: OrderCanceled,
     }
@@ -282,6 +288,35 @@ pub mod payments {
             self.tokens.read(token)
         }
 
+        // These functions are used by the contract to control who is allowed to take part in order
+        // matching and settlement. Only whitelisted addresses are permitted to trade.
+
+        fn whitelist_address(ref self: ContractState, address: ContractAddress) {
+            self.roles.only_app_governor();
+
+            assert(!self.is_whitelisted(address), ADDRESS_ALREADY_WHITELISTED);
+            self.whitelist.write(address, true);
+
+            // Emit an event.
+            self.emit(AddressWhitelisted { address });
+        }
+
+        fn remove_from_whitelist(ref self: ContractState, address: ContractAddress) {
+            self.roles.only_app_governor();
+
+            assert(self.is_whitelisted(address), ADDRESS_NOT_WHITELISTED);
+            self.whitelist.write(address, false);
+
+            // Emit an event.
+            self.emit(AddressRemovedFromWhitelist { address });
+        }
+
+        fn is_whitelisted(self: @ContractState, address: ContractAddress) -> bool {
+            assert(address.is_non_zero(), INVALID_ZERO_ADDRESS);
+            self.whitelist.read(address)
+        }
+
+
         fn cancel_orders(ref self: ContractState, orders: Span<Order>) {
             self.roles.only_operator();
 
@@ -370,7 +405,7 @@ pub mod payments {
         /// amounts.
         fn _validate_order(self: @ContractState, order: Order) {
             assert(order.expiry >= Time::now(), ORDER_EXPIRED);
-            assert(order.user.is_non_zero(), INVALID_ZERO_ADDRESS);
+            assert(self.is_whitelisted(order.user), ADDRESS_NOT_WHITELISTED);
 
             assert(order.sell_amount.is_non_zero(), INVALID_ZERO_AMOUNT);
             assert(order.buy_amount.is_non_zero(), INVALID_ZERO_AMOUNT);

src/tests/test_payments.cairo

@@ -102,6 +102,78 @@ fn test_failed_register_token() {
     assert_panic_with_felt_error(:result, expected_error: errors::TOKEN_ALREADY_REGISTERED);
 }
 
+#[test]
+fn test_successful_address_whitelist() {
+    let contract_address = init_contract_with_roles();
+    let dispatcher = IPaymentsDispatcher { contract_address };
+    let mut spy = snforge_std::spy_events();
+
+    let user_a: ContractAddress = 'user_a'.try_into().unwrap();
+    let user_b: ContractAddress = 'user_b'.try_into().unwrap();
+
+    cheat_caller_address_once(:contract_address, caller_address: testing_constants::APP_GOVERNOR);
+    dispatcher.whitelist_address(address: user_a);
+    assert!(dispatcher.is_whitelisted(address: user_a));
+    assert!(!dispatcher.is_whitelisted(address: user_b));
+
+    cheat_caller_address_once(:contract_address, caller_address: testing_constants::APP_GOVERNOR);
+    dispatcher.whitelist_address(address: user_b);
+    assert!(dispatcher.is_whitelisted(address: user_a));
+    assert!(dispatcher.is_whitelisted(address: user_b));
+
+    cheat_caller_address_once(:contract_address, caller_address: testing_constants::APP_GOVERNOR);
+    dispatcher.remove_from_whitelist(address: user_a);
+    assert!(!dispatcher.is_whitelisted(address: user_a));
+    assert!(dispatcher.is_whitelisted(address: user_b));
+
+    cheat_caller_address_once(:contract_address, caller_address: testing_constants::APP_GOVERNOR);
+    dispatcher.remove_from_whitelist(address: user_b);
+    assert!(!dispatcher.is_whitelisted(address: user_a));
+    assert!(!dispatcher.is_whitelisted(address: user_b));
+
+    // Catch the events.
+    let events = spy.get_events().emitted_by(contract_address).events;
+    let expected_register_token_event = events::AddressWhitelisted { address: user_a };
+    assert_expected_event_emitted(
+        spied_event: events[0],
+        expected_event: expected_register_token_event,
+        expected_event_selector: @selector!("AddressWhitelisted"),
+        expected_event_name: "AddressWhitelisted",
+    );
+    let expected_remove_token_event = events::AddressRemovedFromWhitelist { address: user_b };
+    assert_expected_event_emitted(
+        spied_event: events[3],
+        expected_event: expected_remove_token_event,
+        expected_event_selector: @selector!("AddressRemovedFromWhitelist"),
+        expected_event_name: "AddressRemovedFromWhitelist",
+    );
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_failed_address_whitelist() {
+    let contract_address = init_contract_with_roles();
+    let dispatcher = IPaymentsSafeDispatcher { contract_address };
+    let user: ContractAddress = 'user'.try_into().unwrap();
+
+    let result = dispatcher.remove_from_whitelist(address: user);
+    assert_panic_with_error(:result, expected_error: "ONLY_APP_GOVERNOR");
+
+    let result = dispatcher.whitelist_address(address: user);
+    assert_panic_with_error(:result, expected_error: "ONLY_APP_GOVERNOR");
+
+    cheat_caller_address_once(:contract_address, caller_address: testing_constants::APP_GOVERNOR);
+    let result = dispatcher.remove_from_whitelist(address: user);
+    assert_panic_with_felt_error(:result, expected_error: errors::ADDRESS_NOT_WHITELISTED);
+
+    cheat_caller_address_once(:contract_address, caller_address: testing_constants::APP_GOVERNOR);
+    dispatcher.whitelist_address(address: user).unwrap();
+
+    cheat_caller_address_once(:contract_address, caller_address: testing_constants::APP_GOVERNOR);
+    let result = dispatcher.whitelist_address(address: user);
+    assert_panic_with_felt_error(:result, expected_error: errors::ADDRESS_ALREADY_WHITELISTED);
+}
+
 #[test]
 fn test_successful_set_fee() {
     let contract_address = init_contract_with_roles();