channeldb: add balance at height lookup

Add a balance at height lookup function which can be used to
obtain local/remote balance at a given height. The current in memory
commits and revocation log are used to source this information.

channeldb/channel.go

@@ -144,6 +144,15 @@ var (
 	// ErrChanBorked is returned when a caller attempts to mutate a borked
 	// channel.
 	ErrChanBorked = fmt.Errorf("cannot mutate borked channel")
+
+	// errLogEntryNotFound is returned when we cannot find a log entry at
+	// the height requested in the revocation log.
+	errLogEntryNotFound = fmt.Errorf("log entry not found")
+
+	// errHeightNotFound is returned when a query for channel balances at
+	// a height that we have not reached yet is made.
+	errHeightNotReached = fmt.Errorf("height requested greater than " +
+		"current commit height")
 )
 
 // ChannelType is an enum-like type that describes one of several possible
@@ -1391,6 +1400,44 @@ func (c *OpenChannel) UpdateCommitment(newCommitment *ChannelCommitment,
 	return nil
 }
 
+// BalancesAtHeight returns the local and remote balances on our commitment
+// transactions as of a given height.
+//
+// NOTE: these are our balances *after* subtracting the commitment fee and
+// anchor outputs.
+func (c *OpenChannel) BalancesAtHeight(height uint64) (lnwire.MilliSatoshi,
+	lnwire.MilliSatoshi, error) {
+
+	if height > c.LocalCommitment.CommitHeight &&
+		height > c.RemoteCommitment.CommitHeight {
+
+		return 0, 0, errHeightNotReached
+	}
+
+	// If our current commit is as the desired height, we can return our
+	// current balances.
+	if c.LocalCommitment.CommitHeight == height {
+		return c.LocalCommitment.LocalBalance,
+			c.LocalCommitment.RemoteBalance, nil
+	}
+
+	// If our current remote commit is at the desired height, we can return
+	// the current balances.
+	if c.RemoteCommitment.CommitHeight == height {
+		return c.RemoteCommitment.LocalBalance,
+			c.RemoteCommitment.RemoteBalance, nil
+	}
+
+	// If we are not currently on the height requested, we need to look up
+	// the previous height to obtain our balances at the given height.
+	commit, err := c.FindPreviousState(height)
+	if err != nil {
+		return 0, 0, err
+	}
+
+	return commit.LocalBalance, commit.RemoteBalance, nil
+}
+
 // HTLC is the on-disk representation of a hash time-locked contract. HTLCs are
 // contained within ChannelDeltas which encode the current state of the
 // commitment between state updates.
