invoices: fix amp invoice bug.

We now cancel all HTLCs of an AMP invoice as soon as it expires.
Otherwise because we mark the invoice as cancelled we would not
allow accepted HTLCs to be resolved via the invoiceEventLoop.

channeldb/invoices.go

@@ -650,18 +650,8 @@ func (d *DB) UpdateInvoice(_ context.Context, ref invpkg.InvoiceRef,
 			return err
 		}
 
-		// If the set ID hint is non-nil, then we'll use that to filter
-		// out the HTLCs for AMP invoice so we don't need to read them
-		// all out to satisfy the invoice callback below. If it's nil,
-		// then we pass in the zero set ID which means no HTLCs will be
-		// read out.
-		var invSetID invpkg.SetID
-
-		if setIDHint != nil {
-			invSetID = *setIDHint
-		}
 		invoice, err := fetchInvoice(
-			invoiceNum, invoices, []*invpkg.SetID{&invSetID}, false,
+			invoiceNum, invoices, []*invpkg.SetID{setIDHint}, false,
 		)
 		if err != nil {
 			return err
@@ -691,7 +681,7 @@ func (d *DB) UpdateInvoice(_ context.Context, ref invpkg.InvoiceRef,
 		// If this is an AMP update, then limit the returned AMP state
 		// to only the requested set ID.
 		if setIDHint != nil {
-			filterInvoiceAMPState(updatedInvoice, &invSetID)
+			filterInvoiceAMPState(updatedInvoice, setIDHint)
 		}
 
 		return nil
@@ -848,7 +838,23 @@ func (k *kvInvoiceUpdater) Finalize(updateType invpkg.UpdateType) error {
 		return k.storeSettleHodlInvoiceUpdate()
 
 	case invpkg.CancelInvoiceUpdate:
-		return k.serializeAndStoreInvoice()
+		err := k.serializeAndStoreInvoice()
+		if err != nil {
+			return 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 k.invoice.IsAMP() {
+			err := k.updateAMPInvoices()
+			if err != nil {
+				return err
+			}
+		}
+
+		return nil
 	}
 
 	return fmt.Errorf("unknown update type: %v", updateType)
@@ -1028,7 +1034,12 @@ func (k *kvInvoiceUpdater) serializeAndStoreInvoice() error {
 		return err
 	}
 
-	return k.invoicesBucket.Put(k.invoiceNum, buf.Bytes())
+	err := k.invoicesBucket.Put(k.invoiceNum, buf.Bytes())
+	if err != nil {
+		return err
+	}
+
+	return nil
 }
 
 // InvoicesSettledSince can be used by callers to catch up any settled invoices

invoices/interface.go

@@ -56,6 +56,11 @@ type InvoiceDB interface {
 	// passed payment hash. If an invoice matching the passed payment hash
 	// doesn't exist within the database, then the action will fail with a
 	// "not found" error.
+	// The setIDHint is used to signal whether AMP HTLCs should be fetched
+	// for the invoice. If a blank setID is passed no HTLCs will be fetched
+	// in case of an AMP invoice. Nil means all HTLCs for all sub AMP
+	// invoices will be fetched and if a specific setID is supplied only
+	// HTLCs for that setID will be fetched.
 	//
 	// The update is performed inside the same database transaction that
 	// fetches the invoice and is therefore atomic. The fields to update

invoices/invoiceregistry.go

@@ -705,6 +705,12 @@ func (i *InvoiceRegistry) cancelSingleHtlc(invoiceRef InvoiceRef,
 	// Intercept the update descriptor to set the local updated variable. If
 	// no invoice update is performed, we can return early.
 	setID := (*SetID)(invoiceRef.SetID())
+	// If the setID is nil, we make sure we do NOT fetch any AMP HTLCs.
+	if setID == nil {
+		set := SetID(BlankPayAddr)
+		setID = &set
+	}
+
 	var updated bool
 	invoice, err := i.idb.UpdateInvoice(
 		context.Background(), invoiceRef, setID,
@@ -1014,7 +1020,13 @@ func (i *InvoiceRegistry) notifyExitHopHtlcLocked(
 	HtlcResolution, invoiceExpiry, error) {
 
 	invoiceRef := ctx.invoiceRef()
+
+	// If the setID is nil, we make sure we do NOT fetch any AMP HTLCs.
 	setID := (*SetID)(ctx.setID())
+	if setID == nil {
+		set := SetID(BlankPayAddr)
+		setID = &set
+	}
 
 	// We need to look up the current state of the invoice in order to send
 	// the previously accepted/settled HTLCs to the interceptor.
@@ -1374,7 +1386,14 @@ func (i *InvoiceRegistry) SettleHodlInvoice(ctx context.Context,
 
 	hash := preimage.Hash()
 	invoiceRef := InvoiceRefByHash(hash)
-	invoice, err := i.idb.UpdateInvoice(ctx, invoiceRef, nil, updateInvoice)
+
+	// AMP hold invoices are not supported however if we would support them
+	// we would need to fetch all AMP HTLCs here for the invoice.
+	setID := (*SetID)(nil)
+
+	invoice, err := i.idb.UpdateInvoice(
+		ctx, invoiceRef, setID, updateInvoice,
+	)
 	if err != nil {
 		log.Errorf("SettleHodlInvoice with preimage %v: %v",
 			preimage, err)
@@ -1458,10 +1477,13 @@ func (i *InvoiceRegistry) cancelInvoiceImpl(ctx context.Context,
 		}, nil
 	}
 
+	// If it's an AMP invoice we need to fetch all AMP HTLCs here so that
+	// we can cancel all of them which are in the accepted state.
+	setID := (*SetID)(nil)
 	invoiceRef := InvoiceRefByHash(payHash)
-
-	// We pass a nil setID which means no HTLCs will be read out.
-	invoice, err := i.idb.UpdateInvoice(ctx, invoiceRef, nil, updateInvoice)
+	invoice, err := i.idb.UpdateInvoice(
+		ctx, invoiceRef, setID, updateInvoice,
+	)
 
 	// Implement idempotency by returning success if the invoice was already
 	// canceled.
@@ -1487,8 +1509,8 @@ func (i *InvoiceRegistry) cancelInvoiceImpl(ctx context.Context,
 	// that are waiting for resolution. Any htlcs that were already canceled
 	// before, will be notified again. This isn't necessary but doesn't hurt
 	// either.
-	//
-	// TODO(ziggie): Also consider AMP HTLCs here.
+	// For AMP invoices we fetched all AMP HTLCs for all sub AMP invoices
+	// here so we can clean up all of them.
 	for key, htlc := range invoice.Htlcs {
 		if htlc.State != HtlcStateCanceled {
 			continue
@@ -1500,6 +1522,7 @@ func (i *InvoiceRegistry) cancelInvoiceImpl(ctx context.Context,
 			),
 		)
 	}
+
 	i.notifyClients(payHash, invoice, nil)
 
 	// Attempt to also delete the invoice if requested through the registry

invoices/invoiceregistry_test.go

@@ -117,6 +117,10 @@ func TestInvoiceRegistry(t *testing.T) {
 			name: "FailPartialAMPPayment",
 			test: testFailPartialAMPPayment,
 		},
+		{
+			name: "CancelAMPInvoicePendingHTLCs",
+			test: testCancelAMPInvoicePendingHTLCs,
+		},
 	}
 
 	makeKeyValueDB := func(t *testing.T) (invpkg.InvoiceDB,
@@ -2441,3 +2445,130 @@ func testFailPartialAMPPayment(t *testing.T,
 			"expected HTLC to be canceled")
 	}
 }
+
+// testCancelAMPInvoicePendingHTLCs tests the case where an AMP invoice is
+// canceled and the remaining HTLCs are also canceled so that no HTLCs are left
+// in the accepted state.
+func testCancelAMPInvoicePendingHTLCs(t *testing.T,
+	makeDB func(t *testing.T) (invpkg.InvoiceDB, *clock.TestClock)) {
+
+	t.Parallel()
+
+	ctx := newTestContext(t, nil, makeDB)
+	ctxb := context.Background()
+
+	const (
+		expiry    = uint32(testCurrentHeight + 20)
+		numShards = 4
+	)
+
+	var (
+		shardAmt = testInvoiceAmount / lnwire.MilliSatoshi(numShards)
+		payAddr  [32]byte
+	)
+	_, err := rand.Read(payAddr[:])
+	require.NoError(t, err)
+
+	// Create an AMP invoice we are going to pay via a multi-part payment.
+	ampInvoice := newInvoice(t, false, true)
+
+	// An AMP invoice is referenced by the payment address.
+	ampInvoice.Terms.PaymentAddr = payAddr
+
+	_, err = ctx.registry.AddInvoice(
+		ctxb, ampInvoice, testInvoicePaymentHash,
+	)
+	require.NoError(t, err)
+
+	htlcPayloadSet1 := &mockPayload{
+		mpp: record.NewMPP(testInvoiceAmount, payAddr),
+		// We are not interested in settling the AMP HTLC so we don't
+		// use valid shares.
+		amp: record.NewAMP([32]byte{1}, [32]byte{1}, 1),
+	}
+
+	// Send first HTLC which pays part of the invoice.
+	hodlChan1 := make(chan interface{}, 1)
+	resolution, err := ctx.registry.NotifyExitHopHtlc(
+		lntypes.Hash{1}, shardAmt, expiry, testCurrentHeight,
+		getCircuitKey(1), hodlChan1, nil, htlcPayloadSet1,
+	)
+	require.NoError(t, err)
+	require.Nil(t, resolution, "did not expect direct resolution")
+
+	htlcPayloadSet2 := &mockPayload{
+		mpp: record.NewMPP(testInvoiceAmount, payAddr),
+		// We are not interested in settling the AMP HTLC so we don't
+		// use valid shares.
+		amp: record.NewAMP([32]byte{2}, [32]byte{2}, 1),
+	}
+
+	// Send htlc 2 which should be added to the invoice as expected.
+	hodlChan2 := make(chan interface{}, 1)
+	resolution, err = ctx.registry.NotifyExitHopHtlc(
+		lntypes.Hash{2}, shardAmt, expiry, testCurrentHeight,
+		getCircuitKey(2), hodlChan2, nil, htlcPayloadSet2,
+	)
+	require.NoError(t, err)
+	require.Nil(t, resolution, "did not expect direct resolution")
+
+	require.Eventuallyf(t, func() bool {
+		inv, err := ctx.registry.LookupInvoice(
+			ctxb, testInvoicePaymentHash,
+		)
+		require.NoError(t, err)
+
+		return len(inv.Htlcs) == 2
+	}, testTimeout, time.Millisecond*100, "HTLCs not added to invoice")
+
+	// expire the invoice here.
+	ctx.clock.SetTime(testTime.Add(65 * time.Minute))
+
+	// Expect HLTC 1 to be canceled via the MPPTimeout fail resolution.
+	select {
+	case resolution := <-hodlChan1:
+		htlcResolution, _ := resolution.(invpkg.HtlcResolution)
+		_, ok := htlcResolution.(*invpkg.HtlcFailResolution)
+		require.True(
+			t, ok, "expected fail resolution, got: %T", resolution,
+		)
+
+	case <-time.After(testTimeout):
+		t.Fatal("timeout waiting for HTLC resolution")
+	}
+
+	// Expect HLTC 2 to be canceled via the MPPTimeout fail resolution.
+	select {
+	case resolution := <-hodlChan2:
+		htlcResolution, _ := resolution.(invpkg.HtlcResolution)
+		_, ok := htlcResolution.(*invpkg.HtlcFailResolution)
+		require.True(
+			t, ok, "expected fail resolution, got: %T", resolution,
+		)
+
+	case <-time.After(testTimeout):
+		t.Fatal("timeout waiting for HTLC resolution")
+	}
+
+	require.Eventuallyf(t, func() bool {
+		inv, err := ctx.registry.LookupInvoice(
+			ctxb, testInvoicePaymentHash,
+		)
+		require.NoError(t, err)
+
+		return inv.State == invpkg.ContractCanceled
+	}, testTimeout, time.Millisecond*100, "invoice not canceled")
+
+	// Fetch the invoice again and compare the number of cancelled HTLCs.
+	inv, err := ctx.registry.LookupInvoice(
+		ctxb, testInvoicePaymentHash,
+	)
+	require.NoError(t, err)
+
+	// Make sure all HTLCs are in the cancelled state.
+	require.Len(t, inv.Htlcs, 2)
+	for _, htlc := range inv.Htlcs {
+		require.Equal(t, invpkg.HtlcStateCanceled, htlc.State,
+			"expected HTLC to be canceled")
+	}
+}

invoices/update_invoice.go

@@ -567,6 +567,15 @@ func cancelInvoice(invoice *Invoice, hash *lntypes.Hash,
 			if err != nil {
 				return err
 			}
+
+			if invoice.IsAMP() {
+				err := cancelHtlcsAmp(
+					invoice, key, htlc, updater,
+				)
+				if err != nil {
+					return err
+				}
+			}
 		}
 	}