channeldb: for AMP don't update the main invoice state in updateInvoice

In this commit, we modify the way we handle state updates for AMP
invoices. With the current logic, once an invoice is settled, we'll
update the primary invoice state. This means that a user can't take the
same payment_addr, w/ a new set_id and pay the invoice again.

To remedy this, we'll move to instead _not_ updating the main invoice
state each time a new htlc Set (group of HTLCs according to setID) is
added. Instead, given that each HTLC stores an individual state, we'll
instead just use that directly from now on.

We also update TestSetIDIndex to account for new repeated settle AMP
logic.

channeldb/invoices.go

@@ -1702,7 +1702,7 @@ func fetchFilteredAmpInvoices(invoiceBucket kvdb.RBucket,
 		if htlcSetBytes == nil {
 			// A set ID was passed in, but we don't have this
 			// stored yet, meaning that the setID is being added
-			// for the frist time.
+			// for the first time.
 			return htlcs, ErrInvoiceNotFound
 		}
 
@@ -2369,6 +2369,116 @@ func updateAMPInvoices(invoiceBucket kvdb.RwBucket, invoiceNum []byte,
 	return nil
 }
 
+// updateHtlcsAmp takes an invoice, and a new HTLC to be added (along with its
+// set ID), and update sthe internal AMP state of an invoice, and also tallies
+// the set of HTLCs to be updated on disk.
+func updateHtlcsAmp(invoice *Invoice,
+	updateMap map[SetID]map[CircuitKey]*InvoiceHTLC, htlc *InvoiceHTLC,
+	setID SetID, circuitKey CircuitKey) {
+
+	ampState, ok := invoice.AMPState[setID]
+	if !ok {
+		// If an entry for this set ID doesn't already exist, then
+		// we'll need to create it.
+		ampState = InvoiceStateAMP{
+			State:       HtlcStateAccepted,
+			InvoiceKeys: make(map[CircuitKey]struct{}),
+		}
+	}
+
+	ampState.AmtPaid += htlc.Amt
+	ampState.InvoiceKeys[circuitKey] = struct{}{}
+
+	// Due to the way maps work, we need to read out the value, update it,
+	// then re-assign it into the map.
+	invoice.AMPState[setID] = ampState
+
+	// Now that we've updated the invoice state, we'll inform the caller of
+	// the _neitre_ HTLC set they need to write for this new set ID.
+	if _, ok := updateMap[setID]; !ok {
+		// If we're just now creating the HTLCs for this set then we'll
+		// also pull in the existing HTLCs are part of this set, so we
+		// can write them all to disk together (same value)
+		updateMap[setID] = invoice.HTLCSet(
+			(*[32]byte)(&setID), HtlcStateAccepted,
+		)
+	}
+	updateMap[setID][circuitKey] = htlc
+}
+
+// cancelHtlcsAmp processes a cancellation of an HTLC that belongs to an AMP
+// HTLC set. We'll need to update the meta data in the  main invoice, and also
+// apply the new update to the update MAP, since all the HTLCs for a given HTLC
+// set need to be written in-line with each other.
+func cancelHtlcsAmp(invoice *Invoice,
+	updateMap map[SetID]map[CircuitKey]*InvoiceHTLC, htlc *InvoiceHTLC,
+	circuitKey CircuitKey) {
+
+	setID := htlc.AMP.Record.SetID()
+
+	// First, we'll update the state of the entire HTLC set to cancelled.
+	ampState := invoice.AMPState[setID]
+	ampState.State = HtlcStateCanceled
+
+	ampState.InvoiceKeys[circuitKey] = struct{}{}
+	ampState.AmtPaid -= htlc.Amt
+
+	// With the state update,d we'll set the new value so the struct
+	// changes are propagated.
+	invoice.AMPState[setID] = ampState
+
+	if _, ok := updateMap[setID]; !ok {
+		// Only HTLCs in the accepted state, can be cancelled, but we
+		// also want to merge that with HTLCs that may be canceled as
+		// well since it can be cancelled one by one.
+		updateMap[setID] = invoice.HTLCSet(&setID, HtlcStateAccepted)
+
+		cancelledHtlcs := invoice.HTLCSet(&setID, HtlcStateCanceled)
+		for htlcKey, htlc := range cancelledHtlcs {
+			updateMap[setID][htlcKey] = htlc
+		}
+	}
+
+	// Finally, include the newly cancelled HTLC in the set of HTLCs we
+	// need to cancel.
+	updateMap[setID][circuitKey] = htlc
+
+	// We'll only decrement the total amount paid if the invoice was
+	// already in the accepted state.
+	if invoice.AmtPaid != 0 {
+		invoice.AmtPaid -= htlc.Amt
+	}
+}
+
+// settleHtlcsAmp processes a new settle operation on an HTLC set for an AMP
+// invoice. We'll update some meta data in the main invoice, and also signal
+// that this HTLC set needs to be re-written back to disk.
+func settleHtlcsAmp(invoice *Invoice,
+	settledSetIDs map[SetID]struct{},
+	updateMap map[SetID]map[CircuitKey]*InvoiceHTLC, htlc *InvoiceHTLC,
+	circuitKey CircuitKey) {
+
+	// First, add the set ID to the set that was settled in this invoice
+	// update. We'll use this later to update the settle index.
+	setID := htlc.AMP.Record.SetID()
+	settledSetIDs[setID] = struct{}{}
+
+	// Next update the main AMP meta-data to indicate that this HTLC set
+	// has been fully settled.
+	ampState := invoice.AMPState[setID]
+	ampState.State = HtlcStateSettled
+
+	ampState.InvoiceKeys[circuitKey] = struct{}{}
+
+	invoice.AMPState[setID] = ampState
+
+	// Finally, we'll add this to the set of HTLCs that need to be updated.
+	if _, ok := updateMap[setID]; !ok {
+		updateMap[setID] = make(map[CircuitKey]*InvoiceHTLC)
+	}
+	updateMap[setID][circuitKey] = htlc
+}
+
 // 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, invoices,
@@ -2399,6 +2509,10 @@ func (d *DB) updateInvoice(hash *lntypes.Hash, invoices,
 		newState = invoice.State
 		setID    *[32]byte
 	)
+
+	// We can either get the set ID from the main state update (if the
+	// state is changing), or via the hint passed in returned by the update
+	// call back.
 	if update.State != nil {
 		setID = update.State.SetID
 		newState = update.State.NewState
@@ -2406,7 +2520,12 @@ func (d *DB) updateInvoice(hash *lntypes.Hash, invoices,
 
 	now := d.clock.Now()
 
+	invoiceIsAMP := invoiceCopy.Terms.Features.HasFeature(
+		lnwire.AMPOptional,
+	)
+
 	// Process add actions from update descriptor.
+	htlcsAmpUpdate := make(map[SetID]map[CircuitKey]*InvoiceHTLC)
 	for key, htlcUpdate := range update.AddHtlcs {
 		if _, exists := invoice.Htlcs[key]; exists {
 			return nil, fmt.Errorf("duplicate add of htlc %v", key)
@@ -2421,8 +2540,9 @@ func (d *DB) updateInvoice(hash *lntypes.Hash, invoices,
 		// If a newly added HTLC has an associated set id, use it to
 		// index this invoice in the set id index. An error is returned
 		// if we find the index already points to a different invoice.
+		var setID [32]byte
 		if htlcUpdate.AMP != nil {
-			setID := htlcUpdate.AMP.Record.SetID()
+			setID = htlcUpdate.AMP.Record.SetID()
 			setIDInvNum := setIDIndex.Get(setID[:])
 			if setIDInvNum == nil {
 				err = setIDIndex.Put(setID[:], invoiceNum)
@@ -2446,11 +2566,21 @@ func (d *DB) updateInvoice(hash *lntypes.Hash, invoices,
 		}
 
 		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, setID, key,
+			)
+		}
 	}
 
 	// Process cancel actions from update descriptor.
 	cancelHtlcs := update.CancelHtlcs
 	for key, htlc := range invoice.Htlcs {
+		htlc := htlc
+
 		// Check whether this htlc needs to be canceled. If it does,
 		// update the htlc state to Canceled.
 		_, cancel := cancelHtlcs[key]
@@ -2472,6 +2602,14 @@ func (d *DB) updateInvoice(hash *lntypes.Hash, invoices,
 		// Delete processed cancel action, so that we can check later
 		// that there are no actions left.
 		delete(cancelHtlcs, key)
+
+		// Tally this into the set of HTLCs that need to be updated on
+		// disk, but once again, only if this is an AMP invoice.
+		if invoiceIsAMP {
+			cancelHtlcsAmp(
+				&invoice, htlcsAmpUpdate, htlc, key,
+			)
+		}
 	}
 
 	// Verify that we didn't get an action for htlcs that are not present on
@@ -2486,12 +2624,27 @@ func (d *DB) updateInvoice(hash *lntypes.Hash, invoices,
 	// change, which depends on having an accurate view of the accepted
 	// HTLCs.
 	if update.State != nil {
-		err := updateInvoiceState(&invoice, hash, *update.State)
+		newState, err := updateInvoiceState(
+			&invoice, hash, *update.State,
+		)
 		if err != nil {
 			return nil, err
 		}
 
-		if update.State.NewState == ContractSettled {
+		// 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 == 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.
+		if !invoiceIsAMP && update.State.NewState == ContractSettled {
 			err := setSettleMetaFields(
 				settleIndex, invoiceNum, &invoice, now,
 			)
@@ -2501,23 +2654,34 @@ func (d *DB) updateInvoice(hash *lntypes.Hash, invoices,
 		}
 	}
 
-	// With any invoice level state transitions recorded, we'll now finalize
-	// the process by updating the state transitions for individual HTLCs
-	// and recalculate the total amount paid to the invoice.
-	var amtPaid lnwire.MilliSatoshi
+	// 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[SetID]struct{})
+		amtPaid       lnwire.MilliSatoshi
+	)
 	for key, htlc := range invoice.Htlcs {
 		// Set the HTLC preimage for any AMP HTLCs.
-		if setID != nil {
+		if setID != nil && update.State != nil {
 			preimage, ok := update.State.HTLCPreimages[key]
 			switch {
 
-			// If we don't already have a preiamge for this HTLC, we
+			// 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.
+			// Otherwise, prevent over-writing an existing
+			// preimage.  Ignore the case where the preimage is
+			// identical.
 			case ok && *htlc.AMP.Preimage != preimage:
 				return nil, ErrHTLCPreimageAlreadyExists
 
@@ -2527,21 +2691,68 @@ func (d *DB) updateInvoice(hash *lntypes.Hash, invoices,
 		// 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.
-		err := updateHtlc(now, htlc, invoice.State, setID)
+		//
+		// 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 = ContractSettled
+		}
+		htlcSettled, err := updateHtlc(
+			now, htlc, htlcContextState, setID,
+		)
 		if err != nil {
 			return nil, err
 		}
 
-		// Update the running amount paid to this invoice. We don't
-		// include accepted htlcs when the invoice is still open.
-		if invoice.State != ContractOpen &&
-			(htlc.State == HtlcStateAccepted ||
-				htlc.State == HtlcStateSettled) {
+		// 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,
+			)
+		}
 
-			amtPaid += htlc.Amt
+		invoiceStateReady := (htlc.State == HtlcStateAccepted ||
+			htlc.State == HtlcStateSettled)
+		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 != 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 == ContractOpen && invoiceStateReady {
+				amtPaid += htlc.Amt
+			}
 		}
 	}
-	invoice.AmtPaid = amtPaid
+
+	// 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
+	}
 
 	// Reserialize and update invoice.
 	var buf bytes.Buffer
@@ -2553,16 +2764,28 @@ func (d *DB) updateInvoice(hash *lntypes.Hash, invoices,
 		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
 }
 
-// updateInvoiceState validates and processes an invoice state update.
+// updateInvoiceState validates and processes an invoice state update. The new
+// state to transition to is returned, so the caller is able to select exactly
+// how the invoice state is updated.
 func updateInvoiceState(invoice *Invoice, hash *lntypes.Hash,
-	update InvoiceStateUpdateDesc) error {
+	update InvoiceStateUpdateDesc) (*ContractState, error) {
 
 	// Returning to open is never allowed from any state.
 	if update.NewState == ContractOpen {
-		return ErrInvoiceCannotOpen
+		return nil, ErrInvoiceCannotOpen
 	}
 
 	switch invoice.State {
@@ -2573,7 +2796,7 @@ func updateInvoiceState(invoice *Invoice, hash *lntypes.Hash,
 	// same checks that we apply to open invoices.
 	case ContractAccepted:
 		if update.NewState == ContractAccepted {
-			return ErrInvoiceCannotAccept
+			return nil, ErrInvoiceCannotAccept
 		}
 
 		fallthrough
@@ -2583,14 +2806,13 @@ func updateInvoiceState(invoice *Invoice, hash *lntypes.Hash,
 	// where we ensure the preimage is valid.
 	case ContractOpen:
 		if update.NewState == ContractCanceled {
-			invoice.State = update.NewState
-			return nil
+			return &update.NewState, nil
 		}
 
 		// Sanity check that the user isn't trying to settle or accept a
 		// non-existent HTLC set.
 		if len(invoice.HTLCSet(update.SetID, HtlcStateAccepted)) == 0 {
-			return ErrEmptyHTLCSet
+			return nil, ErrEmptyHTLCSet
 		}
 
 		// For AMP invoices, there are no invoice-level preimage checks.
@@ -2598,10 +2820,10 @@ func updateInvoiceState(invoice *Invoice, hash *lntypes.Hash,
 		// settle an AMP invoice with a preimage.
 		if update.SetID != nil {
 			if update.Preimage != nil {
-				return errors.New("AMP set cannot have preimage")
+				return nil, errors.New("AMP set cannot have " +
+					"preimage")
 			}
-			invoice.State = update.NewState
-			return nil
+			return &update.NewState, nil
 		}
 
 		switch {
@@ -2609,12 +2831,12 @@ func updateInvoiceState(invoice *Invoice, hash *lntypes.Hash,
 		// If an invoice-level preimage was supplied, but the InvoiceRef
 		// doesn't specify a hash (e.g. AMP invoices) we fail.
 		case update.Preimage != nil && hash == nil:
-			return ErrUnexpectedInvoicePreimage
+			return nil, ErrUnexpectedInvoicePreimage
 
 		// Validate the supplied preimage for non-AMP invoices.
 		case update.Preimage != nil:
 			if update.Preimage.Hash() != *hash {
-				return ErrInvoicePreimageMismatch
+				return nil, ErrInvoicePreimageMismatch
 			}
 			invoice.Terms.PaymentPreimage = update.Preimage
 
@@ -2628,23 +2850,21 @@ func updateInvoiceState(invoice *Invoice, hash *lntypes.Hash,
 		case update.NewState == ContractSettled &&
 			invoice.Terms.PaymentPreimage == nil:
 
-			return errors.New("unknown preimage")
+			return nil, errors.New("unknown preimage")
 		}
 
-		invoice.State = update.NewState
-
-		return nil
+		return &update.NewState, nil
 
 	// Once settled, we are in a terminal state.
 	case ContractSettled:
-		return ErrInvoiceAlreadySettled
+		return nil, ErrInvoiceAlreadySettled
 
 	// Once canceled, we are in a terminal state.
 	case ContractCanceled:
-		return ErrInvoiceAlreadyCanceled
+		return nil, ErrInvoiceAlreadyCanceled
 
 	default:
-		return errors.New("unknown state transition")
+		return nil, errors.New("unknown state transition")
 	}
 }
 
@@ -2670,18 +2890,20 @@ func cancelSingleHtlc(resolveTime time.Time, htlc *InvoiceHTLC,
 	return nil
 }
 
-// updateHtlc aligns the state of an htlc with the given invoice state.
+// updateHtlc aligns the state of an htlc with the given invoice state. A
+// boolean is returned if the HTLC was settled.
 func updateHtlc(resolveTime time.Time, htlc *InvoiceHTLC,
-	invState ContractState, setID *[32]byte) error {
+	invState ContractState, setID *[32]byte) (bool, error) {
 
-	trySettle := func(persist bool) error {
+	trySettle := func(persist bool) (bool, error) {
 		if htlc.State != HtlcStateAccepted {
-			return nil
+			return false, nil
 		}
 
-		// Settle the HTLC if it matches the settled set id. Since we
-		// only allow settling of one HTLC set (for now) we cancel any
-		// that do not match the set id.
+		// Settle the HTLC if it matches the settled set id. If
+		// there're other HTLCs with distinct setIDs, then we'll leave
+		// them, as they may eventually be settled as we permit
+		// multiple settles to a single pay_addr for AMP.
 		var htlcState HtlcState
 		if htlc.IsInHTLCSet(setID) {
 			// Non-AMP HTLCs can be settled immediately since we
@@ -2695,18 +2917,18 @@ func updateHtlc(resolveTime time.Time, htlc *InvoiceHTLC,
 			// the invoice level.
 			case setID == nil:
 
-			// At this popint, the setID is non-nil, meaning this is
+			// At this point, the setID is non-nil, meaning this is
 			// an AMP HTLC. We know that htlc.AMP cannot be nil,
 			// otherwise IsInHTLCSet would have returned false.
 			//
 			// Fail if an accepted AMP HTLC has no preimage.
 			case htlc.AMP.Preimage == nil:
-				return ErrHTLCPreimageMissing
+				return false, ErrHTLCPreimageMissing
 
 			// Fail if the accepted AMP HTLC has an invalid
 			// preimage.
 			case !htlc.AMP.Preimage.Matches(htlc.AMP.Hash):
-				return ErrHTLCPreimageMismatch
+				return false, ErrHTLCPreimageMismatch
 			}
 
 			htlcState = HtlcStateSettled
@@ -2720,7 +2942,7 @@ func updateHtlc(resolveTime time.Time, htlc *InvoiceHTLC,
 			htlc.ResolveTime = resolveTime
 		}
 
-		return nil
+		return persist && htlcState == HtlcStateSettled, nil
 	}
 
 	if invState == ContractSettled {
@@ -2734,7 +2956,7 @@ func updateHtlc(resolveTime time.Time, htlc *InvoiceHTLC,
 	// We should never find a settled HTLC on an invoice that isn't in
 	// ContractSettled.
 	if htlc.State == HtlcStateSettled {
-		return ErrHTLCAlreadySettled
+		return false, ErrHTLCAlreadySettled
 	}
 
 	switch invState {
@@ -2744,8 +2966,9 @@ func updateHtlc(resolveTime time.Time, htlc *InvoiceHTLC,
 			htlc.State = HtlcStateCanceled
 			htlc.ResolveTime = resolveTime
 		}
-		return nil
+		return false, nil
 
+	// TODO(roasbeef): never fully passed thru now?
 	case ContractAccepted:
 		// Check that we can settle the HTLCs. For legacy and MPP HTLCs
 		// this will be a NOP, but for AMP HTLCs this asserts that we
@@ -2755,10 +2978,10 @@ func updateHtlc(resolveTime time.Time, htlc *InvoiceHTLC,
 		return trySettle(false)
 
 	case ContractOpen:
-		return nil
+		return false, nil
 
 	default:
-		return errors.New("unknown state transition")
+		return false, errors.New("unknown state transition")
 	}
 }