Fix contract state and chain monitor

dlc-manager/src/chain_monitor.rs

@@ -43,8 +43,10 @@ pub(crate) enum TxType {
         is_offer: bool,
         revoked_tx_type: RevokedTxType,
     },
-    Current,
+    BufferTx,
     CollaborativeClose,
+    SplitTx,
+    SettleTx,
 }
 
 impl_dlc_writeable_enum!(TxType,;
@@ -54,10 +56,10 @@ impl_dlc_writeable_enum!(TxType,;
         (is_offer, writeable),
         (revoked_tx_type, writeable)
     });;
-    (1, Current), (2, CollaborativeClose)
+    (1, BufferTx), (2, CollaborativeClose), (3, SplitTx), (4, SettleTx)
 );
 
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq, Copy)]
 pub(crate) enum RevokedTxType {
     Buffer,
     Settle,
@@ -76,6 +78,11 @@ impl ChainMonitor {
         }
     }
 
+    /// Returns true if the monitor doesn't contain any transaction to be watched.
+    pub fn is_empty(&self) -> bool {
+        self.watched_tx.is_empty()
+    }
+
     pub(crate) fn add_tx(&mut self, txid: Txid, channel_info: ChannelInfo) {
         self.watched_tx.insert(txid, channel_info);
     }
@@ -84,6 +91,23 @@ impl ChainMonitor {
         self.watched_tx.remove(txid);
     }
 
