Don't reuse public keys, generate puzzles for all wallets (#192)

* Don't reuse public keys, generate puzzles for all wallets (100)
* Use the right start and end indeces
* Handle edge case with all used entries

Author: mariano54

src/wallet/cc_wallet/cc_wallet_puzzles.py

@@ -99,7 +99,7 @@ def get_innerpuzzle_from_puzzle(puzzle: str):
 def create_spend_for_ephemeral(parent_of_e, auditor_coin, spend_amount):
     puzstring = f"(r (r (c (q 0x{auditor_coin.name()}) (c (q {spend_amount}) (q ())))))"
     puzzle = Program(binutils.assemble(puzstring))
-    coin = Coin(parent_of_e.name(), puzzle.get_tree_hash(), 0)
+    coin = Coin(parent_of_e.name(), puzzle.get_tree_hash(), uint64(0))
     solution = Program(binutils.assemble("()"))
     coinsol = CoinSolution(coin, clvm.to_sexp_f([puzzle, solution]))
     return coinsol
@@ -109,7 +109,7 @@ def create_spend_for_ephemeral(parent_of_e, auditor_coin, spend_amount):
 def create_spend_for_auditor(parent_of_a, auditee):
     puzstring = f"(r (c (q 0x{auditee.name()}) (q ())))"
     puzzle = Program(binutils.assemble(puzstring))
-    coin = Coin(parent_of_a.name(), puzzle.get_tree_hash(), 0)
+    coin = Coin(parent_of_a.name(), puzzle.get_tree_hash(), uint64(0))
     solution = Program(binutils.assemble("()"))
     coinsol = CoinSolution(coin, clvm.to_sexp_f([puzzle, solution]))
     return coinsol
@@ -119,7 +119,7 @@ def cc_generate_eve_spend(coin: Coin, full_puzzle: Program):
     solution = cc_make_eve_solution(
         coin.parent_coin_info, coin.puzzle_hash, coin.amount
     )
-    list_of_solutions = [CoinSolution(coin, clvm.to_sexp_f([full_puzzle, solution, ]), )]
+    list_of_solutions = [CoinSolution(coin, clvm.to_sexp_f([full_puzzle, solution]),)]
     aggsig = BLSSignature.aggregate([])
     spend_bundle = SpendBundle(list_of_solutions, aggsig)
     return spend_bundle

src/wallet/trade_manager.py

