From eb4e8d4d0a001fa19655a67d91fd72532dc9d714 Mon Sep 17 00:00:00 2001
From: ghassmo <ghxxm@protonmail.com>
Date: Fri, 24 Sep 2021 21:59:44 +0300
Subject: [PATCH] cashierd: check if it's a valid token id when receive
 deposit/withraw request

---
 src/bin/cashierd.rs | 34 ++++++++++++++++++++++++++++++++++
 src/error.rs        |  4 +++-
 src/service/sol.rs  |  2 +-
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/bin/cashierd.rs b/src/bin/cashierd.rs
index 9e10d4195..d2f588e07 100644
--- a/src/bin/cashierd.rs
+++ b/src/bin/cashierd.rs
@@ -296,6 +296,8 @@ impl Cashierd {
         }
 
         let result: Result<String> = async {
+            Self::check_token_id(&network, token_id.as_str().unwrap())?;
+
             let asset_id = drk::util::parse_id(token_id)?;
 
             let drk_pub_key = bs58::decode(&drk_pub_key).into_vec()?;
@@ -389,6 +391,8 @@ impl Cashierd {
         }
 
         let result: Result<String> = async {
+            Self::check_token_id(&network, token.as_str().unwrap())?;
+
             let asset_id = drk::util::parse_id(&token)?;
             let address = serialize(&address.to_string());
 
@@ -427,6 +431,36 @@ impl Cashierd {
     async fn features(&self, id: Value, _params: Value) -> JsonResult {
         JsonResult::Resp(jsonresp(json!(self.features), id))
     }
+
+    fn check_token_id(network: &str, _token_id: &str) -> Result<()> {
+        match network {
+            #[cfg(feature = "sol")]
+            "sol" | "solana" => {
+                if _token_id != "So11111111111111111111111111111111111111112" {
+                    // This is supposed to be a token mint account now
+                    use drk::service::sol::account_is_initialized_mint;
+                    use drk::service::sol::SolFailed::BadSolAddress;
+                    use solana_sdk::pubkey::Pubkey;
+                    use std::str::FromStr;
+
+                    if !account_is_initialized_mint(
+                        &Pubkey::from_str(_token_id)
+                            .map_err(|err| Error::from(BadSolAddress(err.to_string())))?,
+                    ) {
+                        return Err(Error::CashierInvalidTokenId(
+                            "Given address is not a valid token mint".into(),
+                        ));
+                    }
+                }
+            }
+            #[cfg(feature = "btc")]
+            "btc" | "bitcoin" => {
+                // Handle bitcoin address here if needed
+            }
+            _ => {}
+        }
+        Ok(())
+    }
 }
 
 #[async_std::main]
diff --git a/src/error.rs b/src/error.rs
index e28938ebb..d3b027490 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -59,6 +59,7 @@ pub enum Error {
     TomlDeserializeError(String),
     TomlSerializeError(String),
     CashierNoReply,
+    CashierInvalidTokenId(String),
     Base58EncodeError(String),
     Base58DecodeError(String),
     ConfigNotFound,
@@ -124,7 +125,8 @@ impl fmt::Display for Error {
             Error::TomlSerializeError(ref err) => write!(f, "Toml parsing error: {}", err),
             Error::Base58EncodeError(ref err) => write!(f, "bs58 encode error: {}", err),
             Error::Base58DecodeError(ref err) => write!(f, "bs58 decode error: {}", err),
-            Error::CashierNoReply => f.write_str("Cashier did not reply with BTC address"),
+            Error::CashierInvalidTokenId(ref err) => write!(f, "Cashier invalid token id: {}", err),
+            Error::CashierNoReply => f.write_str("Cashier did not reply with token address"),
             Error::ConfigNotFound => {
                 f.write_str("No config file detected. Please create a config file")
             }
diff --git a/src/service/sol.rs b/src/service/sol.rs
index 18e08a35d..ac0e46101 100644
--- a/src/service/sol.rs
+++ b/src/service/sol.rs
@@ -317,7 +317,7 @@ fn get_associated_token_account(owner: &Pubkey, mint: &Pubkey) -> (Pubkey, u8) {
 }
 
 /// Check if given account is a valid token mint
-fn account_is_initialized_mint(mint: &Pubkey) -> bool {
+pub fn account_is_initialized_mint(mint: &Pubkey) -> bool {
     let rpc = RpcClient::new(RPC_SERVER.to_string());
     match rpc.get_token_supply(mint) {
         Ok(_) => return true,