1use core::fmt;
2use std::collections::BTreeMap;
3
4use anyhow::anyhow;
5use fedimint_api_client::api::{FederationApiExt, ServerError};
6use fedimint_api_client::query::FilterMapThreshold;
7use fedimint_client_module::DynGlobalClientContext;
8use fedimint_client_module::sm::{ClientSMDatabaseTransaction, State, StateTransition};
9use fedimint_client_module::transaction::{ClientInput, ClientInputBundle};
10use fedimint_core::core::OperationId;
11use fedimint_core::encoding::{Decodable, Encodable};
12use fedimint_core::module::{Amounts, ApiRequestErased};
13use fedimint_core::secp256k1::Keypair;
14use fedimint_core::{NumPeersExt, OutPoint, PeerId};
15use fedimint_lnv2_common::contracts::IncomingContract;
16use fedimint_lnv2_common::endpoint_constants::DECRYPTION_KEY_SHARE_ENDPOINT;
17use fedimint_lnv2_common::{LightningInput, LightningInputV0};
18use fedimint_logging::LOG_CLIENT_MODULE_GW;
19use tpe::{AggregatePublicKey, DecryptionKeyShare, PublicKeyShare, aggregate_dk_shares};
20use tracing::warn;
21
22use super::events::{IncomingPaymentFailed, IncomingPaymentSucceeded};
23use crate::GatewayClientContextV2;
24
25#[derive(Debug, Clone, Eq, PartialEq, Hash, Decodable, Encodable)]
26pub struct ReceiveStateMachine {
27 pub common: ReceiveSMCommon,
28 pub state: ReceiveSMState,
29}
30
31impl ReceiveStateMachine {
32 pub fn update(&self, state: ReceiveSMState) -> Self {
33 Self {
34 common: self.common.clone(),
35 state,
36 }
37 }
38}
39
40impl fmt::Display for ReceiveStateMachine {
41 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42 write!(
43 f,
44 "Receive State Machine Operation ID: {:?} State: {}",
45 self.common.operation_id, self.state
46 )
47 }
48}
49
50#[derive(Debug, Clone, Eq, PartialEq, Hash, Decodable, Encodable)]
51pub struct ReceiveSMCommon {
52 pub operation_id: OperationId,
53 pub contract: IncomingContract,
54 pub outpoint: OutPoint,
55 pub refund_keypair: Keypair,
56}
57
58#[derive(Debug, Clone, Eq, PartialEq, Hash, Decodable, Encodable)]
59pub enum ReceiveSMState {
60 Funding,
61 Rejected(String),
62 Success([u8; 32]),
63 Failure,
64 Refunding(Vec<OutPoint>),
65}
66
67impl fmt::Display for ReceiveSMState {
68 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69 match self {
70 ReceiveSMState::Funding => write!(f, "Funding"),
71 ReceiveSMState::Rejected(_) => write!(f, "Rejected"),
72 ReceiveSMState::Success(_) => write!(f, "Success"),
73 ReceiveSMState::Failure => write!(f, "Failure"),
74 ReceiveSMState::Refunding(_) => write!(f, "Refunding"),
75 }
76 }
77}
78
79#[cfg_attr(doc, aquamarine::aquamarine)]
80impl State for ReceiveStateMachine {
92 type ModuleContext = GatewayClientContextV2;
93
94 fn transitions(
95 &self,
96 context: &Self::ModuleContext,
97 global_context: &DynGlobalClientContext,
98 ) -> Vec<StateTransition<Self>> {
99 let gc = global_context.clone();
100 let tpe_agg_pk = context.tpe_agg_pk;
101 let gateway_context_ready = context.clone();
102
103 match &self.state {
104 ReceiveSMState::Funding => {
105 vec![StateTransition::new(
106 Self::await_decryption_shares(
107 global_context.clone(),
108 context.tpe_pks.clone(),
109 self.common.outpoint,
110 self.common.contract.clone(),
111 ),
112 move |dbtx, output_outcomes, old_state| {
113 Box::pin(Self::transition_decryption_shares(
114 dbtx,
115 output_outcomes,
116 old_state,
117 gc.clone(),
118 tpe_agg_pk,
119 gateway_context_ready.clone(),
120 ))
121 },
122 )]
123 }
124 ReceiveSMState::Success(..)
125 | ReceiveSMState::Rejected(..)
126 | ReceiveSMState::Refunding(..)
127 | ReceiveSMState::Failure => {
128 vec![]
129 }
130 }
131 }
132
133 fn operation_id(&self) -> OperationId {
134 self.common.operation_id
135 }
136}
137
138impl ReceiveStateMachine {
139 async fn await_decryption_shares(
140 global_context: DynGlobalClientContext,
141 tpe_pks: BTreeMap<PeerId, PublicKeyShare>,
142 outpoint: OutPoint,
143 contract: IncomingContract,
144 ) -> Result<BTreeMap<PeerId, DecryptionKeyShare>, String> {
145 global_context.await_tx_accepted(outpoint.txid).await?;
146
147 Ok(global_context
148 .module_api()
149 .request_with_strategy_retry(
150 FilterMapThreshold::new(
151 move |peer_id, share: DecryptionKeyShare| {
152 if !contract.verify_decryption_share(
153 tpe_pks
154 .get(&peer_id)
155 .ok_or(ServerError::InternalClientError(anyhow!(
156 "Missing TPE PK for peer {peer_id}?!"
157 )))?,
158 &share,
159 ) {
160 return Err(fedimint_api_client::api::ServerError::InvalidResponse(
161 anyhow!("Invalid decryption share"),
162 ));
163 }
164
165 Ok(share)
166 },
167 global_context.api().all_peers().to_num_peers(),
168 ),
169 DECRYPTION_KEY_SHARE_ENDPOINT.to_owned(),
170 ApiRequestErased::new(outpoint),
171 )
172 .await)
173 }
174
175 async fn transition_decryption_shares(
176 dbtx: &mut ClientSMDatabaseTransaction<'_, '_>,
177 decryption_shares: Result<BTreeMap<PeerId, DecryptionKeyShare>, String>,
178 old_state: ReceiveStateMachine,
179 global_context: DynGlobalClientContext,
180 tpe_agg_pk: AggregatePublicKey,
181 client_ctx: GatewayClientContextV2,
182 ) -> ReceiveStateMachine {
183 let decryption_shares = match decryption_shares {
184 Ok(decryption_shares) => decryption_shares
185 .into_iter()
186 .map(|(peer, share)| (peer.to_usize() as u64, share))
187 .collect(),
188 Err(error) => {
189 client_ctx
190 .module
191 .client_ctx
192 .log_event(
193 &mut dbtx.module_tx(),
194 IncomingPaymentFailed {
195 payment_image: old_state
196 .common
197 .contract
198 .commitment
199 .payment_image
200 .clone(),
201 error: error.clone(),
202 },
203 )
204 .await;
205
206 return old_state.update(ReceiveSMState::Rejected(error));
207 }
208 };
209
210 let agg_decryption_key = aggregate_dk_shares(&decryption_shares);
211
212 if !old_state
213 .common
214 .contract
215 .verify_agg_decryption_key(&tpe_agg_pk, &agg_decryption_key)
216 {
217 warn!(target: LOG_CLIENT_MODULE_GW, "Failed to obtain decryption key. Client config's public keys are inconsistent");
218
219 client_ctx
220 .module
221 .client_ctx
222 .log_event(
223 &mut dbtx.module_tx(),
224 IncomingPaymentFailed {
225 payment_image: old_state.common.contract.commitment.payment_image.clone(),
226 error: "Client config's public keys are inconsistent".to_string(),
227 },
228 )
229 .await;
230
231 return old_state.update(ReceiveSMState::Failure);
232 }
233
234 if let Some(preimage) = old_state
235 .common
236 .contract
237 .decrypt_preimage(&agg_decryption_key)
238 {
239 client_ctx
240 .module
241 .client_ctx
242 .log_event(
243 &mut dbtx.module_tx(),
244 IncomingPaymentSucceeded {
245 payment_image: old_state.common.contract.commitment.payment_image.clone(),
246 },
247 )
248 .await;
249
250 return old_state.update(ReceiveSMState::Success(preimage));
251 }
252
253 let client_input = ClientInput::<LightningInput> {
254 input: LightningInput::V0(LightningInputV0::Incoming(
255 old_state.common.outpoint,
256 agg_decryption_key,
257 )),
258 amounts: Amounts::new_bitcoin(old_state.common.contract.commitment.amount),
259 keys: vec![old_state.common.refund_keypair],
260 };
261
262 let outpoints = global_context
263 .claim_inputs(
264 dbtx,
265 ClientInputBundle::new_no_sm(vec![client_input]),
267 )
268 .await
269 .expect("Cannot claim input, additional funding needed")
270 .into_iter()
271 .collect();
272
273 client_ctx
274 .module
275 .client_ctx
276 .log_event(
277 &mut dbtx.module_tx(),
278 IncomingPaymentFailed {
279 payment_image: old_state.common.contract.commitment.payment_image.clone(),
280 error: "Failed to decrypt preimage".to_string(),
281 },
282 )
283 .await;
284
285 old_state.update(ReceiveSMState::Refunding(outpoints))
286 }
287}