From 6ff6c45a6be80714e0d6f84af35949c2c4a758b9 Mon Sep 17 00:00:00 2001 From: positiveblue Date: Thu, 2 Feb 2023 03:43:01 -0800 Subject: [PATCH] channeldb: split `addHTLCs` logic in the `UpdateInvoice` method --- channeldb/invoice_test.go | 27 +++- channeldb/invoices.go | 258 +++++++++++++++++++++++++++++++++++++- invoices/update.go | 7 +- 3 files changed, 283 insertions(+), 9 deletions(-) diff --git a/channeldb/invoice_test.go b/channeldb/invoice_test.go index 1520fcee5..ce2d960b6 100644 --- a/channeldb/invoice_test.go +++ b/channeldb/invoice_test.go @@ -461,7 +461,8 @@ func TestInvoiceCancelSingleHtlc(t *testing.T) { } return &invpkg.InvoiceUpdateDesc{ - AddHtlcs: htlcs, + UpdateType: invpkg.AddHTLCsUpdate, + AddHtlcs: htlcs, }, nil } @@ -1533,6 +1534,7 @@ func getUpdateInvoice(amt lnwire.MilliSatoshi) invpkg.InvoiceUpdateCallback { }, } update := &invpkg.InvoiceUpdateDesc{ + UpdateType: invpkg.AddHTLCsUpdate, State: &invpkg.InvoiceStateUpdateDesc{ Preimage: invoice.Terms.PaymentPreimage, NewState: invpkg.ContractSettled, @@ -1590,7 +1592,10 @@ func TestCustomRecords(t *testing.T) { }, } - return &invpkg.InvoiceUpdateDesc{AddHtlcs: htlcs}, nil + return &invpkg.InvoiceUpdateDesc{ + AddHtlcs: htlcs, + UpdateType: invpkg.AddHTLCsUpdate, + }, nil } _, err = db.UpdateInvoice(ref, nil, callback) @@ -1667,7 +1672,10 @@ func testInvoiceHtlcAMPFields(t *testing.T, isAMP bool) { }, } - return &invpkg.InvoiceUpdateDesc{AddHtlcs: htlcs}, nil + return &invpkg.InvoiceUpdateDesc{ + AddHtlcs: htlcs, + UpdateType: invpkg.AddHTLCsUpdate, + }, nil } ref := invpkg.InvoiceRefByHash(payHash) @@ -2132,8 +2140,9 @@ func updateAcceptAMPHtlc(id uint64, amt lnwire.MilliSatoshi, } update := &invpkg.InvoiceUpdateDesc{ - State: state, - AddHtlcs: htlcs, + State: state, + AddHtlcs: htlcs, + UpdateType: invpkg.AddHTLCsUpdate, } return update, nil @@ -2156,6 +2165,10 @@ func getUpdateInvoiceAMPSettle(setID *[32]byte, preimage [32]byte, } update := &invpkg.InvoiceUpdateDesc{ + // TODO(positiveblue): this would be an invalid update + // because tires to settle an AMP invoice without adding + // any new htlc. + UpdateType: invpkg.AddHTLCsUpdate, State: &invpkg.InvoiceStateUpdateDesc{ Preimage: nil, NewState: invpkg.ContractSettled, @@ -2276,12 +2289,16 @@ func testUpdateHTLCPreimages(t *testing.T, test updateHTLCPreimageTestCase) { invoice *invpkg.Invoice) (*invpkg.InvoiceUpdateDesc, error) { update := &invpkg.InvoiceUpdateDesc{ + // TODO(positiveblue): this would be an invalid update + // because tires to settle an AMP invoice without adding + // any new htlc. State: &invpkg.InvoiceStateUpdateDesc{ Preimage: nil, NewState: invpkg.ContractSettled, HTLCPreimages: htlcPreimages, SetID: setID, }, + UpdateType: invpkg.AddHTLCsUpdate, } return update, nil diff --git a/channeldb/invoices.go b/channeldb/invoices.go index ce43d1063..b8780600a 100644 --- a/channeldb/invoices.go +++ b/channeldb/invoices.go @@ -1858,7 +1858,7 @@ func settleHtlcsAmp(invoice *invpkg.Invoice, // updateInvoice fetches the invoice, obtains the update descriptor from the // callback and applies the updates in a single db transaction. -func (d *DB) updateInvoice(hash *lntypes.Hash, refSetID *invpkg.SetID, invoices, //nolint:lll,funlen +func (d *DB) updateInvoice(hash *lntypes.Hash, refSetID *invpkg.SetID, invoices, settleIndex, setIDIndex kvdb.RwBucket, invoiceNum []byte, callback invpkg.InvoiceUpdateCallback) (*invpkg.Invoice, error) { @@ -1897,6 +1897,12 @@ func (d *DB) updateInvoice(hash *lntypes.Hash, refSetID *invpkg.SetID, invoices, case invpkg.CancelHTLCsUpdate: return d.cancelHTLCs(invoices, invoiceNum, &invoice, update) + case invpkg.AddHTLCsUpdate: + return d.addHTLCs( + invoices, settleIndex, setIDIndex, invoiceNum, &invoice, + hash, update, + ) + case invpkg.SettleHodlInvoiceUpdate: return d.settleHodlInvoice( invoices, settleIndex, invoiceNum, &invoice, hash, @@ -2225,6 +2231,256 @@ func (d *DB) cancelHTLCs(invoices kvdb.RwBucket, invoiceNum []byte, return invoice, nil } +// addHTLCs tries to add the htlcs in the given InvoiceUpdateDesc. +func (d *DB) addHTLCs(invoices, settleIndex, //nolint:funlen + setIDIndex kvdb.RwBucket, invoiceNum []byte, invoice *invpkg.Invoice, + hash *lntypes.Hash, update *invpkg.InvoiceUpdateDesc) (*invpkg.Invoice, + error) { + + var setID *[32]byte + invoiceIsAMP := invoice.IsAMP() + if invoiceIsAMP && update.State != nil { + setID = update.State.SetID + } + timestamp := d.clock.Now() + + // Process add actions from update descriptor. + htlcsAmpUpdate := make(map[invpkg.SetID]map[models.CircuitKey]*invpkg.InvoiceHTLC) //nolint:lll + for key, htlcUpdate := range update.AddHtlcs { + if _, exists := invoice.Htlcs[key]; exists { + return nil, fmt.Errorf("duplicate add of htlc %v", key) + } + + // Force caller to supply htlc without custom records in a + // consistent way. + if htlcUpdate.CustomRecords == nil { + return nil, errors.New("nil custom records map") + } + + if invoiceIsAMP { + if htlcUpdate.AMP == nil { + return nil, fmt.Errorf("unable to add htlc "+ + "without AMP data to AMP invoice(%v)", + invoice.AddIndex) + } + + // Check if this SetID already exist. + htlcSetID := htlcUpdate.AMP.Record.SetID() + setIDInvNum := setIDIndex.Get(htlcSetID[:]) + + if setIDInvNum == nil { + err := setIDIndex.Put(htlcSetID[:], invoiceNum) + if err != nil { + return nil, err + } + } else if !bytes.Equal(setIDInvNum, invoiceNum) { + return nil, invpkg.ErrDuplicateSetID{ + SetID: htlcSetID, + } + } + } + + htlc := &invpkg.InvoiceHTLC{ + Amt: htlcUpdate.Amt, + MppTotalAmt: htlcUpdate.MppTotalAmt, + Expiry: htlcUpdate.Expiry, + AcceptHeight: uint32(htlcUpdate.AcceptHeight), + AcceptTime: timestamp, + State: invpkg.HtlcStateAccepted, + CustomRecords: htlcUpdate.CustomRecords, + AMP: htlcUpdate.AMP.Copy(), + } + + invoice.Htlcs[key] = htlc + + // Collect the set of new HTLCs so we can write them properly + // below, but only if this is an AMP invoice. + if invoiceIsAMP { + updateHtlcsAmp( + invoice, htlcsAmpUpdate, htlc, + htlcUpdate.AMP.Record.SetID(), key, + ) + } + } + + // At this point, the set of accepted HTLCs should be fully + // populated with added HTLCs or removed of canceled ones. Update + // invoice state if the update descriptor indicates an invoice state + // change, which depends on having an accurate view of the accepted + // HTLCs. + if update.State != nil { + newState, err := updateInvoiceState( + invoice, hash, *update.State, + ) + if err != nil { + return nil, err + } + + // If this isn't an AMP invoice, then we'll go ahead and update + // the invoice state directly here. For AMP invoices, we + // instead will keep the top-level invoice open, and instead + // update the state of each _htlc set_ instead. However, we'll + // allow the invoice to transition to the cancelled state + // regardless. + if !invoiceIsAMP || *newState == invpkg.ContractCanceled { + invoice.State = *newState + } + + // If this is a non-AMP invoice, then the state can eventually + // go to ContractSettled, so we pass in nil value as part of + // setSettleMetaFields. + isSettled := update.State.NewState == invpkg.ContractSettled + if !invoiceIsAMP && isSettled { + err := setSettleMetaFields( + settleIndex, invoiceNum, invoice, timestamp, + nil, + ) + if err != nil { + return nil, err + } + } + } + + // The set of HTLC pre-images will only be set if we were actually able + // to reconstruct all the AMP pre-images. + var settleEligibleAMP bool + if update.State != nil { + settleEligibleAMP = len(update.State.HTLCPreimages) != 0 + } + + // With any invoice level state transitions recorded, we'll now + // finalize the process by updating the state transitions for + // individual HTLCs + var ( + settledSetIDs = make(map[invpkg.SetID]struct{}) + amtPaid lnwire.MilliSatoshi + ) + for key, htlc := range invoice.Htlcs { + // Set the HTLC preimage for any AMP HTLCs. + if setID != nil && update.State != nil { + preimage, ok := update.State.HTLCPreimages[key] + switch { + // If we don't already have a preimage for this HTLC, we + // can set it now. + case ok && htlc.AMP.Preimage == nil: + htlc.AMP.Preimage = &preimage + + // Otherwise, prevent over-writing an existing + // preimage. Ignore the case where the preimage is + // identical. + case ok && *htlc.AMP.Preimage != preimage: + return nil, invpkg.ErrHTLCPreimageAlreadyExists + } + } + + // The invoice state may have changed and this could have + // implications for the states of the individual htlcs. Align + // the htlc state with the current invoice state. + // + // If we have all the pre-images for an AMP invoice, then we'll + // act as if we're able to settle the entire invoice. We need + // to do this since it's possible for us to settle AMP invoices + // while the contract state (on disk) is still in the accept + // state. + htlcContextState := invoice.State + if settleEligibleAMP { + htlcContextState = invpkg.ContractSettled + } + htlcSettled, err := updateHtlc( + timestamp, htlc, htlcContextState, setID, + ) + if err != nil { + return nil, err + } + + // If the HTLC has being settled for the first time, and this + // is an AMP invoice, then we'll need to update some additional + // meta data state. + if htlcSettled && invoiceIsAMP { + settleHtlcsAmp( + invoice, settledSetIDs, htlcsAmpUpdate, htlc, + key, + ) + } + + accepted := htlc.State == invpkg.HtlcStateAccepted + settled := htlc.State == invpkg.HtlcStateSettled + invoiceStateReady := accepted || settled + + if !invoiceIsAMP { + // Update the running amount paid to this invoice. We + // don't include accepted htlcs when the invoice is + // still open. + if invoice.State != invpkg.ContractOpen && + invoiceStateReady { + + amtPaid += htlc.Amt + } + } else { + // For AMP invoices, since we won't always be reading + // out the total invoice set each time, we'll instead + // accumulate newly added invoices to the total amount + // paid. + if _, ok := update.AddHtlcs[key]; !ok { + continue + } + + // Update the running amount paid to this invoice. AMP + // invoices never go to the settled state, so if it's + // open, then we tally the HTLC. + if invoice.State == invpkg.ContractOpen && + invoiceStateReady { + + amtPaid += htlc.Amt + } + } + } + + // For non-AMP invoices we recalculate the amount paid from scratch + // each time, while for AMP invoices, we'll accumulate only based on + // newly added HTLCs. + if !invoiceIsAMP { + invoice.AmtPaid = amtPaid + } else { + invoice.AmtPaid += amtPaid + } + + // As we don't update the settle index above for AMP invoices, we'll do + // it here for each sub-AMP invoice that was settled. + for settledSetID := range settledSetIDs { + settledSetID := settledSetID + err := setSettleMetaFields( + settleIndex, invoiceNum, invoice, timestamp, + &settledSetID, + ) + if err != nil { + return nil, err + } + } + + // Reserialize and update invoice. + var buf bytes.Buffer + if err := serializeInvoice(&buf, invoice); err != nil { + return nil, err + } + + if err := invoices.Put(invoiceNum, buf.Bytes()); err != nil { + return nil, err + } + + // If this is an AMP invoice, then we'll actually store the rest of the + // HTLCs in-line with the invoice, using the invoice ID as a prefix, + // and the AMP key as a suffix: invoiceNum || setID. + if invoiceIsAMP { + err := updateAMPInvoices(invoices, invoiceNum, htlcsAmpUpdate) + if err != nil { + return nil, err + } + } + + return invoice, nil +} + // settleHodlInvoice marks a hodl invoice as settled. // // NOTE: Currently it is not possible to have HODL AMP invoices. diff --git a/invoices/update.go b/invoices/update.go index 2709d5f03..390f325a7 100644 --- a/invoices/update.go +++ b/invoices/update.go @@ -238,7 +238,8 @@ func updateMpp(ctx *invoiceUpdateCtx, inv *Invoice) (*InvoiceUpdateDesc, } update := InvoiceUpdateDesc{ - AddHtlcs: newHtlcs, + UpdateType: AddHTLCsUpdate, + AddHtlcs: newHtlcs, } // If the invoice cannot be settled yet, only record the htlc. @@ -252,7 +253,6 @@ func updateMpp(ctx *invoiceUpdateCtx, inv *Invoice) (*InvoiceUpdateDesc, if inv.HodlInvoice { update.State = &InvoiceStateUpdateDesc{ NewState: ContractAccepted, - SetID: setID, } return &update, ctx.acceptRes(resultAccepted), nil } @@ -428,7 +428,8 @@ func updateLegacy(ctx *invoiceUpdateCtx, } update := InvoiceUpdateDesc{ - AddHtlcs: newHtlcs, + AddHtlcs: newHtlcs, + UpdateType: AddHTLCsUpdate, } // Don't update invoice state if we are accepting a duplicate payment.