fedimint_gwv2_client/

receive_sm.rs

use core::fmt;
use std::collections::BTreeMap;
use std::sync::Arc;

use fedimint_api_client::api::{FederationApiExt, ServerError};
use fedimint_api_client::query::FilterMapThreshold;
use fedimint_client_module::DynGlobalClientContext;
use fedimint_client_module::sm::{ClientSMDatabaseTransaction, State, StateTransition};
use fedimint_client_module::transaction::{ClientInput, ClientInputBundle};
use fedimint_core::core::OperationId;
use fedimint_core::encoding::{Decodable, Encodable};
use fedimint_core::module::{Amounts, ApiRequestErased};
use fedimint_core::secp256k1::Keypair;
use fedimint_core::util::FmtCompact;
use fedimint_core::{NumPeersExt, OutPoint, PeerId, runtime};
use fedimint_lnv2_common::contracts::IncomingContract;
use fedimint_lnv2_common::endpoint_constants::DECRYPTION_KEY_SHARE_ENDPOINT;
use fedimint_lnv2_common::{LightningInput, LightningInputV0};
use fedimint_logging::LOG_CLIENT_MODULE_GW;
use serde::{Deserialize, Serialize};
use tpe::{
    AggregateDecryptionKey, AggregatePublicKey, DecryptionKeyShare, PublicKeyShare,
    aggregate_dk_shares,
};
use tracing::warn;

use super::events::{IncomingPaymentFailed, IncomingPaymentSucceeded};
use crate::GatewayClientContextV2;

#[derive(Debug, Clone, Eq, PartialEq, Hash, Decodable, Encodable)]
pub struct ReceiveStateMachine {
    pub common: ReceiveSMCommon,
    pub state: ReceiveSMState,
}

impl ReceiveStateMachine {
    pub fn update(&self, state: ReceiveSMState) -> Self {
        Self {
            common: self.common.clone(),
            state,
        }
    }
}

impl fmt::Display for ReceiveStateMachine {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Receive State Machine Operation ID: {:?} State: {}",
            self.common.operation_id, self.state
        )
    }
}

#[derive(Debug, Clone, Eq, PartialEq, Hash, Decodable, Encodable)]
pub struct ReceiveSMCommon {
    pub operation_id: OperationId,
    pub contract: IncomingContract,
    pub outpoint: OutPoint,
    pub refund_keypair: Keypair,
}

#[derive(Debug, Clone, Eq, PartialEq, Hash, Decodable, Encodable)]
pub enum ReceiveSMState {
    Funding,
    Rejected(String),
    Success([u8; 32]),
    Failure,
    Refunding(Vec<OutPoint>),
}

impl fmt::Display for ReceiveSMState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ReceiveSMState::Funding => write!(f, "Funding"),
            ReceiveSMState::Rejected(_) => write!(f, "Rejected"),
            ReceiveSMState::Success(_) => write!(f, "Success"),
            ReceiveSMState::Failure => write!(f, "Failure"),
            ReceiveSMState::Refunding(_) => write!(f, "Refunding"),
        }
    }
}

/// What the funding trigger resolves to.
///
/// The shares are verified and aggregated in the trigger, off the async
/// workers, so the transition only records the outcome.
#[derive(Debug, Serialize, Deserialize)]
pub enum DecryptionOutcome {
    /// The funding transaction was rejected.
    Rejected(String),
    /// The shares aggregate to a key the config's aggregate key does not
    /// verify, so the config's public keys are inconsistent.
    InconsistentKeys,
    /// The aggregate decryption key, and the preimage if the ciphertext holds
    /// one that matches the payment image.
    Decrypted(AggregateDecryptionKey, Option<[u8; 32]>),
}

#[cfg_attr(doc, aquamarine::aquamarine)]
/// State machine that handles the relay of an incoming Lightning payment.
///
/// ```mermaid
/// graph LR
/// classDef virtual fill:#fff,stroke-dasharray: 5 5
///
///     Funding -- funding transaction is rejected --> Rejected
///     Funding -- aggregated decryption key is invalid --> Failure
///     Funding -- decrypted preimage is valid --> Success
///     Funding -- decrypted preimage is invalid --> Refunding
/// ```
impl State for ReceiveStateMachine {
    type ModuleContext = GatewayClientContextV2;

