staticaddr: address and deposit adjustments for withdrawals

Author: hieblmi

staticaddr/address/manager.go

@@ -11,7 +11,6 @@ import (
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/lightninglabs/lndclient"
-	"github.com/lightninglabs/loop/staticaddr"
 	"github.com/lightninglabs/loop/staticaddr/script"
 	"github.com/lightninglabs/loop/staticaddr/version"
 	"github.com/lightninglabs/loop/swap"

staticaddr/address/manager_test.go

@@ -32,6 +32,17 @@ type mockStaticAddressClient struct {
 	mock.Mock
 }
 
+func (m *mockStaticAddressClient) ServerWithdrawDeposits(ctx context.Context,
+	in *swapserverrpc.ServerWithdrawRequest,
+	opts ...grpc.CallOption) (*swapserverrpc.ServerWithdrawResponse,
+	error) {
+
+	args := m.Called(ctx, in, opts)
+
+	return args.Get(0).(*swapserverrpc.ServerWithdrawResponse),
+		args.Error(1)
+}
+
 func (m *mockStaticAddressClient) ServerNewAddress(ctx context.Context,
 	in *swapserverrpc.ServerNewAddressRequest, opts ...grpc.CallOption) (
 	*swapserverrpc.ServerNewAddressResponse, error) {

staticaddr/deposit/actions.go

@@ -147,3 +147,17 @@ func (f *FSM) SweptExpiredDepositAction(_ fsm.EventContext) fsm.EventType {
 
 	return fsm.NoOp
 }
+
+// WithdrawnDepositAction is the final action after a withdrawal. It signals to
+// the manager that the deposit has been swept and the FSM can be removed.
+func (f *FSM) WithdrawnDepositAction(_ fsm.EventContext) fsm.EventType {
+	select {
+	case <-f.ctx.Done():
+		return fsm.OnError
+
+	default:
+		f.finalizedDepositChan <- f.deposit.OutPoint
+	}
+
+	return fsm.NoOp
+}

staticaddr/deposit/deposit.go

@@ -35,7 +35,7 @@ type Deposit struct {
 	// state is the current state of the deposit.
 	state fsm.StateType
 
-	// The outpoint of the deposit.
+	// Outpoint of the deposit.
 	wire.OutPoint
 
 	// Value is the amount of the deposit.
@@ -52,6 +52,10 @@ type Deposit struct {
 	// ExpirySweepTxid is the transaction id of the expiry sweep.
 	ExpirySweepTxid chainhash.Hash
 
+	// WithdrawalSweepAddress is the address that is used to
+	// cooperatively sweep the deposit to before it is expired.
+	WithdrawalSweepAddress string
+
 	sync.Mutex
 }
 
@@ -68,7 +72,7 @@ func (d *Deposit) IsInFinalState() bool {
 	d.Lock()
 	defer d.Unlock()
 
-	return d.state == Expired || d.state == Failed
+	return d.state == Expired || d.state == Withdrawn || d.state == Failed
 }
 
 func (d *Deposit) IsExpired(currentHeight, expiry uint32) bool {

staticaddr/deposit/fsm.go

@@ -29,6 +29,10 @@ var (
 var (
 	Deposited = fsm.StateType("Deposited")
 
+	Withdrawing = fsm.StateType("Withdrawing")
+
+	Withdrawn = fsm.StateType("Withdrawn")
+
 	PublishExpiredDeposit = fsm.StateType("PublishExpiredDeposit")
 
 	WaitForExpirySweep = fsm.StateType("WaitForExpirySweep")
@@ -40,11 +44,13 @@ var (
 
 // Events.
 var (
-	OnStart           = fsm.EventType("OnStart")
-	OnExpiry          = fsm.EventType("OnExpiry")
-	OnExpiryPublished = fsm.EventType("OnExpiryPublished")
-	OnExpirySwept     = fsm.EventType("OnExpirySwept")
-	OnRecover         = fsm.EventType("OnRecover")
+	OnStart             = fsm.EventType("OnStart")
+	OnWithdrawInitiated = fsm.EventType("OnWithdrawInitiated")
+	OnWithdrawn         = fsm.EventType("OnWithdrawn")
+	OnExpiry            = fsm.EventType("OnExpiry")
+	OnExpiryPublished   = fsm.EventType("OnExpiryPublished")
+	OnExpirySwept       = fsm.EventType("OnExpirySwept")
+	OnRecover           = fsm.EventType("OnRecover")
 )
 
 // FSM is the state machine that handles the instant out.
@@ -118,13 +124,7 @@ func NewFSM(ctx context.Context, deposit *Deposit, cfg *ManagerConfig,
 		for {
 			select {
 			case currentHeight := <-depoFsm.blockNtfnChan:
-				err := depoFsm.handleBlockNotification(
-					currentHeight,
-				)
-				if err != nil {
-					log.Errorf("error handling block "+
-						"notification: %v", err)
-				}
+				depoFsm.handleBlockNotification(currentHeight)
 
 			case <-ctx.Done():
 				return
@@ -138,15 +138,10 @@ func NewFSM(ctx context.Context, deposit *Deposit, cfg *ManagerConfig,
 // handleBlockNotification inspects the current block height and sends the
 // OnExpiry event to publish the expiry sweep transaction if the deposit timed
 // out, or it republishes the expiry sweep transaction if it was not yet swept.
-func (f *FSM) handleBlockNotification(currentHeight uint32) error {
-	params, err := f.cfg.AddressManager.GetStaticAddressParameters(f.ctx)
-	if err != nil {
-		return err
-	}
-
+func (f *FSM) handleBlockNotification(currentHeight uint32) {
 	// If the deposit is expired but not yet sufficiently confirmed, we
 	// republish the expiry sweep transaction.
-	if f.deposit.IsExpired(currentHeight, params.Expiry) {
+	if f.deposit.IsExpired(currentHeight, f.params.Expiry) {
 		if f.deposit.IsInState(WaitForExpirySweep) {
 			f.PublishDepositExpirySweepAction(nil)
 		} else {
@@ -159,8 +154,6 @@ func (f *FSM) handleBlockNotification(currentHeight uint32) error {
 			}()
 		}
 	}
-
-	return nil
 }
 
 // DepositStatesV0 returns the states a deposit can be in.
@@ -174,8 +167,9 @@ func (f *FSM) DepositStatesV0() fsm.States {
 		},
 		Deposited: fsm.State{
 			Transitions: fsm.Transitions{
-				OnExpiry:  PublishExpiredDeposit,
-				OnRecover: Deposited,
+				OnExpiry:            PublishExpiredDeposit,
+				OnWithdrawInitiated: Withdrawing,
+				OnRecover:           Deposited,
 			},
 			Action: fsm.NoOpAction,
 		},
@@ -210,6 +204,36 @@ func (f *FSM) DepositStatesV0() fsm.States {
 			},
 			Action: f.SweptExpiredDepositAction,
 		},
+		Withdrawing: fsm.State{
+			Transitions: fsm.Transitions{
+				OnWithdrawn: Withdrawn,
+				// Upon recovery, we go back to the Deposited
+				// state. The deposit by then has a withdrawal
+				// address stamped to it which will cause it to
+				// transition into the Withdrawing state again.
+				OnRecover: Deposited,
+
+				// A precondition for the Withdrawing state is
+				// that the withdrawal transaction has been
+				// broadcast. If the deposit expires while the
+				// withdrawal isn't confirmed, we can ignore the
+				// expiry.
+				OnExpiry: Withdrawing,
+
+				// If the withdrawal failed we go back to
+				// Deposited, hoping that another withdrawal
+				// attempt will be successful. Alternatively,
+				// the client can wait for the timeout sweep.
+				fsm.OnError: Deposited,
+			},
+			Action: fsm.NoOpAction,
+		},
+		Withdrawn: fsm.State{
+			Transitions: fsm.Transitions{
+				OnExpiry: Expired,
+			},
+			Action: f.WithdrawnDepositAction,
+		},
 		Failed: fsm.State{
 			Transitions: fsm.Transitions{
 				OnExpiry: Failed,

staticaddr/deposit/manager.go

@@ -3,6 +3,7 @@ package deposit
 import (
 	"context"
 	"fmt"
+	"sort"
 	"sync"
 	"time"
 
@@ -11,6 +12,7 @@ import (
 	"github.com/btcsuite/btcd/wire"
 	"github.com/lightninglabs/lndclient"
 	"github.com/lightninglabs/loop"
+	"github.com/lightninglabs/loop/fsm"
 	staticaddressrpc "github.com/lightninglabs/loop/swapserverrpc"
 	"github.com/lightningnetwork/lnd/lnrpc/walletrpc"
 	"github.com/lightningnetwork/lnd/lnwallet"
@@ -29,6 +31,10 @@ const (
 	// MaxConfs is unset since we don't require a max number of
 	// confirmations for deposits.
 	MaxConfs = 0
+
+	// DefaultTransitionTimeout is the default timeout for transitions in
+	// the deposit state machine.
+	DefaultTransitionTimeout = 1 * time.Minute
 )
 
 // ManagerConfig holds the configuration for the address manager.
@@ -128,7 +134,7 @@ func (m *Manager) Run(ctx context.Context, currentHeight uint32) error {
 	}
 
 	// Start the deposit notifier.
-	m.pollDeposits(ctx)
+	m.pollDeposits(m.runCtx)
 
 	// Communicate to the caller that the address manager has completed its
 	// initialization.
@@ -420,3 +426,89 @@ func (m *Manager) finalizeDeposit(outpoint wire.OutPoint) {
 	delete(m.deposits, outpoint)
 	m.Unlock()
 }
+
+// GetActiveDepositsInState returns all active deposits.
+func (m *Manager) GetActiveDepositsInState(stateFilter fsm.StateType) (
+	[]*Deposit, error) {
+
+	m.Lock()
+	defer m.Unlock()
+
+	var deposits []*Deposit
+	for _, fsm := range m.activeDeposits {
+		if fsm.deposit.GetState() != stateFilter {
+			continue
+		}
+		deposits = append(deposits, fsm.deposit)
+	}
+
+	sort.Slice(deposits, func(i, j int) bool {
+		return deposits[i].ConfirmationHeight <
+			deposits[j].ConfirmationHeight
+	})
+
+	return deposits, nil
+}
+
+// GetAllDeposits returns all active deposits.
+func (m *Manager) GetAllDeposits() ([]*Deposit, error) {
+	return m.cfg.Store.AllDeposits(m.runCtx)
+}
+
+// AllOutpointsActiveDeposits checks if all deposits referenced by the outpoints
+// are active and in the specified state.
+func (m *Manager) AllOutpointsActiveDeposits(outpoints []wire.OutPoint,
+	stateFilter fsm.StateType) ([]*Deposit, bool) {
+
+	m.Lock()
+	defer m.Unlock()
+
+	deposits := make([]*Deposit, 0, len(outpoints))
+	for _, o := range outpoints {
+		if _, ok := m.activeDeposits[o]; !ok {
+			return nil, false
+		}
+
+		deposit := m.deposits[o]
+		if deposit.GetState() != stateFilter {
+			return nil, false
+		}
+
+		deposits = append(deposits, m.deposits[o])
+	}
+
+	return deposits, true
+}
+
+// TransitionDeposits allows a caller to transition a set of deposits to a new
+// state.
+func (m *Manager) TransitionDeposits(deposits []*Deposit, event fsm.EventType,
+	expectedFinalState fsm.StateType) error {
+
+	for _, d := range deposits {
+		m.Lock()
+		sm, ok := m.activeDeposits[d.OutPoint]
+		m.Unlock()
+		if !ok {
+			return fmt.Errorf("deposit not found")
+		}
+
+		err := sm.SendEvent(event, nil)
+		if err != nil {
+			return err
+		}
+		err = sm.DefaultObserver.WaitForState(
+			m.runCtx, DefaultTransitionTimeout, expectedFinalState,
+		)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// UpdateDeposit overrides all fields of the deposit with given ID in the store.
+func (m *Manager) UpdateDeposit(d *Deposit) error {
+	return m.cfg.Store.UpdateDeposit(m.runCtx, d)
+}

staticaddr/deposit/manager_test.go

@@ -49,6 +49,17 @@ type mockStaticAddressClient struct {
 	mock.Mock
 }
 
+func (m *mockStaticAddressClient) ServerWithdrawDeposits(ctx context.Context,
+	in *swapserverrpc.ServerWithdrawRequest,
+	opts ...grpc.CallOption) (*swapserverrpc.ServerWithdrawResponse,
+	error) {
+
+	args := m.Called(ctx, in, opts)
+
+	return args.Get(0).(*swapserverrpc.ServerWithdrawResponse),
+		args.Error(1)
+}
+
 func (m *mockStaticAddressClient) ServerNewAddress(ctx context.Context,
 	in *swapserverrpc.ServerNewAddressRequest, opts ...grpc.CallOption) (
 	*swapserverrpc.ServerNewAddressResponse, error) {
@@ -149,16 +160,15 @@ func (m *MockChainNotifier) RegisterSpendNtfn(ctx context.Context,
 // TestManager checks that the manager processes the right channel notifications
 // while a deposit is expiring.
 func TestManager(t *testing.T) {
-	ctxb, cancel := context.WithCancel(context.Background())
-	defer cancel()
+	ctx := context.Background()
 
 	// Create the test context with required mocks.
 	testContext := newManagerTestContext(t)
 
 	// Start the deposit manager.
 	go func() {
 		err := testContext.manager.Run(
-			ctxb, uint32(testContext.mockLnd.Height),
+			ctx, uint32(testContext.mockLnd.Height),
 		)
 		require.NoError(t, err)
 	}()
@@ -191,6 +201,9 @@ func TestManager(t *testing.T) {
 		BlockHeight: defaultDepositConfirmations + defaultExpiry + 3,
 		Tx:          expiryTx,
 	}
+
+	// Ensure that the deposit is finalized.
+	<-finalizedDepositChan
 }
 
 // ManagerTestContext is a helper struct that contains all the necessary