Merge 1c3edee into 09bb9db

channeldb/db.go

@@ -144,6 +144,12 @@ var (
 			number:    14,
 			migration: mig.CreateTLB(payAddrIndexBucket),
 		},
+		{
+			// Initialize set id index so that invoices can be
+			// queried by individual htlc sets.
+			number:    15,
+			migration: mig.CreateTLB(setIDIndexBucket),
+		},
 	}
 
 	// Big endian is the preferred byte order, due to cursor scans over
@@ -257,6 +263,7 @@ var topLevelBuckets = [][]byte{
 	fwdPackagesKey,
 	invoiceBucket,
 	payAddrIndexBucket,
+	setIDIndexBucket,
 	nodeInfoBucket,
 	nodeBucket,
 	edgeBucket,

channeldb/invoice_test.go

@@ -1123,6 +1123,69 @@ func TestCustomRecords(t *testing.T) {
 	}
 }
 
+// TestInvoiceHtlcAMPFields asserts that the set id and preimage fields are
+// properly recorded when updating an invoice.
+func TestInvoiceHtlcAMPFields(t *testing.T) {
+	t.Run("amp", func(t *testing.T) {
+		testInvoiceHtlcAMPFields(t, true)
+	})
+	t.Run("no amp", func(t *testing.T) {
+		testInvoiceHtlcAMPFields(t, false)
+	})
+}
+
+func testInvoiceHtlcAMPFields(t *testing.T, isAMP bool) {
+	db, cleanUp, err := makeTestDB()
+	defer cleanUp()
+	assert.Nil(t, err)
+
+	testInvoice, err := randInvoice(1000)
+	assert.Nil(t, err)
+
+	payHash := testInvoice.Terms.PaymentPreimage.Hash()
+	_, err = db.AddInvoice(testInvoice, payHash)
+	assert.Nil(t, err)
+
+	// Accept an htlc with custom records on this invoice.
+	key := CircuitKey{ChanID: lnwire.NewShortChanIDFromInt(1), HtlcID: 4}
+	records := make(map[uint64][]byte)
+
+	var (
+		amp      *record.AMP
+		preimage *lntypes.Preimage
+	)
+	if isAMP {
+		amp = record.NewAMP([32]byte{1}, [32]byte{2}, 3)
+		preimage = &lntypes.Preimage{4}
+	}
+
+	ref := InvoiceRefByHash(payHash)
+	_, err = db.UpdateInvoice(ref,
+		func(invoice *Invoice) (*InvoiceUpdateDesc, error) {
+			return &InvoiceUpdateDesc{
+				AddHtlcs: map[CircuitKey]*HtlcAcceptDesc{
+					key: {
+						Amt:           500,
+						AMP:           amp,
+						Preimage:      preimage,
+						CustomRecords: records,
+					},
+				},
+			}, nil
+		},
+	)
+	assert.Nil(t, err)
+
+	// Retrieve the invoice from that database and verify that the AMP
+	// fields are as expected.
+	dbInvoice, err := db.LookupInvoice(ref)
+	assert.Nil(t, err)
+
+	assert.Equal(t, 1, len(dbInvoice.Htlcs))
+	assert.Equal(t, amp, dbInvoice.Htlcs[key].AMP)
+	assert.Equal(t, preimage, dbInvoice.Htlcs[key].Preimage)
+}
+
 // TestInvoiceRef asserts that the proper identifiers are returned from an
 // InvoiceRef depending on the constructor used.
 func TestInvoiceRef(t *testing.T) {
@@ -1141,3 +1204,301 @@ func TestInvoiceRef(t *testing.T) {
 	assert.Equal(t, payHash, refByHashAndAddr.PayHash())
 	assert.Equal(t, &payAddr, refByHashAndAddr.PayAddr())
 }
+
+// TestHTLCSet asserts that HTLCSet returns the proper set of accepted HTLCs
+// that can be considered for settlement. It asserts that MPP and AMP HTLCs do
+// not comingle, and also that HTLCs with disjoint set ids appear in different
+// sets.
+func TestHTLCSet(t *testing.T) {
+	inv := &Invoice{
+		Htlcs: make(map[CircuitKey]*InvoiceHTLC),
+	}
+
+	// Construct two distinct set id's, in this test we'll also track the
+	// nil set id as a third group.
+	setID1 := &[32]byte{1}
+	setID2 := &[32]byte{2}
+
+	// Create the expected htlc sets for each group, these will be updated
+	// as the invoice is modified.
+	expSetNil := make(map[CircuitKey]*InvoiceHTLC)
+	expSet1 := make(map[CircuitKey]*InvoiceHTLC)
+	expSet2 := make(map[CircuitKey]*InvoiceHTLC)
+
+	checkHTLCSets := func() {
+		assert.Equal(t, expSetNil, inv.HTLCSet(nil))
+		assert.Equal(t, expSet1, inv.HTLCSet(setID1))
+		assert.Equal(t, expSet2, inv.HTLCSet(setID2))
+	}
+
+	// All HTLC sets should be empty initially.
+	checkHTLCSets()
+
+	// Add the following sequence of HTLCs to the invoice, sanity checking
+	// all three HTLC sets after each transition. This sequence asserts:
+	//   - both nil and non-nil set ids can have multiple htlcs.
+	//   - there may be distinct htlc sets with non-nil set ids.
+	//   - only accepted htlcs are returned as part of the set.
+	htlcs := []struct {
+		setID *[32]byte
+		state HtlcState
+	}{
+		{nil, HtlcStateAccepted},
+		{nil, HtlcStateAccepted},
+		{setID1, HtlcStateAccepted},
+		{setID1, HtlcStateAccepted},
+		{setID2, HtlcStateAccepted},
+		{setID2, HtlcStateAccepted},
+		{nil, HtlcStateCanceled},
+		{setID1, HtlcStateCanceled},
+		{setID2, HtlcStateCanceled},
+		{nil, HtlcStateSettled},
+		{setID1, HtlcStateSettled},
+		{setID2, HtlcStateSettled},
+	}
+
+	for i, h := range htlcs {
+		var amp *record.AMP
+		if h.setID != nil {
+			amp = record.NewAMP([32]byte{0}, *h.setID, 0)
+		}
+
+		// Add the HTLC to the invoice's set of HTLCs.
+		key := CircuitKey{HtlcID: uint64(i)}
+		htlc := &InvoiceHTLC{
+			AMP:   amp,
+			State: h.state,
+		}
+		inv.Htlcs[key] = htlc
+
+		// Update our expected htlc set if the htlc is accepted,
+		// otherwise it shouldn't be reflected.
+		if h.state == HtlcStateAccepted {
+			switch h.setID {
+			case nil:
+				expSetNil[key] = htlc
+			case setID1:
+				expSet1[key] = htlc
+			case setID2:
+				expSet2[key] = htlc
+			default:
+				t.Fatalf("unexpected set id")
+			}
+		}
+
+		checkHTLCSets()
+	}
+}
+
+// TestAddInvoiceWithHTLCs asserts that you can't insert an invoice that already
+// has HTLCs.
+func TestAddInvoiceWithHTLCs(t *testing.T) {
+	db, cleanUp, err := makeTestDB()
+	defer cleanUp()
+	assert.Nil(t, err)
+
+	testInvoice, err := randInvoice(1000)
+	assert.Nil(t, err)
+
+	key := CircuitKey{HtlcID: 1}
+	testInvoice.Htlcs[key] = &InvoiceHTLC{}
+
+	payHash := testInvoice.Terms.PaymentPreimage.Hash()
+	_, err = db.AddInvoice(testInvoice, payHash)
+	assert.Equal(t, ErrInvoiceHasHtlcs, err)
+}
+
+// TestSetIDIndex asserts that the set id index properly adds new invoices as we
+// accept HTLCs, that they can be queried by their set id after accepting, and
+// that invoices with duplicate set ids are disallowed.
+func TestSetIDIndex(t *testing.T) {
+	db, cleanUp, err := makeTestDB()
+	defer cleanUp()
+	assert.Nil(t, err)
+
+	// We'll start out by creating an invoice and writing it to the DB.
+	amt := lnwire.NewMSatFromSatoshis(1000)
+	invoice, err := randInvoice(amt)
+	assert.Nil(t, err)
+
+	payHash := invoice.Terms.PaymentPreimage.Hash()
+	_, err = db.AddInvoice(invoice, payHash)
+	assert.Nil(t, err)
+
+	setID := &[32]byte{1}
+
+	// Update the invoice with an accepted HTLC that also accepts the
+	// invoice.
+	ref := InvoiceRefByHashAndAddr(payHash, invoice.Terms.PaymentAddr)
+	dbInvoice, err := db.UpdateInvoice(ref, updateAcceptAMPHtlc(0, amt, setID, true))
+	assert.Nil(t, err)
+
+	// We'll update what we expect the accepted invoice to be so that our
+	// comparison below has the correct assumption.
+	invoice.State = ContractAccepted
+	invoice.AmtPaid = amt
+	invoice.SettleDate = dbInvoice.SettleDate
+	invoice.Htlcs = map[CircuitKey]*InvoiceHTLC{
+		{HtlcID: 0}: {
+			Amt:           amt,
+			AcceptTime:    time.Unix(1, 0),
+			ResolveTime:   time.Time{},
+			State:         HtlcStateAccepted,
+			CustomRecords: make(record.CustomSet),
+			AMP:           record.NewAMP([32]byte{}, *setID, 0),
+			Preimage:      invoice.Terms.PaymentPreimage,
+		},
+	}
+
+	// We should get back the exact same invoice that we just inserted.
+	assert.Equal(t, invoice, dbInvoice)
+
+	// Now lookup the invoice by set id and see that we get the same one.
+	refBySetID := InvoiceRefBySetID(*setID)
+	dbInvoiceBySetID, err := db.LookupInvoice(refBySetID)
+	assert.Nil(t, err)
+	assert.Equal(t, invoice, &dbInvoiceBySetID)
+
+	// Trying to accept an HTLC to a different invoice, but using the same
+	// set id should fail.
+	invoice2, err := randInvoice(amt)
+	assert.Nil(t, err)
+
+	payHash2 := invoice2.Terms.PaymentPreimage.Hash()
+	_, err = db.AddInvoice(invoice2, payHash2)
+	assert.Nil(t, err)
+
+	ref2 := InvoiceRefByHashAndAddr(payHash2, invoice2.Terms.PaymentAddr)
+	_, err = db.UpdateInvoice(ref2, updateAcceptAMPHtlc(0, amt, setID, true))
+	assert.Equal(t, ErrDuplicateSetID{*setID}, err)
+
+	// Now, begin constructing a second htlc set under a different set id.
+	// This set will contain two distinct HTLCs.
+	setID2 := &[32]byte{2}
+
+	_, err = db.UpdateInvoice(ref, updateAcceptAMPHtlc(1, amt, setID2, false))
+	assert.Nil(t, err)
+	dbInvoice, err = db.UpdateInvoice(ref, updateAcceptAMPHtlc(2, amt, setID2, false))
+	assert.Nil(t, err)
+
+	// We'll update what we expect the settle invoice to be so that our
+	// comparison below has the correct assumption.
+	invoice.State = ContractAccepted
+	invoice.AmtPaid += 2 * amt
+	invoice.SettleDate = dbInvoice.SettleDate
+	invoice.Htlcs = map[CircuitKey]*InvoiceHTLC{
+		{HtlcID: 0}: {
+			Amt:           amt,
+			AcceptTime:    time.Unix(1, 0),
+			ResolveTime:   time.Time{},
+			State:         HtlcStateAccepted,
+			CustomRecords: make(record.CustomSet),
+			AMP:           record.NewAMP([32]byte{}, *setID, 0),
+			Preimage:      invoice.Terms.PaymentPreimage,
+		},
+		{HtlcID: 1}: {
+			Amt:           amt,
+			AcceptTime:    time.Unix(1, 0),
+			ResolveTime:   time.Time{},
+			State:         HtlcStateAccepted,
+			CustomRecords: make(record.CustomSet),
+			AMP:           record.NewAMP([32]byte{}, *setID2, 0),
+			Preimage:      invoice.Terms.PaymentPreimage,
+		},
+		{HtlcID: 2}: {
+			Amt:           amt,
+			AcceptTime:    time.Unix(1, 0),
+			ResolveTime:   time.Time{},
+			State:         HtlcStateAccepted,
+			CustomRecords: make(record.CustomSet),
+			AMP:           record.NewAMP([32]byte{}, *setID2, 0),
+			Preimage:      invoice.Terms.PaymentPreimage,
+		},
+	}
+
+	// We should get back the exact same invoice that we just inserted.
+	assert.Equal(t, invoice, dbInvoice)
+
+	// Now lookup the invoice by second set id and see that we get the same
+	// index, including the htlcs under the first set id.
+	refBySetID = InvoiceRefBySetID(*setID2)
+	dbInvoiceBySetID, err = db.LookupInvoice(refBySetID)
+	assert.Nil(t, err)
+	assert.Equal(t, invoice, &dbInvoiceBySetID)
+
+	// Now settle the first htlc set, asserting that the two htlcs with set
+	// id 2 get canceled as a result.
+	dbInvoice, err = db.UpdateInvoice(ref, getUpdateInvoiceAMPSettle(setID))
+	assert.Nil(t, err)
+
+	invoice.State = ContractSettled
+	invoice.SettleDate = dbInvoice.SettleDate
+	invoice.SettleIndex = 1
+	invoice.AmtPaid = amt
+	invoice.Htlcs[CircuitKey{HtlcID: 0}].ResolveTime = time.Unix(1, 0)
+	invoice.Htlcs[CircuitKey{HtlcID: 0}].State = HtlcStateSettled
+	invoice.Htlcs[CircuitKey{HtlcID: 1}].ResolveTime = time.Unix(1, 0)
+	invoice.Htlcs[CircuitKey{HtlcID: 1}].State = HtlcStateCanceled
+	invoice.Htlcs[CircuitKey{HtlcID: 2}].ResolveTime = time.Unix(1, 0)
+	invoice.Htlcs[CircuitKey{HtlcID: 2}].State = HtlcStateCanceled
+	assert.Equal(t, invoice, dbInvoice)
+
+	// Lastly, querying for an unknown set id should fail.
+	refUnknownSetID := InvoiceRefBySetID([32]byte{})
+	_, err = db.LookupInvoice(refUnknownSetID)
+	assert.Equal(t, ErrInvoiceNotFound, err)
+}
+
+// updateAcceptAMPHtlc returns an invoice update callback that, when called,
+// settles the invoice with the given amount.
+func updateAcceptAMPHtlc(id uint64, amt lnwire.MilliSatoshi,
+	setID *[32]byte, accept bool) InvoiceUpdateCallback {
+
+	return func(invoice *Invoice) (*InvoiceUpdateDesc, error) {
+		if invoice.State == ContractSettled {
+			return nil, ErrInvoiceAlreadySettled
+		}
+
+		noRecords := make(record.CustomSet)
+
+		var state *InvoiceStateUpdateDesc
+		if accept {
+			state = &InvoiceStateUpdateDesc{
+				Preimage: invoice.Terms.PaymentPreimage,
+				NewState: ContractAccepted,
+			}
+		}
+
+		update := &InvoiceUpdateDesc{
+			State: state,
+			AddHtlcs: map[CircuitKey]*HtlcAcceptDesc{
+				{HtlcID: id}: {
+					Amt:           amt,
+					CustomRecords: noRecords,
+					AMP:           record.NewAMP([32]byte{}, *setID, 0),
+					Preimage:      invoice.Terms.PaymentPreimage,
+				},
+			},
+		}
+
+		return update, nil
+	}
+}
+
+func getUpdateInvoiceAMPSettle(setID *[32]byte) InvoiceUpdateCallback {
+	return func(invoice *Invoice) (*InvoiceUpdateDesc, error) {
+		if invoice.State == ContractSettled {
+			return nil, ErrInvoiceAlreadySettled
+		}
+
+		update := &InvoiceUpdateDesc{
+			State: &InvoiceStateUpdateDesc{
+				Preimage: nil,
+				NewState: ContractSettled,
+				SetID:    setID,
+			},
+		}
+
+		return update, nil
+	}
+}