channeldb+migration: add migration test for patching balances

channeldb/migration25/migration_test.go

@@ -0,0 +1,391 @@
+package migration25
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+
+	"github.com/btcsuite/btcd/btcec/v2"
+	"github.com/btcsuite/btcd/btcutil"
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/btcsuite/btcd/wire"
+
+	lnwire "github.com/lightningnetwork/lnd/channeldb/migration/lnwire21"
+	mig24 "github.com/lightningnetwork/lnd/channeldb/migration24"
+	mig "github.com/lightningnetwork/lnd/channeldb/migration_01_to_11"
+	"github.com/lightningnetwork/lnd/channeldb/migtest"
+	"github.com/lightningnetwork/lnd/kvdb"
+)
+
+var (
+	// Create dummy values to be stored in db.
+	dummyPrivKey, _ = btcec.NewPrivateKey()
+	dummyPubKey     = dummyPrivKey.PubKey()
+	dummySig        = []byte{1, 2, 3}
+	dummyTx         = &wire.MsgTx{
+		Version: 1,
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: wire.OutPoint{
+					Hash:  chainhash.Hash{},
+					Index: 0xffffffff,
+				},
+				Sequence: 0xffffffff,
+			},
+		},
+		TxOut: []*wire.TxOut{
+			{
+				Value: 5000000000,
+			},
+		},
+		LockTime: 5,
+	}
+
+	dummyOp = wire.OutPoint{
+		Hash:  chainhash.Hash{},
+		Index: 9,
+	}
+	dummyHTLC = mig.HTLC{
+		Signature:     dummySig,
+		RHash:         [32]byte{},
+		Amt:           100_000,
+		RefundTimeout: 731583,
+		OutputIndex:   1,
+		Incoming:      true,
+		HtlcIndex:     1,
+		LogIndex:      1,
+	}
+
+	// ourAmt and theirAmt are the initial balances found in the local
+	// channel commitment at height 0.
+	ourAmt   = lnwire.MilliSatoshi(500_000)
+	theirAmt = lnwire.MilliSatoshi(1000_000)
+
+	// ourAmtRevoke and theirAmtRevoke are the initial balances found in
+	// the revocation log at height 0.
+	//
+	// NOTE: they are made differently such that we can easily check the
+	// source when patching the balances.
+	ourAmtRevoke   = lnwire.MilliSatoshi(501_000)
+	theirAmtRevoke = lnwire.MilliSatoshi(1001_000)
+
+	// remoteCommit0 is the channel commitment at commit height 0. This is
+	// also the revocation log we will use to patch the initial balances.
+	remoteCommit0 = mig.ChannelCommitment{
+		LocalBalance:  ourAmtRevoke,
+		RemoteBalance: theirAmtRevoke,
+		CommitFee:     btcutil.Amount(1),
+		FeePerKw:      btcutil.Amount(1),
+		CommitTx:      dummyTx,
+		CommitSig:     dummySig,
+		Htlcs:         []mig.HTLC{},
+	}
+
+	// localCommit0 is the channel commitment at commit height 0. This is
+	// the channel commitment we will use to patch the initial balances
+	// when there are no revocation logs.
+	localCommit0 = mig.ChannelCommitment{
+		LocalBalance:  ourAmt,
+		RemoteBalance: theirAmt,
+		CommitFee:     btcutil.Amount(1),
+		FeePerKw:      btcutil.Amount(1),
+		CommitTx:      dummyTx,
+		CommitSig:     dummySig,
+		Htlcs:         []mig.HTLC{},
+	}
+
+	// remoteCommit1 and localCommit1 are the channel commitment at commit
+	// height 1.
+	remoteCommit1 = mig.ChannelCommitment{
+		CommitHeight:    1,
+		LocalLogIndex:   1,
+		LocalHtlcIndex:  1,
+		RemoteLogIndex:  1,
+		RemoteHtlcIndex: 1,
+		LocalBalance:    ourAmt - dummyHTLC.Amt,
+		RemoteBalance:   theirAmt + dummyHTLC.Amt,
+		CommitFee:       btcutil.Amount(1),
+		FeePerKw:        btcutil.Amount(1),
+		CommitTx:        dummyTx,
+		CommitSig:       dummySig,
+		Htlcs:           []mig.HTLC{dummyHTLC},
+	}
+	localCommit1 = mig.ChannelCommitment{
+		CommitHeight:    1,
+		LocalLogIndex:   1,
+		LocalHtlcIndex:  1,
+		RemoteLogIndex:  1,
+		RemoteHtlcIndex: 1,
+		LocalBalance:    ourAmt - dummyHTLC.Amt,
+		RemoteBalance:   theirAmt + dummyHTLC.Amt,
+		CommitFee:       btcutil.Amount(1),
+		FeePerKw:        btcutil.Amount(1),
+		CommitTx:        dummyTx,
+		CommitSig:       dummySig,
+		Htlcs:           []mig.HTLC{dummyHTLC},
+	}
+
+	// openChannel0 is the OpenChannel at commit height 0. When this
+	// variable is used, we expect to patch the initial balances from its
+	// commitments.
+	openChannel0 = &OpenChannel{
+		OpenChannel: mig.OpenChannel{
+			IdentityPub:      dummyPubKey,
+			FundingOutpoint:  dummyOp,
+			LocalCommitment:  localCommit0,
+			RemoteCommitment: remoteCommit0,
+		},
+	}
+
+	// openChannel1 is the OpenChannel at commit height 1. When this
+	// variable is used, we expect to patch the initial balances from the
+	// remote commitment at height 0.
+	openChannel1 = &OpenChannel{
+		OpenChannel: mig.OpenChannel{
+			IdentityPub:      dummyPubKey,
+			FundingOutpoint:  dummyOp,
+			LocalCommitment:  localCommit1,
+			RemoteCommitment: remoteCommit1,
+		},
+	}
+)
+
+// TestMigrateInitialBalances checks that the proper initial balances are
+// patched to the channel info.
+func TestMigrateInitialBalances(t *testing.T) {
+	testCases := []struct {
+		name                string
+		beforeMigrationFunc func(kvdb.RwTx) error
+		afterMigrationFunc  func(kvdb.RwTx) error
+		shouldFail          bool
+	}{
+		{
+			// Test that we patch the initial balances using the
+			// revocation log.
+			name: "patch balance from revoke log",
+			beforeMigrationFunc: genBeforeMigration(
+				openChannel1, &remoteCommit0,
+			),
+			afterMigrationFunc: genAfterMigration(
+				ourAmtRevoke, theirAmtRevoke, openChannel1,
+			),
+		},
+		{
+			// Test that we patch the initial balances using the
+			// channel's local commitment since at height 0,
+			// balances found in LocalCommitment reflect the
+			// initial balances.
+			name: "patch balance from local commit",
+			beforeMigrationFunc: genBeforeMigration(
+				openChannel0, nil,
+			),
+			afterMigrationFunc: genAfterMigration(
+				ourAmt, theirAmt, openChannel0,
+			),
+		},
+		{
+			// Test that we patch the initial balances using the
+			// channel's local commitment even when there is a
+			// revocation log available.
+			name: "patch balance from local commit only",
+			beforeMigrationFunc: genBeforeMigration(
+				openChannel0, &remoteCommit0,
+			),
+			afterMigrationFunc: genAfterMigration(
+				ourAmt, theirAmt, openChannel0,
+			),
+		},
+		{
+			// Test that when there is no revocation log the
+			// migration would fail.
+			name: "patch balance error on no revoke log",
+			beforeMigrationFunc: genBeforeMigration(
+				// Use nil to specify no revocation log will be
+				// created.
+				openChannel1, nil,
+			),
+			afterMigrationFunc: genAfterMigration(
+				// Use nil to specify skipping the
+				// afterMigrationFunc.
+				0, 0, nil,
+			),
+			shouldFail: true,
+		},
+		{
+			// Test that when the saved revocation log is not what
+			// we want the migration would fail.
+			name: "patch balance error on wrong revoke log",
+			beforeMigrationFunc: genBeforeMigration(
+				// Use the revocation log with the wrong
+				// height.
+				openChannel1, &remoteCommit1,
+			),
+			afterMigrationFunc: genAfterMigration(
+				// Use nil to specify skipping the
+				// afterMigrationFunc.
+				0, 0, nil,
+			),
+			shouldFail: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			migtest.ApplyMigration(
+				t,
+				tc.beforeMigrationFunc,
+				tc.afterMigrationFunc,
+				MigrateInitialBalances,
+				tc.shouldFail,
+			)
+		})
+	}
+}
+
+func genBeforeMigration(c *OpenChannel,
+	commit *mig.ChannelCommitment) func(kvdb.RwTx) error {
+
+	return func(tx kvdb.RwTx) error {
+		if c.InitialLocalBalance != 0 {
+			return fmt.Errorf("non zero initial local amount")
+		}
+
+		if c.InitialRemoteBalance != 0 {
+			return fmt.Errorf("non zero initial local amount")
+		}
+
+		// Create the channel bucket.
+		chanBucket, err := createChanBucket(tx, c)
+		if err != nil {
+			return err
+		}
+
+		// Save the channel info using legacy format.
+		if err := putChanInfo(chanBucket, c, true); err != nil {
+			return err
+		}
+
+		// Save the channel commitments.
+		if err := PutChanCommitments(chanBucket, c); err != nil {
+			return err
+		}
+
+		// If we have a remote commitment, save it as our revocation
+		// log.
+		if commit != nil {
+			err := putChannelLogEntryLegacy(chanBucket, commit)
+			if err != nil {
+				return err
+			}
+		}
+
+		return nil
+	}
+}
+
+func genAfterMigration(ourAmt, theirAmt lnwire.MilliSatoshi,
+	c *OpenChannel) func(kvdb.RwTx) error {
+
+	return func(tx kvdb.RwTx) error {
+		// If the passed OpenChannel is nil, we will skip the
+		// afterMigrationFunc as it indicates an error is expected
+		// during the migration.
+		if c == nil {
+			return nil
+		}
+
+		chanBucket, err := fetchChanBucket(tx, c)
+		if err != nil {
+			return err
+		}
+
+		newChan := &OpenChannel{}
+
+		// Fetch the channel info using the new format.
+		err = fetchChanInfo(chanBucket, newChan, false)
+		if err != nil {
+			return err
+		}
+
+		// Check our initial amount is correct.
+		if newChan.InitialLocalBalance != ourAmt {
+			return fmt.Errorf("wrong local balance, got %d, "+
+				"want %d", newChan.InitialLocalBalance, ourAmt)
+		}
+
+		// Check their initial amount is correct.
+		if newChan.InitialRemoteBalance != theirAmt {
+			return fmt.Errorf("wrong remote balance, got %d, "+
+				"want %d", newChan.InitialRemoteBalance,
+				theirAmt)
+		}
+
+		// We also check the relevant channel info fields stay the
+		// same.
+		if !newChan.IdentityPub.IsEqual(c.IdentityPub) {
+			return fmt.Errorf("wrong IdentityPub")
+		}
+		if newChan.FundingOutpoint != c.FundingOutpoint {
+			return fmt.Errorf("wrong FundingOutpoint")
+		}
+
+		return nil
+	}
+}
+
+func createChanBucket(tx kvdb.RwTx, c *OpenChannel) (kvdb.RwBucket, error) {
+	// First fetch the top level bucket which stores all data related to
+	// current, active channels.
+	openChanBucket, err := tx.CreateTopLevelBucket(openChannelBucket)
+	if err != nil {
+		return nil, err
+	}
+
+	// Within this top level bucket, fetch the bucket dedicated to storing
+	// open channel data specific to the remote node.
+	nodePub := c.IdentityPub.SerializeCompressed()
+	nodeChanBucket, err := openChanBucket.CreateBucketIfNotExists(nodePub)
+	if err != nil {
+		return nil, err
+	}
+
+	// We'll then recurse down an additional layer in order to fetch the
+	// bucket for this particular chain.
+	chainBucket, err := nodeChanBucket.CreateBucketIfNotExists(
+		c.ChainHash[:],
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	var chanPointBuf bytes.Buffer
+	err = mig.WriteOutpoint(&chanPointBuf, &c.FundingOutpoint)
+	if err != nil {
+		return nil, err
+	}
+
+	// With the bucket for the node fetched, we can now go down another
+	// level, creating the bucket for this channel itself.
+	return chainBucket.CreateBucketIfNotExists(chanPointBuf.Bytes())
+}
+
+// putChannelLogEntryLegacy saves an old format revocation log to the bucket.
+func putChannelLogEntryLegacy(chanBucket kvdb.RwBucket,
+	commit *mig.ChannelCommitment) error {
+
+	logBucket, err := chanBucket.CreateBucketIfNotExists(
+		revocationLogBucketLegacy,
+	)
+	if err != nil {
+		return err
+	}
+
+	var b bytes.Buffer
+	if err := mig.SerializeChanCommit(&b, commit); err != nil {
+		return err
+	}
+
+	logEntrykey := mig24.MakeLogKey(commit.CommitHeight)
+	return logBucket.Put(logEntrykey[:], b.Bytes())
+}