Merge pull request #74 from bennyhodl/contract-balances

contract balances

ddk-manager/src/contract/accepted_contract.rs

@@ -109,7 +109,7 @@ mod tests {
 
     #[test]
     fn pnl_compute_test() {
-        let buf = include_bytes!("../../../ddk/tests/data/dlc_storage/sled/Accepted");
+        let buf = include_bytes!("../../../ddk/tests/data/dlc_storage/Accepted");
         let accepted_contract: AcceptedContract = Readable::read(&mut Cursor::new(&buf)).unwrap();
         let cets = &accepted_contract.dlc_transactions.cets;
         assert_eq!(accepted_contract.compute_pnl(&cets[0]), 90000000);

ddk-node/src/cli_opts.rs

@@ -12,6 +12,8 @@ pub enum CliCommand {
     AcceptOffer(Accept),
     /// List contracts.
     Contracts,
+    #[command(about = "Get the wallet balance.")]
+    Balance,
     /// Wallet commands
     #[clap(subcommand)]
     Wallet(WalletCommand),
@@ -38,8 +40,6 @@ pub struct Offer {
 
 #[derive(Clone, Debug, Subcommand)]
 pub enum WalletCommand {
-    #[command(about = "Get the wallet balance.")]
-    Balance,
     #[command(about = "Generate a new, unused address from the wallet.")]
     NewAddress,
     #[command(about = "Get the wallet transactions.")]

ddk-node/src/command.rs

@@ -215,15 +215,15 @@ pub async fn cli_command(
                 .collect::<Vec<Value>>();
             print!("{}", serde_json::to_string_pretty(&contract_values)?)
         }
+        CliCommand::Balance => {
+            let balance = client
+                .wallet_balance(WalletBalanceRequest::default())
+                .await?
+                .into_inner();
+            let pretty_string = serde_json::to_string_pretty(&balance)?;
+            println!("{}", pretty_string);
+        }
         CliCommand::Wallet(wallet) => match wallet {
-            WalletCommand::Balance => {
-                let balance = client
-                    .wallet_balance(WalletBalanceRequest::default())
-                    .await?
-                    .into_inner();
-                let pretty_string = serde_json::to_string_pretty(&balance)?;
-                println!("{}", pretty_string);
-            }
             WalletCommand::NewAddress => {
                 let address = client
                     .new_address(NewAddressRequest::default())

ddk-node/src/ddkrpc.rs

@@ -81,7 +81,11 @@ pub struct WalletBalanceResponse {
     #[prost(uint64, tag = "1")]
     pub confirmed: u64,
     #[prost(uint64, tag = "2")]
-    pub unconfirmed: u64,
+    pub foreign_unconfirmed: u64,
+    #[prost(uint64, tag = "3")]
+    pub change_unconfirmed: u64,
+    #[prost(int64, tag = "4")]
+    pub contract_balance: i64,
 }
 #[derive(serde::Serialize, serde::Deserialize)]
 #[allow(clippy::derive_partial_eq_without_eq)]

ddk-node/src/lib.rs

@@ -217,12 +217,13 @@ impl DdkRpc for DdkNode {
         _request: Request<WalletBalanceRequest>,
     ) -> Result<Response<WalletBalanceResponse>, Status> {
         tracing::info!("Request for wallet balance.");
-        let wallet_balance = self.node.wallet.get_balance().unwrap();
+        let wallet_balance = self.node.balance().unwrap();
 
         let response = WalletBalanceResponse {
             confirmed: wallet_balance.confirmed.to_sat(),
-            unconfirmed: (wallet_balance.trusted_pending + wallet_balance.untrusted_pending)
-                .to_sat(),
+            foreign_unconfirmed: wallet_balance.foreign_unconfirmed.to_sat(),
+            change_unconfirmed: wallet_balance.change_unconfirmed.to_sat(),
+            contract_balance: wallet_balance.contract_pnl,
         };
         Ok(Response::new(response))
     }

ddk-node/src/proto/ddkrpc.proto

@@ -61,7 +61,9 @@ message WalletBalanceRequest {}
 
 message WalletBalanceResponse {
   uint64 confirmed = 1;
-  uint64 unconfirmed = 2;
+  uint64 foreign_unconfirmed = 2;
+  uint64 change_unconfirmed = 3;
+  int64 contract_balance = 4;
 }
 
 message GetWalletTransactionsRequest {}

ddk/src/ddk.rs

@@ -5,8 +5,9 @@ use crate::{nostr::marketplace::*, DEFAULT_NOSTR_RELAY};
 use crate::{Oracle, Storage, Transport};
 use anyhow::anyhow;
 use bitcoin::secp256k1::PublicKey;
-use bitcoin::Network;
+use bitcoin::{Amount, Network};
 use crossbeam::channel::{unbounded, Receiver, Sender};
+use ddk_manager::contract::Contract;
 use ddk_manager::error::Error;
 use ddk_manager::{
     contract::contract_input::ContractInput, CachedContractSignerProvider, ContractId,
@@ -241,4 +242,50 @@ where
 
         Ok((contract_id, counter_party, accept_dlc))
     }
+
+    pub fn balance(&self) -> anyhow::Result<crate::Balance> {
+        let wallet_balance = self.wallet.get_balance()?;
+        let contracts = self.storage.get_contracts()?;
+
+        let contract = &contracts
+            .iter()
+            .map(|contract| match contract {
+                Contract::Confirmed(c) => {
+                    let accept_party_collateral = c.accepted_contract.accept_params.collateral;
+                    if c.accepted_contract.offered_contract.is_offer_party {
+                        Amount::from_sat(
+                            c.accepted_contract.offered_contract.total_collateral
+                                - accept_party_collateral,
+                        )
+                    } else {
+                        Amount::from_sat(c.accepted_contract.accept_params.collateral)
+                    }
+                }
+                _ => Amount::ZERO,
+            })
+            .sum::<Amount>();
+
+        let contract_pnl = &contracts
+            .iter()
+            .map(|contract| match contract {
+                Contract::Closed(c) => {
+                    println!("Closed contract: {}", c.pnl);
+                    0_i64
+                }
+                Contract::PreClosed(p) => p
+                    .signed_contract
+                    .accepted_contract
+                    .compute_pnl(&p.signed_cet),
+                _ => 0_i64,
+            })
+            .sum::<i64>();
+
+        Ok(crate::Balance {
+            confirmed: wallet_balance.confirmed,
+            change_unconfirmed: wallet_balance.immature + wallet_balance.trusted_pending,
+            foreign_unconfirmed: wallet_balance.untrusted_pending,
+            contract: contract.to_owned(),
+            contract_pnl: contract_pnl.to_owned(),
+        })
+    }
 }

ddk/src/lib.rs

@@ -23,16 +23,18 @@ pub mod util;
 pub mod wallet;
 use std::sync::Arc;
 
-use bdk_wallet::ChangeSet;
 /// DDK object with all services
 pub use ddk::DlcDevKit;
 pub use ddk::DlcManagerMessage;
+pub use ddk_manager;
 
 /// Default nostr relay.
 pub const DEFAULT_NOSTR_RELAY: &str = "wss://nostr.dlcdevkit.com";
 
 use async_trait::async_trait;
+use bdk_wallet::ChangeSet;
 use bitcoin::secp256k1::{PublicKey, SecretKey};
+use bitcoin::Amount;
 use ddk::DlcDevKitDlcManager;
 use dlc_messages::oracle_msgs::OracleAnnouncement;
 use dlc_messages::Message;
@@ -85,3 +87,17 @@ pub trait KeyStorage {
 pub trait Oracle: ddk_manager::Oracle + Send + Sync + 'static {
     fn name(&self) -> String;
 }
+
+#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)]
+pub struct Balance {
+    /// Confirmed in the wallet.
+    pub confirmed: Amount,
+    /// Unconfirmed UTXO that is owned by the wallet. Typically change.
+    pub change_unconfirmed: Amount,
+    /// Unconfirmed UTXO not owned by the wallet.
+    pub foreign_unconfirmed: Amount,
+    /// UTXOs in an active contract.
+    pub contract: Amount,
+    /// Profit and loss in all closed contracts.
+    pub contract_pnl: i64,
+}

ddk/src/storage/sled/contract.rs

@@ -400,11 +400,7 @@ mod tests {
         ($name: ident, $body: expr) => {
             #[test]
             fn $name() {
-                let path = format!(
-                    "{}{}",
-                    "tests/data/dlc_storage/sleddb/",
-                    std::stringify!($name)
-                );
+                let path = format!("{}{}", "tests/data/dlc_storagedb/", std::stringify!($name));
                 {
                     let storage = SledStorage::new(&path).expect("Error opening sled DB");
                     #[allow(clippy::redundant_closure_call)]
@@ -424,7 +420,7 @@ mod tests {
     }
 
     sled_test!(create_contract_can_be_retrieved, |storage: SledStorage| {
-        let serialized = include_bytes!("../../../tests/data/dlc_storage/sled/Offered");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/Offered");
         let contract = deserialize_object(serialized);
 
         storage
@@ -443,9 +439,9 @@ mod tests {
     });
 
     sled_test!(update_contract_is_updated, |storage: SledStorage| {
-        let serialized = include_bytes!("../../../tests/data/dlc_storage/sled/Offered");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/Offered");
         let offered_contract = deserialize_object(serialized);
-        let serialized = include_bytes!("../../../tests/data/dlc_storage/sled/Accepted");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/Accepted");
         let accepted_contract = deserialize_object(serialized);
         let accepted_contract = Contract::Accepted(accepted_contract);
 
@@ -467,7 +463,7 @@ mod tests {
     });
 
     sled_test!(delete_contract_is_deleted, |storage: SledStorage| {
-        let serialized = include_bytes!("../../../tests/data/dlc_storage/sled/Offered");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/Offered");
         let contract = deserialize_object(serialized);
         storage
             .create_contract(&contract)
@@ -484,45 +480,45 @@ mod tests {
     });
 
     fn insert_offered_signed_and_confirmed(storage: &mut SledStorage) {
-        let serialized = include_bytes!("../../../tests/data/dlc_storage/sled/Offered");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/Offered");
         let offered_contract = deserialize_object(serialized);
         storage
             .create_contract(&offered_contract)
             .expect("Error creating contract");
 
-        let serialized = include_bytes!("../../../tests/data/dlc_storage/sled/Signed");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/Signed");
         let signed_contract = Contract::Signed(deserialize_object(serialized));
         storage
             .update_contract(&signed_contract)
             .expect("Error creating contract");
-        let serialized = include_bytes!("../../../tests/data/dlc_storage/sled/Signed1");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/Signed1");
         let signed_contract = Contract::Signed(deserialize_object(serialized));
         storage
             .update_contract(&signed_contract)
             .expect("Error creating contract");
 
-        let serialized = include_bytes!("../../../tests/data/dlc_storage/sled/Confirmed");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/Confirmed");
         let confirmed_contract = Contract::Confirmed(deserialize_object(serialized));
         storage
             .update_contract(&confirmed_contract)
             .expect("Error creating contract");
-        let serialized = include_bytes!("../../../tests/data/dlc_storage/sled/Confirmed1");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/Confirmed1");
         let confirmed_contract = Contract::Confirmed(deserialize_object(serialized));
         storage
             .update_contract(&confirmed_contract)
             .expect("Error creating contract");
 
-        let serialized = include_bytes!("../../../tests/data/dlc_storage/sled/PreClosed");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/PreClosed");
         let preclosed_contract = Contract::PreClosed(deserialize_object(serialized));
         storage
             .update_contract(&preclosed_contract)
             .expect("Error creating contract");
     }
 
     fn insert_offered_and_signed_channels(storage: &mut SledStorage) {
-        let serialized = include_bytes!("../../../tests/data/dlc_storage/sled/Offered");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/Offered");
         let offered_contract = deserialize_object(serialized);
-        let serialized = include_bytes!("../../../tests/data/dlc_storage/sled/OfferedChannel");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/OfferedChannel");
         let offered_channel = deserialize_object(serialized);
         storage
             .upsert_channel(
@@ -531,15 +527,13 @@ mod tests {
             )
             .expect("Error creating contract");
 
-        let serialized =
-            include_bytes!("../../../tests/data/dlc_storage/sled/SignedChannelEstablished");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/SignedChannelEstablished");
         let signed_channel = Channel::Signed(deserialize_object(serialized));
         storage
             .upsert_channel(signed_channel, None)
             .expect("Error creating contract");
 
-        let serialized =
-            include_bytes!("../../../tests/data/dlc_storage/sled/SignedChannelSettled");
+        let serialized = include_bytes!("../../../tests/data/dlc_storage/SignedChannelSettled");
         let signed_channel = Channel::Signed(deserialize_object(serialized));
         storage
             .upsert_channel(signed_channel, None)

ddk/tests/balance.rs

@@ -0,0 +1,60 @@
+mod test_util;
+
+use std::sync::Arc;
+
+use bitcoin::key::Secp256k1;
+use bitcoin::Amount;
+use bitcoincore_rpc::RpcApi;
+use ddk::oracle::memory::MemoryOracle;
+use ddk_manager::contract::Contract;
+use ddk_manager::contract::{ser::Serializable, PreClosedContract};
+use ddk_manager::Storage;
+use lightning::io::Cursor;
+use test_util::generate_blocks;
+
+#[tokio::test]
+async fn contract_balance() {
+    let contract_bytes = include_bytes!("../tests/data/dlc_storage/PreClosed");
+    let mut cursor = Cursor::new(contract_bytes);
+    let preclosed = PreClosedContract::deserialize(&mut cursor).unwrap();
+    let secp = Secp256k1::new();
+    let oracle = Arc::new(MemoryOracle::default());
+    let bob = test_util::TestSuite::new(&secp, "balance", oracle).await;
+
+    bob.ddk
+        .storage
+        .update_contract(&Contract::PreClosed(preclosed.clone()))
+        .unwrap();
+
+    let address = bob.ddk.wallet.new_external_address().unwrap().address;
+
+    let auth = bitcoincore_rpc::Auth::UserPass("ddk".to_string(), "ddk".to_string());
+    let client = bitcoincore_rpc::Client::new("http://127.0.0.1:18443", auth).unwrap();
+    client
+        .send_to_address(
+            &address,
+            Amount::ONE_BTC,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+        .unwrap();
+    tokio::time::sleep(std::time::Duration::from_secs(3)).await;
+
+    bob.ddk.wallet.sync().unwrap();
+    let balance = bob.ddk.balance().unwrap();
+    assert_eq!(balance.foreign_unconfirmed, Amount::ONE_BTC);
+    assert_eq!(balance.contract_pnl, -11000000);
+
+    generate_blocks(2);
+
+    tokio::time::sleep(std::time::Duration::from_secs(2)).await;
+    bob.ddk.wallet.sync().unwrap();
+    let balance = bob.ddk.balance().unwrap();
+    assert_eq!(balance.confirmed, Amount::ONE_BTC);
+    assert_eq!(balance.foreign_unconfirmed, Amount::ZERO);
+    assert_eq!(balance.contract_pnl, -11000000);
+}