@@ -59,9 +59,9 @@ async def create_offer_for_ids(
                             to_exclude: List[Coin] = []
                         else:
                             to_exclude = spend_bundle.removals()
-                        zero_spend_bundle: Optional[SpendBundle] = await wallet.generate_zero_val_coin(
-                            False, to_exclude
-                        )
+                        zero_spend_bundle: Optional[
+                            SpendBundle
+                        ] = await wallet.generate_zero_val_coin(False, to_exclude)
 
                         if zero_spend_bundle is None:
                             raise ValueError(
@@ -290,7 +290,7 @@ async def respond_to_offer(self, file_path: Path) -> bool:
                 )
             else:
                 if chia_spend_bundle is None:
-                    to_exclude = []
+                    to_exclude: List = []
                 else:
                     to_exclude = chia_spend_bundle.removals()
                 zero_spend_bundle: SpendBundle = await wallets[

src/wallet/wallet_puzzle_store.py

@@ -293,20 +293,3 @@ async def get_unused_derivation_path(self) -> Optional[uint32]:
             return uint32(row[0])
 
         return None
-
-    async def get_unused_derivation_path_for_wallet(
-        self, wallet_id: int
-    ) -> Optional[uint32]:
-        """
-        Returns the first unused derivation path by derivation_index.
-        """
-        cursor = await self.db_connection.execute(
-            f"SELECT MIN(derivation_index) FROM derivation_paths WHERE used=0 and wallet_id={wallet_id};"
-        )
-        row = await cursor.fetchone()
-        await cursor.close()
-
-        if row is not None and row[0] is not None:
-            return uint32(row[0])
-
-        return None

src/wallet/wallet_state_manager.py

@@ -207,48 +207,41 @@ async def get_keys(
         private = self.private_key.private_child(index_for_puzzlehash).get_private_key()
         return pubkey, private
 
-    async def create_more_puzzle_hashes(
-        self, from_zero: bool = False, for_wallet: Any = None
-    ):
+    async def create_more_puzzle_hashes(self, from_zero: bool = False):
         """
         For all wallets in the user store, generates the first few puzzle hashes so
         that we can restore the wallet from only the private keys.
         """
+        targets = list(self.wallets.keys())
 
-        if for_wallet is None:
-            targets = list(self.wallets.keys())
-        else:
-            targets = [for_wallet]
+        unused: Optional[uint32] = await self.puzzle_store.get_unused_derivation_path()
+        if unused is None:
+            # This handles the case where the database has entries but they have all been used
+            unused = await self.puzzle_store.get_last_derivation_path()
+            if unused is None:
+                # This handles the case where the database is empty
+                unused = uint32(0)
+
+        to_generate = 100
 
         for wallet_id in targets:
             target_wallet = self.wallets[wallet_id]
-            unused: Optional[
-                uint32
-            ] = await self.puzzle_store.get_unused_derivation_path_for_wallet(wallet_id)
+
             last: Optional[
                 uint32
             ] = await self.puzzle_store.get_last_derivation_path_for_wallet(wallet_id)
 
-            if target_wallet.wallet_info.type == WalletType.COLOURED_COIN:
-                to_generate = 100
-            else:
-                to_generate = 500
             start_index = 0
             derivation_paths: List[DerivationRecord] = []
 
-            if last is None:
-                assert unused is None
-            if unused is not None:
-                assert last is not None
+            if last is not None:
                 start_index = last + 1
-                to_generate -= last - unused
 
             # If the key was replaced (from_zero=True), we should generate the puzzle hashes for the new key
-            end = start_index + to_generate
             if from_zero:
                 start_index = 0
 
-            for index in range(start_index, end):
+            for index in range(start_index, unused + to_generate):
                 pubkey: PublicKey = self.get_public_key(uint32(index))
                 puzzle: Program = target_wallet.puzzle_for_pk(bytes(pubkey))
                 puzzlehash: bytes32 = puzzle.get_tree_hash()
@@ -266,8 +259,8 @@ async def create_more_puzzle_hashes(
                 )
 
             await self.puzzle_store.add_derivation_paths(derivation_paths)
-            if from_zero and unused is not None and unused > 0:
-                await self.puzzle_store.set_used_up_to(uint32(unused - 1))
+        if unused > 0:
+            await self.puzzle_store.set_used_up_to(uint32(unused - 1))
 
     async def get_unused_derivation_record(self, wallet_id: uint32) -> DerivationRecord:
         """
@@ -279,9 +272,9 @@ async def get_unused_derivation_record(self, wallet_id: uint32) -> DerivationRec
             # If we have no unused public keys, we will create new ones
             unused: Optional[
                 uint32
-            ] = await self.puzzle_store.get_unused_derivation_path_for_wallet(wallet_id)
+            ] = await self.puzzle_store.get_unused_derivation_path()
             if unused is None:
-                await self.create_more_puzzle_hashes(for_wallet=wallet_id)
+                await self.create_more_puzzle_hashes()
 
             # Now we must have unused public keys
             unused = await self.puzzle_store.get_unused_derivation_path()
@@ -1336,7 +1329,7 @@ async def get_wallet_for_colour(self, colour):
 
     async def add_new_wallet(self, wallet: Any, id: int):
         self.wallets[uint32(id)] = wallet
-        await self.create_more_puzzle_hashes(for_wallet=id)
+        await self.create_more_puzzle_hashes()
 
     async def get_coin_records_by_spent(self, spent: bool):
         return await self.wallet_store.get_coin_records_by_spent(spent)

src/wallet/websocket_server.py

@@ -26,6 +26,7 @@
 from src.simulator.simulator_protocol import FarmNewBlockProtocol
 from src.util.config import load_config_cli, load_config
 from src.util.ints import uint64
+from src.types.sized_bytes import bytes32
 from src.util.logging import initialize_logging
 from src.wallet.util.wallet_types import WalletType
 from src.wallet.rl_wallet.rl_wallet import RLWallet
@@ -224,7 +225,7 @@ async def create_new_wallet(self, websocket, request, response_api):
                 response = {"success": True, "type": cc_wallet.wallet_info.type.name}
                 return await websocket.send(format_response(response_api, response))
             elif request["mode"] == "existing":
-                cc_wallet: CCWallet = await CCWallet.create_wallet_for_cc(
+                cc_wallet = await CCWallet.create_wallet_for_cc(
                     wallet_state_manager, main_wallet, request["colour"]
                 )
                 response = {"success": True, "type": cc_wallet.wallet_info.type.name}
@@ -405,7 +406,7 @@ async def cc_get_new_innerpuzzlehash(self, websocket, request, response_api):
 
         wallet_id = int(request["wallet_id"])
         wallet: CCWallet = self.wallet_node.wallet_state_manager.wallets[wallet_id]
-        innerpuz: str = await wallet.get_new_inner_hash()
+        innerpuz: bytes32 = await wallet.get_new_inner_hash()
         response = {"innerpuz": innerpuz.hex()}
         return await websocket.send(format_response(response_api, response))
 

tests/cc_wallet/test_cc_wallet.py

@@ -133,6 +133,7 @@ async def test_cc_spend(self, two_wallet_nodes):
         await self.time_out_assert(15, cc_wallet.get_confirmed_balance, 100)
         await self.time_out_assert(15, cc_wallet.get_unconfirmed_balance, 100)
 
+        assert cc_wallet.cc_info.my_core
         colour = cc_wallet_puzzles.get_genesis_from_core(cc_wallet.cc_info.my_core)
 
         cc_wallet_2: CCWallet = await CCWallet.create_wallet_for_cc(
@@ -232,10 +233,10 @@ async def test_generate_zero_val(self, two_wallet_nodes):
         for i in range(1, num_blocks):
             await full_node_1.farm_new_block(FarmNewBlockProtocol(ph))
 
-
         await self.time_out_assert(15, cc_wallet.get_confirmed_balance, 100)
         await self.time_out_assert(15, cc_wallet.get_unconfirmed_balance, 100)
 
+        assert cc_wallet.cc_info.my_core
         colour = cc_wallet_puzzles.get_genesis_from_core(cc_wallet.cc_info.my_core)
 
         cc_wallet_2: CCWallet = await CCWallet.create_wallet_for_cc(
@@ -295,6 +296,7 @@ async def test_cc_trade(self, two_wallet_nodes):
         await self.time_out_assert(15, cc_wallet.get_confirmed_balance, 100)
         await self.time_out_assert(15, cc_wallet.get_unconfirmed_balance, 100)
 
+        assert cc_wallet.cc_info.my_core
         colour = cc_wallet_puzzles.get_genesis_from_core(cc_wallet.cc_info.my_core)
 
         cc_wallet_2: CCWallet = await CCWallet.create_wallet_for_cc(
@@ -382,6 +384,7 @@ async def test_cc_spend_uncoloured(self, two_wallet_nodes):
         await self.time_out_assert(15, cc_wallet.get_confirmed_balance, 100)
         await self.time_out_assert(15, cc_wallet.get_unconfirmed_balance, 100)
 
+        assert cc_wallet.cc_info.my_core
         colour = cc_wallet_puzzles.get_genesis_from_core(cc_wallet.cc_info.my_core)
 
         cc_wallet_2: CCWallet = await CCWallet.create_wallet_for_cc(
@@ -396,7 +399,6 @@ async def test_cc_spend_uncoloured(self, two_wallet_nodes):
         for i in range(1, num_blocks):
             await full_node_1.farm_new_block(FarmNewBlockProtocol(ph))
 
-
         await self.time_out_assert(15, cc_wallet.get_confirmed_balance, 40)
         await self.time_out_assert(15, cc_wallet.get_unconfirmed_balance, 40)
 
@@ -412,7 +414,6 @@ async def test_cc_spend_uncoloured(self, two_wallet_nodes):
         for i in range(0, num_blocks):
             await full_node_1.farm_new_block(FarmNewBlockProtocol(token_bytes()))
 
-
         id = cc_wallet_2.wallet_info.id
         wsm = cc_wallet_2.wallet_state_manager
         await self.time_out_assert(15, wsm.get_confirmed_balance_for_wallet, 70, id)
@@ -457,6 +458,7 @@ async def test_cc_trade_with_multiple_colours(self, two_wallet_nodes):
         await self.time_out_assert(15, red_wallet.get_confirmed_balance, 100)
         await self.time_out_assert(15, red_wallet.get_unconfirmed_balance, 100)
 
+        assert red_wallet.cc_info.my_core
         red = cc_wallet_puzzles.get_genesis_from_core(red_wallet.cc_info.my_core)
 
         await full_node_1.farm_new_block(FarmNewBlockProtocol(ph2))
@@ -469,6 +471,7 @@ async def test_cc_trade_with_multiple_colours(self, two_wallet_nodes):
         for i in range(1, num_blocks):
             await full_node_1.farm_new_block(FarmNewBlockProtocol(ph))
 
+        assert blue_wallet_2.cc_info.my_core
         blue = cc_wallet_puzzles.get_genesis_from_core(blue_wallet_2.cc_info.my_core)
 
         red_wallet_2: CCWallet = await CCWallet.create_wallet_for_cc(
@@ -570,11 +573,10 @@ async def test_create_offer_with_zero_val(self, two_wallet_nodes):
         for i in range(1, num_blocks):
             await full_node_1.farm_new_block(FarmNewBlockProtocol(ph))
 
-
         await self.time_out_assert(15, cc_wallet.get_unconfirmed_balance, 100)
         await self.time_out_assert(15, cc_wallet.get_confirmed_balance, 100)
 
-
+        assert cc_wallet.cc_info.my_core
         colour = cc_wallet_puzzles.get_genesis_from_core(cc_wallet.cc_info.my_core)
 
         cc_wallet_2: CCWallet = await CCWallet.create_wallet_for_cc(
@@ -622,6 +624,5 @@ async def test_create_offer_with_zero_val(self, two_wallet_nodes):
         for i in range(0, num_blocks):
             await full_node_1.farm_new_block(FarmNewBlockProtocol(token_bytes()))
 
-
         await self.time_out_assert(15, cc_wallet_2.get_confirmed_balance, 30)
         await self.time_out_assert(15, cc_wallet_2.get_confirmed_balance, 30)