    fn transitions(
        &self,
        context: &Self::ModuleContext,
        global_context: &DynGlobalClientContext,
    ) -> Vec<StateTransition<Self>> {
        let gc = global_context.clone();
        let tpe_agg_pk = context.tpe_agg_pk;
        let gateway_context_ready = context.clone();

        match &self.state {
            ReceiveSMState::Funding => {
                vec![StateTransition::new(
                    Self::await_decryption_shares(
                        global_context.clone(),
                        context.tpe_pks.clone(),
                        tpe_agg_pk,
                        self.common.outpoint,
                        self.common.contract.clone(),
                    ),
                    move |dbtx, outcome, old_state| {
                        Box::pin(Self::transition_decryption_shares(
                            dbtx,
                            outcome,
                            old_state,
                            gc.clone(),
                            gateway_context_ready.clone(),
                        ))
                    },
                )]
            }
            ReceiveSMState::Success(..)
            | ReceiveSMState::Rejected(..)
            | ReceiveSMState::Refunding(..)
            | ReceiveSMState::Failure => {
                vec![]
            }
        }
    }

    fn operation_id(&self) -> OperationId {
        self.common.operation_id
    }
}

impl ReceiveStateMachine {
    async fn await_decryption_shares(
        global_context: DynGlobalClientContext,
        tpe_pks: BTreeMap<PeerId, PublicKeyShare>,
        tpe_agg_pk: AggregatePublicKey,
        outpoint: OutPoint,
        contract: IncomingContract,
    ) -> DecryptionOutcome {
        let num_peers = global_context.api().all_peers().to_num_peers();
        let module_api = global_context.module_api();
        let tpe_pks = Arc::new(tpe_pks);
        let contract = Arc::new(contract);

        // The decryption key share endpoint long-polls until the share exists, which
        // happens atomically when the funding transaction is accepted. We can therefore
        // fire the request up front and let it overlap with awaiting transaction
        // acceptance instead of serializing the two, saving a round trip on the happy
        // path. Acceptance is still awaited to detect a rejected funding transaction.
        let decryption_shares = module_api.request_with_strategy_retry(
            FilterMapThreshold::new(
                {
                    let tpe_pks = tpe_pks.clone();
                    let contract = contract.clone();

                    move |peer_id, share: DecryptionKeyShare| {
                        let tpe_pks = tpe_pks.clone();
                        let contract = contract.clone();

                        // Four pairings per share; keep them off the async workers.
                        runtime::spawn_blocking(move || {
                            let pk =
                                tpe_pks
                                    .get(&peer_id)
                                    .ok_or(ServerError::InternalClientError(format!(
                                        "Missing TPE PK for peer {peer_id}?!"
                                    )))?;

                            if !contract.verify_decryption_share(pk, &share) {
                                return Err(ServerError::InvalidResponse(
                                    "Invalid decryption share".to_string(),
                                ));
                            }

                            Ok(share)
                        })
                    }
                },
                num_peers,
            ),
            DECRYPTION_KEY_SHARE_ENDPOINT.to_owned(),
            ApiRequestErased::new(outpoint),
        );

        let decryption_shares = std::pin::pin!(decryption_shares);
        let tx_accepted = std::pin::pin!(global_context.await_tx_accepted(outpoint.txid));

        let decryption_shares = match futures::future::select(decryption_shares, tx_accepted).await
        {
            futures::future::Either::Left((shares, _)) => shares,
            futures::future::Either::Right((accepted, decryption_shares)) => {
                if let Err(error) = accepted {
                    return DecryptionOutcome::Rejected(error);
                }

                decryption_shares.await
            }
        };

        // Aggregating the shares and verifying the key is four more pairings.
        runtime::spawn_blocking(move || {
            let agg_decryption_key = aggregate_dk_shares(
                &decryption_shares
                    .into_iter()
                    .map(|(peer, share)| (peer.to_usize() as u64, share))
                    .collect(),
            );

            if !contract.verify_agg_decryption_key(&tpe_agg_pk, &agg_decryption_key) {
                return DecryptionOutcome::InconsistentKeys;
            }

            let preimage = contract.decrypt_preimage(&agg_decryption_key);

            DecryptionOutcome::Decrypted(agg_decryption_key, preimage)
        })
        .await
    }