+    pub(crate) fn cleanup_channel(&mut self, channel_id: ChannelId) {
+        let to_remove = self
+            .watched_tx
+            .iter()
+            .filter_map(|x| {
+                if x.1.channel_id == channel_id {
+                    Some(*x.0)
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>();
+        for txid in to_remove {
+            self.watched_tx.remove(&txid);
+        }
+    }
+
     pub(crate) fn process_block(
         &self,
         block: &Block,
@@ -94,12 +118,16 @@ impl ChainMonitor {
         assert_eq!(self.last_height + 1, height);
 
         for tx in &block.txdata {
-            let txid = tx.txid();
-            if self.watched_tx.contains_key(&txid) {
-                let channel_info = self
-                    .watched_tx
-                    .get(&txid)
-                    .expect("to be able to retrieve the channel info");
+            let channel_info = self.watched_tx.get(&tx.txid()).or_else(|| {
+                for txid in tx.input.iter().map(|x| &x.previous_output.txid) {
+                    let info = self.watched_tx.get(txid);
+                    if info.is_some() {
+                        return info;
+                    }
+                }
+                None
+            });
+            if let Some(channel_info) = channel_info {
                 res.push((tx.clone(), channel_info.clone()));
             }
         }

dlc-manager/src/channel/ser.rs

@@ -62,10 +62,10 @@ impl_dlc_writeable_enum!(
     (8, RenewConfirmed, {(contract_id, writeable), (offer_per_update_point, writeable), (accept_per_update_point, writeable), (buffer_transaction, writeable), (buffer_script_pubkey, writeable), (offer_buffer_adaptor_signature, {cb_writeable, write_ecdsa_adaptor_signature, read_ecdsa_adaptor_signature}), (accept_buffer_adaptor_signature, {cb_writeable, write_ecdsa_adaptor_signature, read_ecdsa_adaptor_signature}), (timeout, writeable), (own_payout, writeable), (total_collateral, writeable)}),
     (9, Closing, {(buffer_transaction, writeable), (signed_cet, writeable), (contract_id, writeable), (attestations, vec)}),
     (10, ClosedPunished, { (punishment_txid, writeable) }),
-    (11, CollaborativeCloseOffered, { (counter_payout, writeable), (offer_signature, writeable), (close_tx, writeable), (timeout, writeable) })
+    (11, CollaborativeCloseOffered, { (counter_payout, writeable), (offer_signature, writeable), (close_tx, writeable), (timeout, writeable) }),
+    (15, CounterClosing, { (buffer_txid, writeable) })
     ;;(12, Closed), (13, CounterClosed), (14, CollaborativelyClosed)
 );
 
 impl_dlc_writeable!(FailedAccept, {(temporary_channel_id, writeable), (error_message, {cb_writeable, write_string, read_string}), (accept_message, writeable), (counter_party, writeable)});
 impl_dlc_writeable!(FailedSign, {(channel_id, writeable), (error_message, {cb_writeable, write_string, read_string}), (sign_message, writeable), (counter_party, writeable)});
-

dlc-manager/src/channel/signed_channel.rs

@@ -284,6 +284,12 @@ typed_enum!(
         /// A [`SignedChannel`] is in `Closed` state when it was force closed by
         /// the local party.
         Closed,
+        /// A [`SignedChannel`] is in `CounterClosing` state when the counter party broadcasted a
+        /// buffer transaction.
+        CounterClosing {
+            /// The transaction id of the buffer transaction that was found on chain.
+            buffer_txid: Txid,
+        },
         /// A [`SignedChannel`] is in `CounterClosed` state when it was force
         /// closed by the counter party.
         CounterClosed,
@@ -322,6 +328,19 @@ impl SignedChannel {
     /// a contract is established or under establishment.
     pub fn get_contract_id(&self) -> Option<ContractId> {
         match &self.state {
+            SignedChannelState::CounterClosing { .. } => self
+                .roll_back_state
+                .as_ref()
+                .expect("To have a rollback state")
+                .get_contract_id(),
+            _ => self.state.get_contract_id(),
+        }
+    }
+}
+
+impl SignedChannelState {
+    fn get_contract_id(&self) -> Option<ContractId> {
+        match &self {
             SignedChannelState::Established {
                 signed_contract_id, ..
             } => Some(*signed_contract_id),

dlc-manager/src/channel_updater.rs

@@ -1,8 +1,9 @@
 //! # This module contains static functions to update the state of a DLC channel.
 
-use std::ops::Deref;
+use std::{ops::Deref, sync::Mutex};
 
 use crate::{
+    chain_monitor::{ChainMonitor, ChannelInfo, TxType},
     channel::{
         accepted_channel::AcceptedChannel,
         offered_channel::OfferedChannel,
@@ -364,6 +365,7 @@ pub fn verify_and_sign_accepted_channel<S: Deref>(
     accept_channel: &AcceptChannel,
     cet_nsequence: u32,
     signer: &S,
+    chain_monitor: &Mutex<ChainMonitor>,
 ) -> Result<(SignedChannel, SignedContract, SignChannel), Error>
 where
     S::Target: Signer,
@@ -376,6 +378,7 @@ where
         cet_nsequence,
         signer,
         None,
+        chain_monitor,
     )
 }
 
@@ -387,6 +390,7 @@ pub(crate) fn verify_and_sign_accepted_channel_internal<S: Deref>(
     cet_nsequence: u32,
     signer: &S,
     sub_channel_info: Option<SubChannelSignVerifyInfo>,
+    chain_monitor: &Mutex<ChainMonitor>,
 ) -> Result<(SignedChannel, SignedContract, SignChannel), Error>
 where
     S::Target: Signer,
@@ -560,6 +564,14 @@ where
         &accept_revoke_params.publish_pk.inner,
     )?;
 
+    chain_monitor.lock().unwrap().add_tx(
+        buffer_transaction.txid(),
+        ChannelInfo {
+            channel_id,
+            tx_type: TxType::BufferTx,
+        },
+    );
+
     let signed_channel = SignedChannel {
         counter_party: signed_contract
             .accepted_contract
@@ -617,6 +629,7 @@ pub fn verify_signed_channel<S: Deref>(
     accepted_contract: &AcceptedContract,
     sign_channel: &SignChannel,
     signer: &S,
+    chain_monitor: &Mutex<ChainMonitor>,
 ) -> Result<(SignedChannel, SignedContract), Error>
 where
     S::Target: Signer,
@@ -628,6 +641,7 @@ where
         sign_channel,
         signer,
         None,
+        chain_monitor,
     )
 }
 
@@ -638,6 +652,7 @@ pub(crate) fn verify_signed_channel_internal<S: Deref>(
     sign_channel: &SignChannel,
     signer: &S,
     sub_channel_info: Option<SubChannelVerifyInfo>,
+    chain_monitor: &Mutex<ChainMonitor>,
 ) -> Result<(SignedChannel, SignedContract), Error>
 where
     S::Target: Signer,
@@ -708,6 +723,14 @@ where
         accepted_contract.dlc_transactions.get_fund_output_index()
     };
 
+    chain_monitor.lock().unwrap().add_tx(
+        accepted_channel.buffer_transaction.txid(),
+        ChannelInfo {
+            channel_id: accepted_channel.channel_id,
+            tx_type: TxType::BufferTx,
+        },
+    );
+
     let signed_channel = SignedChannel {
         counter_party: signed_contract
             .accepted_contract
@@ -846,6 +869,7 @@ pub fn settle_channel_accept<S: Deref, T: Deref>(
     peer_timeout: u64,
     signer: &S,
     time: &T,
+    chain_monitor: &Mutex<ChainMonitor>,
 ) -> Result<SettleAccept, Error>
 where
     S::Target: Signer,
@@ -860,6 +884,7 @@ where
         signer,
         time,
         None,
+        chain_monitor,
     )
 }
 
@@ -872,6 +897,7 @@ pub(crate) fn settle_channel_accept_internal<S: Deref, T: Deref>(
     signer: &S,
     time: &T,
     own_settle_adaptor_sk: Option<SecretKey>,
+    chain_monitor: &Mutex<ChainMonitor>,
 ) -> Result<SettleAccept, Error>
 where
     S::Target: Signer,
@@ -939,6 +965,14 @@ where
         channel.fee_rate_per_vb,
     )?;
 
+    chain_monitor.lock().unwrap().add_tx(
+        settle_tx.txid(),
+        ChannelInfo {
+            channel_id: channel.channel_id,
+            tx_type: TxType::SettleTx,
+        },
+    );
+
     channel.state = SignedChannelState::SettledAccepted {
         counter_next_per_update_point,
         own_next_per_update_point,
@@ -971,6 +1005,7 @@ pub fn settle_channel_confirm<T: Deref, S: Deref>(
     peer_timeout: u64,
     signer: &S,
     time: &T,
+    chain_monitor: &Mutex<ChainMonitor>,
 ) -> Result<SettleConfirm, Error>
 where
     T::Target: Time,
@@ -987,6 +1022,7 @@ where
         time,
         None,
         None,
+        chain_monitor,
     )
 }
 
@@ -1001,6 +1037,7 @@ pub(crate) fn settle_channel_confirm_internal<T: Deref, S: Deref>(
     time: &T,
     own_settle_adaptor_sk: Option<SecretKey>,
     counter_settle_adaptor_pk: Option<PublicKey>,
+    chain_monitor: &Mutex<ChainMonitor>,
 ) -> Result<SettleConfirm, Error>
 where
     T::Target: Time,
@@ -1061,6 +1098,14 @@ where
         channel.fee_rate_per_vb,
     )?;
 
+    chain_monitor.lock().unwrap().add_tx(
+        settle_tx.txid(),
+        ChannelInfo {
+            channel_id: channel.channel_id,
+            tx_type: TxType::SettleTx,
+        },
+    );
+
     let per_update_seed_pk = channel.own_per_update_seed;
     let per_update_seed = signer.get_secret_key_for_pubkey(&per_update_seed_pk)?;