@@ -3160,7 +3207,7 @@ func fetchChannelLogEntry(log kvdb.ReadBucket,
 	logEntrykey := makeLogKey(updateNum)
 	commitBytes := log.Get(logEntrykey[:])
 	if commitBytes == nil {
-		return ChannelCommitment{}, fmt.Errorf("log entry not found")
+		return ChannelCommitment{}, errLogEntryNotFound
 	}
 
 	commitReader := bytes.NewReader(commitBytes)

channeldb/channel_test.go

@@ -10,6 +10,8 @@ import (
 	"runtime"
 	"testing"
 
+	"github.com/lightningnetwork/lnd/channeldb/kvdb"
+
 	"github.com/btcsuite/btcd/btcec"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
@@ -131,6 +133,25 @@ type testChannelParams struct {
 // default channel that is creates for testing.
 type testChannelOption func(params *testChannelParams)
 
+// channelCommitmentOption is an option which allows overwriting of the default
+// commitment height and balances. The local boolean can be used to set these
+// balances on the local or remote commit.
+func channelCommitmentOption(height uint64, localBalance,
+	remoteBalance lnwire.MilliSatoshi, local bool) testChannelOption {
+
+	return func(params *testChannelParams) {
+		if local {
+			params.channel.LocalCommitment.CommitHeight = height
+			params.channel.LocalCommitment.LocalBalance = localBalance
+			params.channel.LocalCommitment.RemoteBalance = remoteBalance
+		} else {
+			params.channel.RemoteCommitment.CommitHeight = height
+			params.channel.RemoteCommitment.LocalBalance = localBalance
+			params.channel.RemoteCommitment.RemoteBalance = remoteBalance
+		}
+	}
+}
+
 // pendingHeightOption is an option which can be used to set the height the
 // channel is marked as pending at.
 func pendingHeightOption(height uint32) testChannelOption {
@@ -1393,3 +1414,169 @@ func TestCloseChannelStatus(t *testing.T) {
 		t.Fatalf("channel should have status")
 	}
 }
+
+// TestBalanceAtHeight tests lookup of our local and remote balance at a given
+// height.
+func TestBalanceAtHeight(t *testing.T) {
+	const (
+		// Values that will be set on our current local commit in
+		// memory.
+		localHeight        = 2
+		localLocalBalance  = 1000
+		localRemoteBalance = 1500
+
+		// Values that will be set on our current remote commit in
+		// memory.
+		remoteHeight        = 3
+		remoteLocalBalance  = 2000
+		remoteRemoteBalance = 2500
+
+		// Values that will be written to disk in the revocation log.
+		oldHeight        = 0
+		oldLocalBalance  = 200
+		oldRemoteBalance = 300
+
+		// Heights to test error cases.
+		unknownHeight   = 1
+		unreachedHeight = 4
+	)
+
+	// putRevokedState is a helper function used to put commitments is
+	// the revocation log bucket to test lookup of balances at heights that
+	// are not our current height.
+	putRevokedState := func(c *OpenChannel, height uint64, local,
+		remote lnwire.MilliSatoshi) error {
+
+		err := kvdb.Update(c.Db, func(tx kvdb.RwTx) error {
+			chanBucket, err := fetchChanBucketRw(
+				tx, c.IdentityPub, &c.FundingOutpoint,
+				c.ChainHash,
+			)
+			if err != nil {
+				return err
+			}
+
+			logKey := revocationLogBucket
+			logBucket, err := chanBucket.CreateBucketIfNotExists(
+				logKey,
+			)
+			if err != nil {
+				return err
+			}
+
+			// Make a copy of our current commitment so we do not
+			// need to re-fill all the required fields and copy in
+			// our new desired values.
+			commit := c.LocalCommitment
+			commit.CommitHeight = height
+			commit.LocalBalance = local
+			commit.RemoteBalance = remote
+
+			return appendChannelLogEntry(logBucket, &commit)
+		})
+
+		return err
+	}
+
+	tests := []struct {
+		name                  string
+		targetHeight          uint64
+		expectedLocalBalance  lnwire.MilliSatoshi
+		expectedRemoteBalance lnwire.MilliSatoshi
+		expectedError         error
+	}{
+		{
+			name:                  "target is current local height",
+			targetHeight:          localHeight,
+			expectedLocalBalance:  localLocalBalance,
+			expectedRemoteBalance: localRemoteBalance,
+			expectedError:         nil,
+		},
+		{
+			name:                  "target is current remote height",
+			targetHeight:          remoteHeight,
+			expectedLocalBalance:  remoteLocalBalance,
+			expectedRemoteBalance: remoteRemoteBalance,
+			expectedError:         nil,
+		},
+		{
+			name:                  "need to lookup commit",
+			targetHeight:          oldHeight,
+			expectedLocalBalance:  oldLocalBalance,
+			expectedRemoteBalance: oldRemoteBalance,
+			expectedError:         nil,
+		},
+		{
+			name:                  "height not found",
+			targetHeight:          unknownHeight,
+			expectedLocalBalance:  0,
+			expectedRemoteBalance: 0,
+			expectedError:         errLogEntryNotFound,
+		},
+		{
+			name:                  "height not reached",
+			targetHeight:          unreachedHeight,
+			expectedLocalBalance:  0,
+			expectedRemoteBalance: 0,
+			expectedError:         errHeightNotReached,
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			cdb, cleanUp, err := makeTestDB()
+			if err != nil {
+				t.Fatalf("unable to make test database: %v",
+					err)
+			}
+			defer cleanUp()
+
+			// Create options to set the heights and balances of
+			// our local and remote commitments.
+			localCommitOpt := channelCommitmentOption(
+				localHeight, localLocalBalance,
+				localRemoteBalance, true,
+			)
+
+			remoteCommitOpt := channelCommitmentOption(
+				remoteHeight, remoteLocalBalance,
+				remoteRemoteBalance, false,
+			)
+
+			// Create an open channel.
+			channel := createTestChannel(
+				t, cdb, openChannelOption(),
+				localCommitOpt, remoteCommitOpt,
+			)
+
+			// Write an older commit to disk.
+			err = putRevokedState(channel, oldHeight,
+				oldLocalBalance, oldRemoteBalance)
+			if err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			}
+
+			local, remote, err := channel.BalancesAtHeight(
+				test.targetHeight,
+			)
+			if err != test.expectedError {
+				t.Fatalf("expected: %v, got: %v",
+					test.expectedError, err)
+			}
+
+			if local != test.expectedLocalBalance {
+				t.Fatalf("expected local: %v, got: %v",
+					test.expectedLocalBalance, local)
+			}
+
+			if remote != test.expectedRemoteBalance {
+				t.Fatalf("expected remote: %v, got: %v",
+					test.expectedRemoteBalance, remote)
+			}
+		})
+	}
+}

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/btcsuite/btcwallet/walletdb v1.2.0
 	github.com/btcsuite/btcwallet/wtxmgr v1.0.0
 	github.com/btcsuite/fastsha256 v0.0.0-20160815193821-637e65642941
+	github.com/coreos/bbolt v1.3.3
 	github.com/davecgh/go-spew v1.1.1
 	github.com/go-errors/errors v1.0.1
 	github.com/golang/protobuf v1.3.1