followups: Configurable Forwarding Limit and Stub Improvements

db.go

@@ -96,7 +96,7 @@ type Db struct {
 	fwdHistoryLimit int
 }
 
-func NewDb(dbPath string, opts ...func(*Db)) (*Db, error) {
+func NewDb(ctx context.Context, dbPath string, fwdHistoryLimit int) (*Db, error) {
 	const busyTimeoutMs = 5000
 
 	dsn := dbPath + fmt.Sprintf("?_pragma=busy_timeout=%d", busyTimeoutMs)
@@ -116,10 +116,13 @@ func NewDb(dbPath string, opts ...func(*Db)) (*Db, error) {
 
 	database := &Db{
 		db:              db,
-		fwdHistoryLimit: defaultFwdHistoryLimit,
+		fwdHistoryLimit: fwdHistoryLimit,
 	}
-	for _, opt := range opts {
-		opt(database)
+
+	// Perform a once-off cleanup of the records in the db to update to a potential
+	// change in limit value.
+	if err := database.limitHTLCRecords(ctx); err != nil {
+		return nil, err
 	}
 
 	return database, nil
@@ -244,6 +247,12 @@ type HtlcInfo struct {
 func (d *Db) RecordHtlcResolution(ctx context.Context,
 	htlc *HtlcInfo) error {
 
+	// If the database is configured to not store any records, save the hassle of
+	// writing and deleting a record by returning early.
+	if d.fwdHistoryLimit == 0 {
+		return nil
+	}
+
 	if err := d.insertHtlcResolution(ctx, htlc); err != nil {
 		return err
 	}

db_test.go

@@ -10,11 +10,11 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func setupTestDb(t *testing.T, dbOpts ...func(*Db)) (*Db, func()) {
+func setupTestDb(t *testing.T, fwdingHistoryLimit int) (*Db, func()) {
 	file, err := os.CreateTemp("", "test_db_")
 	require.NoError(t, err)
 
-	db, err := NewDb(file.Name(), dbOpts...)
+	db, err := NewDb(context.Background(), file.Name(), fwdingHistoryLimit)
 	require.NoError(t, err)
 
 	return db, func() {
@@ -24,7 +24,7 @@ func setupTestDb(t *testing.T, dbOpts ...func(*Db)) (*Db, func()) {
 
 func TestDb(t *testing.T) {
 	ctx := context.Background()
-	db, cleanup := setupTestDb(t)
+	db, cleanup := setupTestDb(t, defaultFwdHistoryLimit)
 	defer cleanup()
 
 	expectedDefaultLimit := Limit{
@@ -68,18 +68,12 @@ func TestDb(t *testing.T) {
 	defer db.Close()
 }
 
-func dbWithCustomForwardingHistoryLimit(limit int) func(d *Db) {
-	return func(d *Db) {
-		d.fwdHistoryLimit = limit
-	}
-}
-
 func TestDbForwardingHistory(t *testing.T) {
 	limit := 20
 
 	// Create a test DB that will limit to 10 forwarding history records.
 	ctx := context.Background()
-	db, cleanup := setupTestDb(t, dbWithCustomForwardingHistoryLimit(limit))
+	db, cleanup := setupTestDb(t, limit)
 	defer cleanup()
 
 	// Insert HTLCs just up until our limit.
@@ -122,3 +116,41 @@ func testHtlc(i uint64) *HtlcInfo {
 		},
 	}
 }
+
+func TestDbNoForwardingHistory(t *testing.T) {
+	ctx := context.Background()
+	db, cleanup := setupTestDb(t, 0)
+	defer cleanup()
+
+	htlc := testHtlc(1)
+	require.NoError(t, db.RecordHtlcResolution(ctx, htlc))
+
+	fwds, err := db.ListForwardingHistory(ctx, time.Time{}, time.Unix(1000000, 0))
+	require.NoError(t, err)
+	require.Len(t, fwds, 0)
+}
+
+func TestForwadingHistoryDelete(t *testing.T) {
+	// Create a db that will store HTLCs.
+	ctx := context.Background()
+	db, cleanup := setupTestDb(t, 5)
+	defer cleanup()
+
+	// Write a test HTLC and assert that it's stored.
+	htlc := testHtlc(1)
+	require.NoError(t, db.RecordHtlcResolution(ctx, htlc))
+
+	fwds, err := db.ListForwardingHistory(ctx, time.Time{}, time.Unix(1000000, 0))
+	require.NoError(t, err)
+	require.Len(t, fwds, 1)
+
+	// Modify the db to have a zero limit on forwarding history. We don't recreate
+	// the test db because it would re-create the file. Run limitHTLCRecords once
+	// (as we would on NewDb) to assert that we clean up our records.
+	db.fwdHistoryLimit = 0
+	require.NoError(t, db.limitHTLCRecords(ctx))
+
+	fwds, err = db.ListForwardingHistory(ctx, time.Time{}, time.Unix(1000000, 0))
+	require.NoError(t, err)
+	require.Len(t, fwds, 0)
+}

lndclient.go

@@ -255,6 +255,49 @@ func (l *lndclientGrpc) listChannels() (map[uint64]*channel, error) {
 	return chans, nil
 }
 
+func (l *lndclientGrpc) listClosedChannels() (map[uint64]*channel, error) {
+	ctx, cancel := context.WithTimeout(ctxb, rpcTimeout)
+	defer cancel()
+
+	resp, err := l.main.ClosedChannels(ctx, &lnrpc.ClosedChannelsRequest{})
+	if err != nil {
+		return nil, err
+	}
+
+	chans := make(map[uint64]*channel)
+	for _, rpcChan := range resp.Channels {
+		peer, err := route.NewVertexFromStr(rpcChan.RemotePubkey)
+		if err != nil {
+			return nil, err
+		}
+
+		channel := &channel{
+			peer: peer,
+		}
+
+		// LND didn't always store who initiated the channel, so in some cases
+		// we don't know who initiated the channel (for very old channels). We're
+		// unlikely to hit this case since we're dealing with channels related
+		// to current forwards, so we just log that we don't know this value and
+		// allow initiator to be true.
+		switch rpcChan.OpenInitiator {
+		case lnrpc.Initiator_INITIATOR_LOCAL:
+			channel.initiator = true
+
+		case lnrpc.Initiator_INITIATOR_REMOTE:
+
+		default:
+			channel.initiator = true
+			log.Debugf("Channel initiator for %v with %v unknown",
+				rpcChan.ChanId, peer)
+		}
+
+		chans[rpcChan.ChanId] = channel
+	}
+
+	return chans, nil
+}
+
 func (l *lndclientGrpc) subscribeHtlcEvents(ctx context.Context) (
 	htlcEventsClient, error) {
 

lndclient_mock.go

@@ -27,16 +27,18 @@ type lndclientMock struct {
 	htlcInterceptorRequests  chan *interceptedEvent
 	htlcInterceptorResponses chan *interceptResponse
 
-	channels map[uint64]*channel
+	channels       map[uint64]*channel
+	closedChannels map[uint64]*channel
 }
 
-func newLndclientMock(channels map[uint64]*channel) *lndclientMock {
+func newLndclientMock(channels, closedChannels map[uint64]*channel) *lndclientMock {
 	return &lndclientMock{
 		htlcEvents:               make(chan *resolvedEvent),
 		htlcInterceptorRequests:  make(chan *interceptedEvent),
 		htlcInterceptorResponses: make(chan *interceptResponse),
 
-		channels: channels,
+		channels:       channels,
+		closedChannels: closedChannels,
 	}
 }
 
@@ -50,6 +52,10 @@ func (l *lndclientMock) listChannels() (map[uint64]*channel, error) {
 	return l.channels, nil
 }
 
+func (l *lndclientMock) listClosedChannels() (map[uint64]*channel, error) {
+	return l.closedChannels, nil
+}
+
 func (l *lndclientMock) subscribeHtlcEvents(ctx context.Context) (
 	htlcEventsClient, error) {
 

main.go

@@ -129,6 +129,11 @@ func main() {
 			Value: "127.0.0.1:9234",
 			Usage: "grpc server listen address",
 		},
+		cli.Uint64Flag{
+			Name:  "fwdhistorylimit",
+			Usage: "limit the number of htlc forwards that are persisted",
+			Value: defaultFwdHistoryLimit,
+		},
 		httpListenFlag,
 		stubFlag,
 	}

process.go

@@ -26,6 +26,8 @@ type lndclient interface {
 
 	listChannels() (map[uint64]*channel, error)
 
+	listClosedChannels() (map[uint64]*channel, error)
+
 	getNodeAlias(key route.Vertex) (string, error)
 
 	subscribeHtlcEvents(ctx context.Context) (htlcEventsClient, error)
@@ -524,6 +526,23 @@ func (p *process) getChanInfo(channel uint64) (*channel, error) {
 		return ch, nil
 	}
 
+	// If the channel is not open, fall back to checking our closed
+	// channels.
+	closedChannels, err := p.client.listClosedChannels()
+	if err != nil {
+		return nil, err
+	}
+
+	// Add to cache and try again.
+	for chanId, ch := range closedChannels {
+		p.chanMap[chanId] = ch
+	}
+
+	ch, ok = p.chanMap[channel]
+	if ok {
+		return ch, nil
+	}
+
 	// Channel not found.
 	return nil, fmt.Errorf("%w: %v", errChannelNotFound, channel)
 }

process_test.go

@@ -33,11 +33,11 @@ const (
 )
 
 func testProcess(t *testing.T, event resolveEvent) {
-	client := newLndclientMock(testChannels)
+	client := newLndclientMock(testChannels, nil)
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	db, cleanup := setupTestDb(t)
+	db, cleanup := setupTestDb(t, defaultFwdHistoryLimit)
 	defer cleanup()
 
 	log := zaptest.NewLogger(t).Sugar()
@@ -114,7 +114,7 @@ func TestLimits(t *testing.T) {
 func testRateLimit(t *testing.T, mode Mode) {
 	defer Timeout()()
 
-	db, cleanup := setupTestDb(t)
+	db, cleanup := setupTestDb(t, defaultFwdHistoryLimit)
 	defer cleanup()
 
 	cfg := &Limits{
@@ -130,7 +130,7 @@ func testRateLimit(t *testing.T, mode Mode) {
 		},
 	}
 
-	client := newLndclientMock(testChannels)
+	client := newLndclientMock(testChannels, nil)
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
@@ -204,7 +204,7 @@ func testRateLimit(t *testing.T, mode Mode) {
 func testMaxPending(t *testing.T, mode Mode) {
 	defer Timeout()()
 
-	db, cleanup := setupTestDb(t)
+	db, cleanup := setupTestDb(t, defaultFwdHistoryLimit)
 	defer cleanup()
 
 	cfg := &Limits{
@@ -222,7 +222,7 @@ func testMaxPending(t *testing.T, mode Mode) {
 		},
 	}
 
-	client := newLndclientMock(testChannels)
+	client := newLndclientMock(testChannels, nil)
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
@@ -278,12 +278,12 @@ func testMaxPending(t *testing.T, mode Mode) {
 
 func TestNewPeer(t *testing.T) {
 	// Initialize lnd with test channels.
-	client := newLndclientMock(testChannels)
+	client := newLndclientMock(testChannels, nil)
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	db, cleanup := setupTestDb(t)
+	db, cleanup := setupTestDb(t, defaultFwdHistoryLimit)
 	defer cleanup()
 
 	log := zaptest.NewLogger(t).Sugar()
@@ -323,7 +323,7 @@ func TestNewPeer(t *testing.T) {
 func TestBlocked(t *testing.T) {
 	defer Timeout()()
 
-	db, cleanup := setupTestDb(t)
+	db, cleanup := setupTestDb(t, defaultFwdHistoryLimit)
 	defer cleanup()
 
 	cfg := &Limits{
@@ -335,7 +335,7 @@ func TestBlocked(t *testing.T) {
 		},
 	}
 
-	client := newLndclientMock(testChannels)
+	client := newLndclientMock(testChannels, nil)
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
@@ -370,12 +370,12 @@ func TestBlocked(t *testing.T) {
 // TestChannelNotFound tests that we'll successfully exit when we cannot lookup the
 // channel that a htlc belongs to.
 func TestChannelNotFound(t *testing.T) {
-	client := newLndclientMock(testChannels)
+	client := newLndclientMock(testChannels, nil)
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	db, cleanup := setupTestDb(t)
+	db, cleanup := setupTestDb(t, defaultFwdHistoryLimit)
 	defer cleanup()
 
 	log := zaptest.NewLogger(t).Sugar()
@@ -407,3 +407,47 @@ func TestChannelNotFound(t *testing.T) {
 		t.Fatalf("timeout on process error")
 	}
 }
+
+// TestClosedChannelHtlc tests that we can handle intercepted htlcs that are associated
+// with closed channels.
+func TestClosedChannelHtlc(t *testing.T) {
+	// Initialize lnd with a closed channel.
+	var testClosedChannels = map[uint64]*channel{
+		5: {peer: route.Vertex{2}},
+	}
+	client := newLndclientMock(testChannels, testClosedChannels)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	db, cleanup := setupTestDb(t, defaultFwdHistoryLimit)
+	defer cleanup()
+
+	log := zaptest.NewLogger(t).Sugar()
+
+	cfg := &Limits{}
+
+	p := NewProcess(client, log, cfg, db)
+
+	exit := make(chan error)
+
+	go func() {
+		exit <- p.Run(ctx)
+	}()
+
+	// Send a htlc that is from a closed channel, it should be given the go-ahead to
+	// resume.
+	key := circuitKey{
+		channel: 5,
+		htlc:    3,
+	}
+	client.htlcInterceptorRequests <- &interceptedEvent{
+		circuitKey: key,
+	}
+
+	resp := <-client.htlcInterceptorResponses
+	require.Equal(t, key, resp.key)
+
+	cancel()
+	require.ErrorIs(t, <-exit, context.Canceled)
+}