htlcswitch+channeldb: move control tower to channeldb

channeldb/control_tower.go

@@ -0,0 +1,244 @@
+package channeldb
+
+import (
+	"errors"
+
+	"github.com/coreos/bbolt"
+	"github.com/lightningnetwork/lnd/lnwire"
+)
+
+var (
+	// ErrAlreadyPaid signals we have already paid this payment hash.
+	ErrAlreadyPaid = errors.New("invoice is already paid")
+
+	// ErrPaymentInFlight signals that payment for this payment hash is
+	// already "in flight" on the network.
+	ErrPaymentInFlight = errors.New("payment is in transition")
+
+	// ErrPaymentNotInitiated is returned  if payment wasn't initiated in
+	// switch.
+	ErrPaymentNotInitiated = errors.New("payment isn't initiated")
+
+	// ErrPaymentAlreadyCompleted is returned in the event we attempt to
+	// recomplete a completed payment.
+	ErrPaymentAlreadyCompleted = errors.New("payment is already completed")
+
+	// ErrUnknownPaymentStatus is returned when we do not recognize the
+	// existing state of a payment.
+	ErrUnknownPaymentStatus = errors.New("unknown payment status")
+)
+
+// ControlTower tracks all outgoing payments made by the switch, whose primary
+// purpose is to prevent duplicate payments to the same payment hash. In
+// production, a persistent implementation is preferred so that tracking can
+// survive across restarts. Payments are transition through various payment
+// states, and the ControlTower interface provides access to driving the state
+// transitions.
+type ControlTower interface {
+	// ClearForTakeoff atomically checks that no inflight or completed
+	// payments exist for this payment hash. If none are found, this method
+	// atomically transitions the status for this payment hash as InFlight.
+	ClearForTakeoff(htlc *lnwire.UpdateAddHTLC) error
+
+	// Success transitions an InFlight payment into a Completed payment.
+	// After invoking this method, ClearForTakeoff should always return an
+	// error to prevent us from making duplicate payments to the same
+	// payment hash.
+	Success(paymentHash [32]byte) error
+
+	// Fail transitions an InFlight payment into a Grounded Payment. After
+	// invoking this method, ClearForTakeoff should return nil on its next
+	// call for this payment hash, allowing the switch to make a subsequent
+	// payment.
+	Fail(paymentHash [32]byte) error
+}
+
+// paymentControl is persistent implementation of ControlTower to restrict
+// double payment sending.
+type paymentControl struct {
+	strict bool
+
+	db *DB
+}
+
+// NewPaymentControl creates a new instance of the paymentControl. The strict
+// flag indicates whether the controller should require "strict" state
+// transitions, which would be otherwise intolerant to older databases that may
+// already have duplicate payments to the same payment hash. It should be
+// enabled only after sufficient checks have been made to ensure the db does not
+// contain such payments. In the meantime, non-strict mode enforces a superset
+// of the state transitions that prevent additional payments to a given payment
+// hash from being added.
+func NewPaymentControl(strict bool, db *DB) ControlTower {
+	return &paymentControl{
+		strict: strict,
+		db:     db,
+	}
+}
+
+// ClearForTakeoff checks that we don't already have an InFlight or Completed
+// payment identified by the same payment hash.
+func (p *paymentControl) ClearForTakeoff(htlc *lnwire.UpdateAddHTLC) error {
+	var takeoffErr error
+	err := p.db.Batch(func(tx *bbolt.Tx) error {
+		// Retrieve current status of payment from local database.
+		paymentStatus, err := FetchPaymentStatusTx(
+			tx, htlc.PaymentHash,
+		)
+		if err != nil {
+			return err
+		}
+
+		// Reset the takeoff error, to avoid carrying over an error
+		// from a previous execution of the batched db transaction.
+		takeoffErr = nil
+
+		switch paymentStatus {
+
+		case StatusGrounded:
+			// It is safe to reattempt a payment if we know that we
+			// haven't left one in flight. Since this one is
+			// grounded, Transition the payment status to InFlight
+			// to prevent others.
+			return UpdatePaymentStatusTx(
+				tx, htlc.PaymentHash, StatusInFlight,
+			)
+
+		case StatusInFlight:
+			// We already have an InFlight payment on the network. We will
+			// disallow any more payment until a response is received.
+			takeoffErr = ErrPaymentInFlight
+
+		case StatusCompleted:
+			// We've already completed a payment to this payment hash,
+			// forbid the switch from sending another.
+			takeoffErr = ErrAlreadyPaid
+
+		default:
+			takeoffErr = ErrUnknownPaymentStatus
+		}
+
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
+	return takeoffErr
+}
+
+// Success transitions an InFlight payment to Completed, otherwise it returns an
+// error. After calling Success, ClearForTakeoff should prevent any further
+// attempts for the same payment hash.
+func (p *paymentControl) Success(paymentHash [32]byte) error {
+	var updateErr error
+	err := p.db.Batch(func(tx *bbolt.Tx) error {
+		paymentStatus, err := FetchPaymentStatusTx(
+			tx, paymentHash,
+		)
+		if err != nil {
+			return err
+		}
+
+		// Reset the update error, to avoid carrying over an error
+		// from a previous execution of the batched db transaction.
+		updateErr = nil
+
+		switch {
+
+		case paymentStatus == StatusGrounded && p.strict:
+			// Our records show the payment as still being grounded,
+			// meaning it never should have left the switch.
+			updateErr = ErrPaymentNotInitiated
+
+		case paymentStatus == StatusGrounded && !p.strict:
+			// Though our records show the payment as still being
+			// grounded, meaning it never should have left the
+			// switch, we permit this transition in non-strict mode
+			// to handle inconsistent db states.
+			fallthrough
+
+		case paymentStatus == StatusInFlight:
+			// A successful response was received for an InFlight
+			// payment, mark it as completed to prevent sending to
+			// this payment hash again.
+			return UpdatePaymentStatusTx(
+				tx, paymentHash, StatusCompleted,
+			)
+
+		case paymentStatus == StatusCompleted:
+			// The payment was completed previously, alert the
+			// caller that this may be a duplicate call.
+			updateErr = ErrPaymentAlreadyCompleted
+
+		default:
+			updateErr = ErrUnknownPaymentStatus
+		}
+
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
+	return updateErr
+}
+
+// Fail transitions an InFlight payment to Grounded, otherwise it returns an
+// error. After calling Fail, ClearForTakeoff should fail any further attempts
+// for the same payment hash.
+func (p *paymentControl) Fail(paymentHash [32]byte) error {
+	var updateErr error
+	err := p.db.Batch(func(tx *bbolt.Tx) error {
+		paymentStatus, err := FetchPaymentStatusTx(
+			tx, paymentHash,
+		)
+		if err != nil {
+			return err
+		}
+
+		// Reset the update error, to avoid carrying over an error
+		// from a previous execution of the batched db transaction.
+		updateErr = nil
+
+		switch {
+
+		case paymentStatus == StatusGrounded && p.strict:
+			// Our records show the payment as still being grounded,
+			// meaning it never should have left the switch.
+			updateErr = ErrPaymentNotInitiated
+
+		case paymentStatus == StatusGrounded && !p.strict:
+			// Though our records show the payment as still being
+			// grounded, meaning it never should have left the
+			// switch, we permit this transition in non-strict mode
+			// to handle inconsistent db states.
+			fallthrough
+
+		case paymentStatus == StatusInFlight:
+			// A failed response was received for an InFlight
+			// payment, mark it as Grounded again to allow
+			// subsequent attempts.
+			return UpdatePaymentStatusTx(
+				tx, paymentHash, StatusGrounded,
+			)
+
+		case paymentStatus == StatusCompleted:
+			// The payment was completed previously, and we are now
+			// reporting that it has failed. Leave the status as
+			// completed, but alert the user that something is
+			// wrong.
+			updateErr = ErrPaymentAlreadyCompleted
+
+		default:
+			updateErr = ErrUnknownPaymentStatus
+		}
+
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
+	return updateErr
+}

channeldb/control_tower_test.go

@@ -0,0 +1,375 @@
+package channeldb
+
+import (
+	"crypto/rand"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"testing"
+
+	"github.com/btcsuite/fastsha256"
+	"github.com/lightningnetwork/lnd/lnwire"
+)
+
+func initDB() (*DB, error) {
+	tempPath, err := ioutil.TempDir("", "switchdb")
+	if err != nil {
+		return nil, err
+	}
+
+	db, err := Open(tempPath)
+	if err != nil {
+		return nil, err
+	}
+
+	return db, err
+}
+
+func genPreimage() ([32]byte, error) {
+	var preimage [32]byte
+	if _, err := io.ReadFull(rand.Reader, preimage[:]); err != nil {
+		return preimage, err
+	}
+	return preimage, nil
+}
+
+func genHtlc() (*lnwire.UpdateAddHTLC, error) {
+	preimage, err := genPreimage()
+	if err != nil {
+		return nil, fmt.Errorf("unable to generate preimage: %v", err)
+	}
+
+	rhash := fastsha256.Sum256(preimage[:])
+	htlc := &lnwire.UpdateAddHTLC{
+		PaymentHash: rhash,
+		Amount:      1,
+	}
+
+	return htlc, nil
+}
+
+type paymentControlTestCase func(*testing.T, bool)
+
+var paymentControlTests = []struct {
+	name     string
+	strict   bool
+	testcase paymentControlTestCase
+}{
+	{
+		name:     "fail-strict",
+		strict:   true,
+		testcase: testPaymentControlSwitchFail,
+	},
+	{
+		name:     "double-send-strict",
+		strict:   true,
+		testcase: testPaymentControlSwitchDoubleSend,
+	},
+	{
+		name:     "double-pay-strict",
+		strict:   true,
+		testcase: testPaymentControlSwitchDoublePay,
+	},
+	{
+		name:     "fail-not-strict",
+		strict:   false,
+		testcase: testPaymentControlSwitchFail,
+	},
+	{
+		name:     "double-send-not-strict",
+		strict:   false,
+		testcase: testPaymentControlSwitchDoubleSend,
+	},
+	{
+		name:     "double-pay-not-strict",
+		strict:   false,
+		testcase: testPaymentControlSwitchDoublePay,
+	},
+}
+
+// TestPaymentControls runs a set of common tests against both the strict and
+// non-strict payment control instances. This ensures that the two both behave
+// identically when making the expected state-transitions of the stricter
+// implementation. Behavioral differences in the strict and non-strict
+// implementations are tested separately.
+func TestPaymentControls(t *testing.T) {
+	for _, test := range paymentControlTests {
+		t.Run(test.name, func(t *testing.T) {
+			test.testcase(t, test.strict)
+		})
+	}
+}
+
+// testPaymentControlSwitchFail checks that payment status returns to Grounded
+// status after failing, and that ClearForTakeoff allows another HTLC for the
+// same payment hash.
+func testPaymentControlSwitchFail(t *testing.T, strict bool) {
+	t.Parallel()
+
+	db, err := initDB()
+	if err != nil {
+		t.Fatalf("unable to init db: %v", err)
+	}
+
+	pControl := NewPaymentControl(strict, db)
+
+	htlc, err := genHtlc()
+	if err != nil {
+		t.Fatalf("unable to generate htlc message: %v", err)
+	}
+
+	// Sends base htlc message which initiate StatusInFlight.
+	if err := pControl.ClearForTakeoff(htlc); err != nil {
+		t.Fatalf("unable to send htlc message: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusInFlight)
+
+	// Fail the payment, which should moved it to Grounded.
+	if err := pControl.Fail(htlc.PaymentHash); err != nil {
+		t.Fatalf("unable to fail payment hash: %v", err)
+	}
+
+	// Verify the status is indeed Grounded.
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusGrounded)
+
+	// Sends the htlc again, which should succeed since the prior payment
+	// failed.
+	if err := pControl.ClearForTakeoff(htlc); err != nil {
+		t.Fatalf("unable to send htlc message: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusInFlight)
+
+	// Verifies that status was changed to StatusCompleted.
+	if err := pControl.Success(htlc.PaymentHash); err != nil {
+		t.Fatalf("error shouldn't have been received, got: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusCompleted)
+
+	// Attempt a final payment, which should now fail since the prior
+	// payment succeed.
+	if err := pControl.ClearForTakeoff(htlc); err != ErrAlreadyPaid {
+		t.Fatalf("unable to send htlc message: %v", err)
+	}
+}
+
+// testPaymentControlSwitchDoubleSend checks the ability of payment control to
+// prevent double sending of htlc message, when message is in StatusInFlight.
+func testPaymentControlSwitchDoubleSend(t *testing.T, strict bool) {
+	t.Parallel()
+
+	db, err := initDB()
+	if err != nil {
+		t.Fatalf("unable to init db: %v", err)
+	}
+
+	pControl := NewPaymentControl(strict, db)
+
+	htlc, err := genHtlc()
+	if err != nil {
+		t.Fatalf("unable to generate htlc message: %v", err)
+	}
+
+	// Sends base htlc message which initiate base status and move it to
+	// StatusInFlight and verifies that it was changed.
+	if err := pControl.ClearForTakeoff(htlc); err != nil {
+		t.Fatalf("unable to send htlc message: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusInFlight)
+
+	// Try to initiate double sending of htlc message with the same
+	// payment hash, should result in error indicating that payment has
+	// already been sent.
+	if err := pControl.ClearForTakeoff(htlc); err != ErrPaymentInFlight {
+		t.Fatalf("payment control wrong behaviour: " +
+			"double sending must trigger ErrPaymentInFlight error")
+	}
+}
+
+// TestPaymentControlSwitchDoublePay checks the ability of payment control to
+// prevent double payment.
+func testPaymentControlSwitchDoublePay(t *testing.T, strict bool) {
+	t.Parallel()
+
+	db, err := initDB()
+	if err != nil {
+		t.Fatalf("unable to init db: %v", err)
+	}
+
+	pControl := NewPaymentControl(strict, db)
+
+	htlc, err := genHtlc()
+	if err != nil {
+		t.Fatalf("unable to generate htlc message: %v", err)
+	}
+
+	// Sends base htlc message which initiate StatusInFlight.
+	if err := pControl.ClearForTakeoff(htlc); err != nil {
+		t.Fatalf("unable to send htlc message: %v", err)
+	}
+
+	// Verify that payment is InFlight.
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusInFlight)
+
+	// Move payment to completed status, second payment should return error.
+	if err := pControl.Success(htlc.PaymentHash); err != nil {
+		t.Fatalf("error shouldn't have been received, got: %v", err)
+	}
+
+	// Verify that payment is Completed.
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusCompleted)
+
+	if err := pControl.ClearForTakeoff(htlc); err != ErrAlreadyPaid {
+		t.Fatalf("payment control wrong behaviour:" +
+			" double payment must trigger ErrAlreadyPaid")
+	}
+}
+
+// TestPaymentControlNonStrictSuccessesWithoutInFlight checks that a non-strict
+// payment control will allow calls to Success when no payment is in flight. This
+// is necessary to gracefully handle the case in which the switch already sent
+// out a payment for a particular payment hash in a prior db version that didn't
+// have payment statuses.
+func TestPaymentControlNonStrictSuccessesWithoutInFlight(t *testing.T) {
+	t.Parallel()
+
+	db, err := initDB()
+	if err != nil {
+		t.Fatalf("unable to init db: %v", err)
+	}
+
+	pControl := NewPaymentControl(false, db)
+
+	htlc, err := genHtlc()
+	if err != nil {
+		t.Fatalf("unable to generate htlc message: %v", err)
+	}
+
+	if err := pControl.Success(htlc.PaymentHash); err != nil {
+		t.Fatalf("unable to mark payment hash success: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusCompleted)
+
+	err = pControl.Success(htlc.PaymentHash)
+	if err != ErrPaymentAlreadyCompleted {
+		t.Fatalf("unable to remark payment hash failed: %v", err)
+	}
+}
+
+// TestPaymentControlNonStrictFailsWithoutInFlight checks that a non-strict
+// payment control will allow calls to Fail when no payment is in flight. This
+// is necessary to gracefully handle the case in which the switch already sent
+// out a payment for a particular payment hash in a prior db version that didn't
+// have payment statuses.
+func TestPaymentControlNonStrictFailsWithoutInFlight(t *testing.T) {
+	t.Parallel()
+
+	db, err := initDB()
+	if err != nil {
+		t.Fatalf("unable to init db: %v", err)
+	}
+
+	pControl := NewPaymentControl(false, db)
+
+	htlc, err := genHtlc()
+	if err != nil {
+		t.Fatalf("unable to generate htlc message: %v", err)
+	}
+
+	if err := pControl.Fail(htlc.PaymentHash); err != nil {
+		t.Fatalf("unable to mark payment hash failed: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusGrounded)
+
+	err = pControl.Fail(htlc.PaymentHash)
+	if err != nil {
+		t.Fatalf("unable to remark payment hash failed: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusGrounded)
+
+	err = pControl.Success(htlc.PaymentHash)
+	if err != nil {
+		t.Fatalf("unable to remark payment hash success: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusCompleted)
+
+	err = pControl.Fail(htlc.PaymentHash)
+	if err != ErrPaymentAlreadyCompleted {
+		t.Fatalf("unable to remark payment hash failed: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusCompleted)
+}
+
+// TestPaymentControlStrictSuccessesWithoutInFlight checks that a strict payment
+// control will disallow calls to Success when no payment is in flight.
+func TestPaymentControlStrictSuccessesWithoutInFlight(t *testing.T) {
+	t.Parallel()
+
+	db, err := initDB()
+	if err != nil {
+		t.Fatalf("unable to init db: %v", err)
+	}
+
+	pControl := NewPaymentControl(true, db)
+
+	htlc, err := genHtlc()
+	if err != nil {
+		t.Fatalf("unable to generate htlc message: %v", err)
+	}
+
+	err = pControl.Success(htlc.PaymentHash)
+	if err != ErrPaymentNotInitiated {
+		t.Fatalf("expected ErrPaymentNotInitiated, got %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusGrounded)
+}
+
+// TestPaymentControlStrictFailsWithoutInFlight checks that a strict payment
+// control will disallow calls to Fail when no payment is in flight.
+func TestPaymentControlStrictFailsWithoutInFlight(t *testing.T) {
+	t.Parallel()
+
+	db, err := initDB()
+	if err != nil {
+		t.Fatalf("unable to init db: %v", err)
+	}
+
+	pControl := NewPaymentControl(true, db)
+
+	htlc, err := genHtlc()
+	if err != nil {
+		t.Fatalf("unable to generate htlc message: %v", err)
+	}
+
+	err = pControl.Fail(htlc.PaymentHash)
+	if err != ErrPaymentNotInitiated {
+		t.Fatalf("expected ErrPaymentNotInitiated, got %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, StatusGrounded)
+}
+
+func assertPaymentStatus(t *testing.T, db *DB,
+	hash [32]byte, expStatus PaymentStatus) {
+
+	t.Helper()
+
+	pStatus, err := db.FetchPaymentStatus(hash)
+	if err != nil {
+		t.Fatalf("unable to fetch payment status: %v", err)
+	}
+
+	if pStatus != expStatus {
+		t.Fatalf("payment status mismatch: expected %v, got %v",
+			expStatus, pStatus)
+	}
+}