From a1e1f79d02071cbe72e1280d40d680a11406c855 Mon Sep 17 00:00:00 2001
From: dragosrebegea <dragos.rebegea@multiversx.com>
Date: Fri, 18 Oct 2024 22:34:11 +0300
Subject: [PATCH] add protection

---
 contracts/faucet/src/faucet.rs   | 44 +++++++++++++++++++++++++++++---
 contracts/faucet/wasm/Cargo.toml | 34 ++++++++++++++++++++++++
 contracts/faucet/wasm/src/lib.rs | 29 +++++++++++++++++++++
 3 files changed, 104 insertions(+), 3 deletions(-)
 create mode 100644 contracts/faucet/wasm/Cargo.toml
 create mode 100644 contracts/faucet/wasm/src/lib.rs

diff --git a/contracts/faucet/src/faucet.rs b/contracts/faucet/src/faucet.rs
index 4a1cb36e..e88cc4b3 100644
--- a/contracts/faucet/src/faucet.rs
+++ b/contracts/faucet/src/faucet.rs
@@ -3,12 +3,18 @@
 #[allow(unused_imports)]
 use multiversx_sc::{derive_imports::*, imports::*};
 
+pub type Signature<M> = ManagedByteArray<M, ED25519_SIGNATURE_BYTE_LEN>;
+pub const ED25519_SIGNATURE_BYTE_LEN: usize = 64;
+
 /// An empty contract. To be used as a template when starting a new contract from scratch.
 #[multiversx_sc::contract]
 pub trait Faucet {
     #[init]
     fn init(&self) {}
 
+    #[upgrade]
+    fn upgrade(&self) {}
+
     #[only_owner]
     #[payable("*")]
     #[endpoint(deposit)]
@@ -21,7 +27,14 @@ pub trait Faucet {
     }
 
     #[endpoint(claim)]
-    fn claim(&self) {
+    fn claim(&self, signature: &Signature<Self::Api>) {
+        let caller = self.blockchain().get_caller();
+        let claimers_mapper = self.claimers();
+        require!(claimers_mapper.contains(&caller) == false, "Caller already claimed");
+
+        self.verify_signature(&caller, signature);
+        self.claimers().add(&caller);
+
         let mut payments = ManagedVec::new();
         for payment in self.payments().iter() {
             payments.push(payment);
@@ -30,9 +43,34 @@ pub trait Faucet {
         self.tx().to(ToCaller).with_multi_token_transfer(payments).transfer();
     }
 
+    fn verify_signature(
+        &self, 
+        caller: &ManagedAddress<Self::Api>,
+        signature: &Signature<Self::Api>,
+    ) {
+        let mut data = ManagedBuffer::new();
+        let _ = caller.dep_encode(&mut data);
+
+        let signer = self.signer().get();
+        self.crypto().verify_ed25519(
+            signer.as_managed_buffer(),
+            &data,
+            signature.as_managed_buffer(),
+        );
+    }
+
+    #[only_owner]
+    #[endpoint(setSigner)]
+    fn set_signer(&self, signer: ManagedAddress<Self::Api>) {
+        self.signer().set(&signer);
+    }
+
     #[storage_mapper("payments")]
     fn payments(&self) -> VecMapper<EsdtTokenPayment<Self::Api>>;
 
-    #[upgrade]
-    fn upgrade(&self) {}
+    #[storage_mapper("signer")]
+    fn signer(&self) -> SingleValueMapper<ManagedAddress<Self::Api>>;
+
+    #[storage_mapper("claimers")]
+    fn claimers(&self) -> WhitelistMapper<ManagedAddress<Self::Api>>;
 }
diff --git a/contracts/faucet/wasm/Cargo.toml b/contracts/faucet/wasm/Cargo.toml
new file mode 100644
index 00000000..8b20a063
--- /dev/null
+++ b/contracts/faucet/wasm/Cargo.toml
@@ -0,0 +1,34 @@
+# Code generated by the multiversx-sc build system. DO NOT EDIT.
+
+# ##########################################
+# ############## AUTO-GENERATED #############
+# ##########################################
+
+[package]
+name = "faucet-wasm"
+version = "0.0.0"
+edition = "2021"
+publish = false
+
+[lib]
+crate-type = ["cdylib"]
+
+[profile.release]
+codegen-units = 1
+opt-level = "z"
+lto = true
+debug = false
+panic = "abort"
+overflow-checks = false
+
+[profile.dev]
+panic = "abort"
+
+[dependencies.faucet]
+path = ".."
+
+[dependencies.multiversx-sc-wasm-adapter]
+version = "0.53.2"
+
+[workspace]
+members = ["."]
diff --git a/contracts/faucet/wasm/src/lib.rs b/contracts/faucet/wasm/src/lib.rs
new file mode 100644
index 00000000..b471a0c2
--- /dev/null
+++ b/contracts/faucet/wasm/src/lib.rs
@@ -0,0 +1,29 @@
+// Code generated by the multiversx-sc build system. DO NOT EDIT.
+
+////////////////////////////////////////////////////
+////////////////// AUTO-GENERATED //////////////////
+////////////////////////////////////////////////////
+
+// Init:                                 1
+// Upgrade:                              1
+// Endpoints:                            3
+// Async Callback (empty):               1
+// Total number of exported functions:   6
+
+#![no_std]
+
+multiversx_sc_wasm_adapter::allocator!();
+multiversx_sc_wasm_adapter::panic_handler!();
+
+multiversx_sc_wasm_adapter::endpoints! {
+    faucet
+    (
+        init => init
+        upgrade => upgrade
+        deposit => deposit
+        claim => claim
+        setSigner => set_signer
+    )
+}
+
+multiversx_sc_wasm_adapter::async_callback_empty! {}