From 4c5ab0da80ab6ca2182425b368ed009646f4abfd Mon Sep 17 00:00:00 2001
From: ghassmo <ghxxm@protonmail.com>
Date: Fri, 15 Oct 2021 11:13:49 +0300
Subject: [PATCH] cashierdb: create MainTokenKey struct

---
 src/bin/cashierd.rs     | 20 +++++++++++++-------
 src/wallet/cashierdb.rs | 38 ++++++++++++++++++++++++--------------
 2 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/src/bin/cashierd.rs b/src/bin/cashierd.rs
index c8fa5b84c..a99b23e44 100644
--- a/src/bin/cashierd.rs
+++ b/src/bin/cashierd.rs
@@ -25,7 +25,7 @@ use drk::{
     serial::{deserialize, serialize},
     service::{bridge, bridge::Bridge},
     util::{expand_path, generate_id, join_config_path, parse::truncate, NetworkName},
-    wallet::{CashierDb, WalletDb},
+    wallet::{cashierdb::MainTokenKey, CashierDb, WalletDb},
     Error, Result,
 };
 
@@ -482,12 +482,15 @@ impl Cashierd {
                         if main_keypairs.is_empty() {
                             main_keypair = Keypair::new();
                             self.cashier_wallet.put_main_keys(
-                                &serialize(&main_keypair),
-                                &serialize(&main_keypair.pubkey()),
+                                &MainTokenKey {
+                                    private_key: serialize(&main_keypair),
+                                    public_key: serialize(&main_keypair.pubkey()),
+                                },
                                 &NetworkName::Solana,
                             )?;
                         } else {
-                            main_keypair = deserialize(&main_keypairs[main_keypairs.len() - 1].0)?;
+                            main_keypair =
+                                deserialize(&main_keypairs[main_keypairs.len() - 1].private_key)?;
                         }
                     } else {
                         let keypair_str = drk::cli::cli_config::load_keypair_to_str(
@@ -518,12 +521,15 @@ impl Cashierd {
                         if main_keypairs.is_empty() {
                             main_keypair = Keypair::new();
                             self.cashier_wallet.put_main_keys(
-                                &serialize(&main_keypair),
-                                &serialize(&main_keypair.pubkey()),
+                                &MainTokenKey {
+                                    private_key: serialize(&main_keypair),
+                                    public_key: serialize(&main_keypair.pubkey()),
+                                },
                                 &NetworkName::Bitcoin,
                             )?;
                         } else {
-                            main_keypair = deserialize(&main_keypairs[main_keypairs.len() - 1].0)?;
+                            main_keypair =
+                                deserialize(&main_keypairs[main_keypairs.len() - 1].private_key)?;
                         }
                     } else {
                         let keypair_str = drk::cli::cli_config::load_keypair_to_str(
diff --git a/src/wallet/cashierdb.rs b/src/wallet/cashierdb.rs
index 78348de72..803452f67 100644
--- a/src/wallet/cashierdb.rs
+++ b/src/wallet/cashierdb.rs
@@ -17,6 +17,11 @@ pub struct CashierDb {
     pub initialized: Mutex<bool>,
 }
 
+pub struct MainTokenKey {
+    pub public_key: Vec<u8>,
+    pub private_key: Vec<u8>,
+}
+
 impl WalletApi for CashierDb {
     fn get_password(&self) -> String {
         self.password.to_owned()
@@ -59,12 +64,7 @@ impl CashierDb {
         Ok(())
     }
 
-    pub fn put_main_keys(
-        &self,
-        token_key_private: &[u8],
-        token_key_public: &[u8],
-        network: &NetworkName,
-    ) -> Result<()> {
+    pub fn put_main_keys(&self, token_key: &MainTokenKey, network: &NetworkName) -> Result<()> {
         debug!(target: "CASHIERDB", "Put main keys");
 
         // open connection
@@ -80,15 +80,15 @@ impl CashierDb {
             VALUES
             (:token_key_private, :token_key_public, :network)",
             named_params! {
-                ":token_key_private": token_key_private,
-                ":token_key_public": token_key_public,
+                ":token_key_private": token_key.private_key,
+                ":token_key_public": token_key.public_key,
                 ":network": &network,
             },
         )?;
         Ok(())
     }
 
-    pub fn get_main_keys(&self, network: &NetworkName) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
+    pub fn get_main_keys(&self, network: &NetworkName) -> Result<Vec<MainTokenKey>> {
         debug!(target: "CASHIERDB", "Get main keys");
         // open connection
         let conn = Connection::open(&self.path)?;
@@ -110,7 +110,11 @@ impl CashierDb {
         let mut keys = vec![];
 
         for k in keys_iter {
-            keys.push(k?);
+            let k = k?;
+            keys.push(MainTokenKey {
+                private_key: k.0,
+                public_key: k.1,
+            });
         }
 
         Ok(keys)
@@ -488,14 +492,20 @@ mod tests {
 
         let network = NetworkName::Bitcoin;
 
-        wallet.put_main_keys(&token_addr_private, &token_addr, &network)?;
+        wallet.put_main_keys(
+            &MainTokenKey {
+                private_key: token_addr_private.clone(),
+                public_key: token_addr.clone(),
+            },
+            &network,
+        )?;
 
         let keys = wallet.get_main_keys(&network)?;
 
         assert_eq!(keys.len(), 1);
 
-        assert_eq!(keys[0].0, token_addr_private);
-        assert_eq!(keys[0].1, token_addr);
+        assert_eq!(keys[0].private_key, token_addr_private);
+        assert_eq!(keys[0].public_key, token_addr);
 
         std::fs::remove_file(walletdb_path)?;
 
@@ -541,7 +551,7 @@ mod tests {
         assert_eq!(resumed_keys[0].1, token_addr_private);
         assert_eq!(resumed_keys[0].2, token_addr);
         assert_eq!(resumed_keys[0].3, token_id);
-        
+
         wallet.confirm_deposit_key_record(&public2, &network)?;
 
         let keys = wallet.get_deposit_token_keys_by_dkey_public(&public2, &network)?;