From fe0058e900022e3d12331ab7e99baaa2b7fcda37 Mon Sep 17 00:00:00 2001
From: Joey Kraut <joey@renegade.fi>
Date: Thu, 25 Jul 2024 14:58:04 -0700
Subject: [PATCH] funds-manager: fee-indexer: Fix misc bugs with fee indexing +
 redemption

---
 .../src/fee_indexer/queries.rs                | 20 ++++++++++++++++
 .../src/fee_indexer/redeem_fees.rs            | 24 ++++++++++++-------
 .../src/relayer_client.rs                     | 14 +++++++++++
 3 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/funds-manager/funds-manager-server/src/fee_indexer/queries.rs b/funds-manager/funds-manager-server/src/fee_indexer/queries.rs
index 309dc20..9295bcb 100644
--- a/funds-manager/funds-manager-server/src/fee_indexer/queries.rs
+++ b/funds-manager/funds-manager-server/src/fee_indexer/queries.rs
@@ -13,6 +13,7 @@ use diesel::sql_types::{Array, Integer, Nullable, Numeric, Text};
 use diesel::PgArrayExpressionMethods;
 use diesel::{ExpressionMethods, QueryDsl};
 use diesel_async::RunQueryDsl;
+use renegade_common::types::wallet::WalletIdentifier;
 use renegade_constants::MAX_BALANCES;
 use tracing::warn;
 
@@ -46,6 +47,11 @@ sql_function! {
     fn coalesce<T: SingleValue>(x: Nullable<T>, y: T) -> T;
 }
 
+sql_function! {
+    /// Append an element to an array
+    fn array_append<T: SingleValue>(arr: Array<T>, elem: T) -> Array<T>;
+}
+
 // ---------------
 // | Query Types |
 // ---------------
@@ -244,4 +250,18 @@ impl Indexer {
             .map_err(|_| FundsManagerError::db("failed to insert wallet"))
             .map(|_| ())
     }
+
+    /// Add a new mint to a wallet's managed mints
+    pub(crate) async fn add_mint_to_wallet(
+        &mut self,
+        wallet_id: &WalletIdentifier,
+        mint: &str,
+    ) -> Result<(), FundsManagerError> {
+        diesel::update(wallet_table.find(wallet_id))
+            .set(managed_mints_col.eq(array_append(managed_mints_col, mint)))
+            .execute(&mut self.db_conn)
+            .await
+            .map_err(|_| FundsManagerError::db("failed to add mint to wallet"))
+            .map(|_| ())
+    }
 }
diff --git a/funds-manager/funds-manager-server/src/fee_indexer/redeem_fees.rs b/funds-manager/funds-manager-server/src/fee_indexer/redeem_fees.rs
index f4c0ac4..4c3862e 100644
--- a/funds-manager/funds-manager-server/src/fee_indexer/redeem_fees.rs
+++ b/funds-manager/funds-manager-server/src/fee_indexer/redeem_fees.rs
@@ -4,6 +4,7 @@ use std::collections::HashMap;
 use std::str::FromStr;
 
 use aws_sdk_secretsmanager::Client as SecretsManagerClient;
+use diesel::IntoSql;
 use ethers::core::rand::thread_rng;
 use ethers::signers::LocalWallet;
 use ethers::types::TxHash;
@@ -66,20 +67,25 @@ impl Indexer {
         &mut self,
         mint: &str,
     ) -> Result<WalletMetadata, FundsManagerError> {
+        // Find a wallet with an existing balance
         let maybe_wallet = self.get_wallet_for_mint(mint).await?;
-        let maybe_wallet = if maybe_wallet.is_none() {
-            self.find_wallet_with_empty_balance().await?
-        } else {
-            maybe_wallet
-        };
+        if let Some(wallet) = maybe_wallet {
+            return Ok(wallet);
+        }
 
-        match maybe_wallet {
-            Some(wallet) => Ok(wallet),
+        // Otherwise find a wallet with an empty balance slot, create a new one if no
+        // such wallet exists
+        let maybe_wallet = self.find_wallet_with_empty_balance().await?;
+        let wallet = match maybe_wallet {
+            Some(wallet) => wallet,
             None => {
                 info!("creating new wallet for {mint}");
-                self.create_new_wallet().await
+                self.create_new_wallet().await?
             },
-        }
+        };
+
+        self.add_mint_to_wallet(&wallet.id, mint).await?;
+        Ok(wallet)
     }
 
     /// Create a new wallet for managing a given mint
diff --git a/funds-manager/funds-manager-server/src/relayer_client.rs b/funds-manager/funds-manager-server/src/relayer_client.rs
index 9e229fd..31835f6 100644
--- a/funds-manager/funds-manager-server/src/relayer_client.rs
+++ b/funds-manager/funds-manager-server/src/relayer_client.rs
@@ -18,6 +18,7 @@ use renegade_api::{
             FIND_WALLET_ROUTE, GET_WALLET_ROUTE, REDEEM_NOTE_ROUTE,
         },
     },
+    types::ApiWallet,
     RENEGADE_AUTH_HEADER_NAME, RENEGADE_SIG_EXPIRATION_HEADER_NAME,
 };
 use renegade_circuit_types::keychain::SecretSigningKey;
@@ -85,6 +86,19 @@ impl RelayerClient {
     // | Wallet Methods |
     // ------------------
 
+    /// Get the wallet associated with the given wallet id
+    pub async fn get_wallet(
+        &self,
+        wallet_id: WalletIdentifier,
+        root_key: &SecretSigningKey,
+    ) -> Result<ApiWallet, FundsManagerError> {
+        let mut path = GET_WALLET_ROUTE.to_string();
+        path = path.replace(":wallet_id", &wallet_id.to_string());
+        let resp: GetWalletResponse = self.get_relayer_with_auth(&path, root_key).await?;
+
+        Ok(resp.wallet)
+    }
+
     /// Check that the relayer has a given wallet, lookup the wallet if not
     pub async fn check_wallet_indexed(
         &self,