From 3d50edf9f821022230adfebf7d501c8233985bef Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Thu, 22 Jul 2021 13:16:55 +0800
Subject: [PATCH] channeldb: add method to wipe all forwarding packages

This commit adds a method, Wipe, to remove all forwarding packages on
disk for a given channel.
---
 channeldb/forwarding_package.go      | 21 +++++++++++++
 channeldb/forwarding_package_test.go | 44 ++++++++++++++++++++++++++++
 htlcswitch/link_test.go              |  4 +++
 3 files changed, 69 insertions(+)

diff --git a/channeldb/forwarding_package.go b/channeldb/forwarding_package.go
index 53e8881d3..09f5671cc 100644
--- a/channeldb/forwarding_package.go
+++ b/channeldb/forwarding_package.go
@@ -433,6 +433,9 @@ type FwdPackager interface {
 	// RemovePkg deletes a forwarding package owned by this channel at
 	// the provided remote `height`.
 	RemovePkg(tx kvdb.RwTx, height uint64) error
+
+	// Wipe deletes all the forwarding packages owned by this channel.
+	Wipe(tx kvdb.RwTx) error
 }
 
 // ChannelPackager is used by a channel to manage the lifecycle of its forwarding
@@ -944,6 +947,24 @@ func (p *ChannelPackager) RemovePkg(tx kvdb.RwTx, height uint64) error {
 	return sourceBkt.DeleteNestedBucket(heightKey[:])
 }
 
+// Wipe deletes all the channel's forwarding packages, if any.
+func (p *ChannelPackager) Wipe(tx kvdb.RwTx) error {
+	// If the root bucket doesn't exist, there's no need to delete.
+	fwdPkgBkt := tx.ReadWriteBucket(fwdPackagesKey)
+	if fwdPkgBkt == nil {
+		return nil
+	}
+
+	sourceBytes := makeLogKey(p.source.ToUint64())
+
+	// If the nested bucket doesn't exist, there's no need to delete.
+	if fwdPkgBkt.NestedReadWriteBucket(sourceBytes[:]) == nil {
+		return nil
+	}
+
+	return fwdPkgBkt.DeleteNestedBucket(sourceBytes[:])
+}
+
 // uint16Key writes the provided 16-bit unsigned integer to a 2-byte slice.
 func uint16Key(i uint16) []byte {
 	key := make([]byte, 2)
diff --git a/channeldb/forwarding_package_test.go b/channeldb/forwarding_package_test.go
index a56d1e530..a1a22fdf6 100644
--- a/channeldb/forwarding_package_test.go
+++ b/channeldb/forwarding_package_test.go
@@ -11,6 +11,7 @@ import (
 	"github.com/lightningnetwork/lnd/channeldb"
 	"github.com/lightningnetwork/lnd/kvdb"
 	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/stretchr/testify/require"
 )
 
 // TestPkgFilterBruteForce tests the behavior of a pkg filter up to size 1000,
@@ -729,6 +730,49 @@ func TestPackagerSettleFailsThenAdds(t *testing.T) {
 	}
 }
 
+// TestPackagerWipeAll checks that when the method is called, all the related
+// forwarding packages will be removed.
+func TestPackagerWipeAll(t *testing.T) {
+	t.Parallel()
+
+	db := makeFwdPkgDB(t, "")
+
+	shortChanID := lnwire.NewShortChanIDFromInt(1)
+	packager := channeldb.NewChannelPackager(shortChanID)
+
+	// To begin, there should be no forwarding packages on disk.
+	fwdPkgs := loadFwdPkgs(t, db, packager)
+	require.Empty(t, fwdPkgs, "no forwarding packages should exist")
+
+	// Now, check we can wipe without error since it's a noop.
+	err := kvdb.Update(db, packager.Wipe, func() {})
+	require.NoError(t, err, "unable to wipe fwdpkg")
+
+	// Next, create and write two forwarding packages with no htlcs.
+	fwdPkg1 := channeldb.NewFwdPkg(shortChanID, 0, nil, nil)
+	fwdPkg2 := channeldb.NewFwdPkg(shortChanID, 1, nil, nil)
+
+	err = kvdb.Update(db, func(tx kvdb.RwTx) error {
+		if err := packager.AddFwdPkg(tx, fwdPkg2); err != nil {
+			return err
+		}
+		return packager.AddFwdPkg(tx, fwdPkg1)
+	}, func() {})
+	require.NoError(t, err, "unable to add fwd pkg")
+
+	// There should now be two fwdpkgs on disk.
+	fwdPkgs = loadFwdPkgs(t, db, packager)
+	require.Equal(t, 2, len(fwdPkgs), "expected 2 fwdpkg")
+
+	// Now, wipe all forwarding packages from disk.
+	err = kvdb.Update(db, packager.Wipe, func() {})
+	require.NoError(t, err, "unable to wipe fwdpkg")
+
+	// Check that the packages were actually removed.
+	fwdPkgs = loadFwdPkgs(t, db, packager)
+	require.Empty(t, fwdPkgs, "no forwarding packages should exist")
+}
+
 // assertFwdPkgState checks the current state of a fwdpkg meets our
 // expectations.
 func assertFwdPkgState(t *testing.T, fwdPkg *channeldb.FwdPkg,
diff --git a/htlcswitch/link_test.go b/htlcswitch/link_test.go
index d228b9509..41f51ad0f 100644
--- a/htlcswitch/link_test.go
+++ b/htlcswitch/link_test.go
@@ -5373,6 +5373,10 @@ func (*mockPackager) RemovePkg(tx kvdb.RwTx, height uint64) error {
 	return nil
 }
 
+func (*mockPackager) Wipe(tx kvdb.RwTx) error {
+	return nil
+}
+
 func (*mockPackager) AckSettleFails(tx kvdb.RwTx,
 	settleFailRefs ...channeldb.SettleFailRef) error {
 	return nil