    async fn transition_decryption_shares(
        dbtx: &mut ClientSMDatabaseTransaction<'_, '_>,
        outcome: DecryptionOutcome,
        old_state: ReceiveStateMachine,
        global_context: DynGlobalClientContext,
        client_ctx: GatewayClientContextV2,
    ) -> ReceiveStateMachine {
        let payment_image = old_state.common.contract.commitment.payment_image.clone();

        let (agg_decryption_key, preimage) = match outcome {
            DecryptionOutcome::Rejected(error) => {
                client_ctx
                    .module
                    .client_ctx
                    .log_event(
                        &mut dbtx.module_tx(),
                        IncomingPaymentFailed {
                            payment_image,
                            error: error.clone(),
                        },
                    )
                    .await;

                return old_state.update(ReceiveSMState::Rejected(error));
            }
            DecryptionOutcome::InconsistentKeys => {
                warn!(target: LOG_CLIENT_MODULE_GW, "Failed to obtain decryption key. Client config's public keys are inconsistent");

                client_ctx
                    .module
                    .client_ctx
                    .log_event(
                        &mut dbtx.module_tx(),
                        IncomingPaymentFailed {
                            payment_image,
                            error: "Client config's public keys are inconsistent".to_string(),
                        },
                    )
                    .await;

                return old_state.update(ReceiveSMState::Failure);
            }
            DecryptionOutcome::Decrypted(agg_decryption_key, preimage) => {
                (agg_decryption_key, preimage)
            }
        };

        if let Some(preimage) = preimage {
            client_ctx
                .module
                .client_ctx
                .log_event(
                    &mut dbtx.module_tx(),
                    IncomingPaymentSucceeded { payment_image },
                )
                .await;

            return old_state.update(ReceiveSMState::Success(preimage));
        }

        let client_input = ClientInput::<LightningInput> {
            input: LightningInput::V0(LightningInputV0::Incoming(
                old_state.common.outpoint,
                agg_decryption_key,
            )),
            amounts: Amounts::new_bitcoin(old_state.common.contract.commitment.amount),
            keys: vec![old_state.common.refund_keypair],
        };

        let outpoints = match global_context
            .claim_inputs(
                dbtx,
                // The input of the refund tx is managed by this state machine
                ClientInputBundle::new_no_sm(vec![client_input]),
            )
            .await
        {
            Ok(outpoints) => outpoints.into_iter().collect(),
            // The contract is the refund transaction's only input, so this fails
            // exactly when the federation's fees exceed the contract, leaving nothing
            // to refund. The gateway keeps the lightning payment it already received
            // for it, so the loss is bounded by the contract; a panic here would not
            // be, because it runs in the executor and would take down every payment
            // this gateway handles for the federation, on every restart.
            Err(err) => {
                warn!(
                    target: LOG_CLIENT_MODULE_GW,
                    err = %err.fmt_compact(),
                    amount = %old_state.common.contract.commitment.amount,
                    "Not refunding incoming contract, its amount does not cover the refund fee"
                );

                client_ctx
                    .module
                    .client_ctx
                    .log_event(
                        &mut dbtx.module_tx(),
                        IncomingPaymentFailed {
                            payment_image: old_state
                                .common
                                .contract
                                .commitment
                                .payment_image
                                .clone(),
                            error: "Contract does not cover the refund fee".to_string(),
                        },
                    )
                    .await;

                return old_state.update(ReceiveSMState::Failure);
            }
        };

        client_ctx
            .module
            .client_ctx
            .log_event(
                &mut dbtx.module_tx(),
                IncomingPaymentFailed {
                    payment_image: old_state.common.contract.commitment.payment_image.clone(),
                    error: "Failed to decrypt preimage".to_string(),
                },
            )
            .await;

        old_state.update(ReceiveSMState::Refunding(outpoints))
    }
}