channeldb/control_tower: add payment information during state changes

This commit gives a new responsibility to the control tower, letting it
populate the payment bucket structure as the payment goes through
its different stages.

The payment will transition states Grounded->InFlight->Success/Failed,
where the CreationInfo/AttemptInfo/Preimage must be set accordingly.

This will be the main driver for the router state machine.

channeldb/control_tower.go

@@ -1,11 +1,12 @@
 package channeldb
 
 import (
+	"bytes"
+	"encoding/binary"
 	"errors"
 
 	"github.com/coreos/bbolt"
 	"github.com/lightningnetwork/lnd/lntypes"
-	"github.com/lightningnetwork/lnd/lnwire"
 )
 
 var (
@@ -33,29 +34,31 @@ var (
 	ErrUnknownPaymentStatus = errors.New("unknown payment status")
 )
 
-// ControlTower tracks all outgoing payments made by the switch, whose primary
+// ControlTower tracks all outgoing payments made, whose primary purpose is to
-// purpose is to prevent duplicate payments to the same payment hash. In
+// prevent duplicate payments to the same payment hash. In production, a
-// production, a persistent implementation is preferred so that tracking can
+// persistent implementation is preferred so that tracking can survive across
-// survive across restarts. Payments are transition through various payment
+// restarts. Payments are transitioned through various payment states, and the
-// states, and the ControlTower interface provides access to driving the state
+// ControlTower interface provides access to driving the state transitions.
-// transitions.
 type ControlTower interface {
-	// ClearForTakeoff atomically checks that no inflight or completed
+	// InitPayment atomically moves the payment into the InFlight state.
-	// payments exist for this payment hash. If none are found, this method
+	// This method checks that no completed payment exist for this payment
-	// atomically transitions the status for this payment hash as InFlight.
+	// hash.
-	ClearForTakeoff(htlc *lnwire.UpdateAddHTLC) error
+	InitPayment(lntypes.Hash, *PaymentCreationInfo) error
 
-	// Success transitions an InFlight payment into a Completed payment.
+	// RegisterAttempt atomically records the provided PaymentAttemptInfo.
-	// After invoking this method, ClearForTakeoff should always return an
+	RegisterAttempt(lntypes.Hash, *PaymentAttemptInfo) error
-	// 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
+	// Success transitions a payment into the Completed state. After
-	// invoking this method, ClearForTakeoff should return nil on its next
+	// invoking this method, InitPayment should always return an error to
-	// call for this payment hash, allowing the switch to make a subsequent
+	// prevent us from making duplicate payments to the same payment hash.
-	// payment.
+	// The provided preimage is atomically saved to the DB for record
-	Fail(paymentHash [32]byte) error
+	// keeping.
+	Success(lntypes.Hash, lntypes.Preimage) error
+
+	// Fail transitions a payment into the Failed state. After invoking
+	// this method, InitPayment should return nil on its next call for this
+	// payment hash, allowing the switch to make a subsequent payment.
+	Fail(lntypes.Hash) error
 }
 
 // paymentControl is persistent implementation of ControlTower to restrict
@@ -71,12 +74,22 @@ func NewPaymentControl(db *DB) ControlTower {
 	}
 }
 
-// ClearForTakeoff checks that we don't already have an InFlight or Completed
+// InitPayment checks or records the given PaymentCreationInfo with the DB,
-// payment identified by the same payment hash.
+// making sure it does not already exist as an in-flight payment. Then this
-func (p *paymentControl) ClearForTakeoff(htlc *lnwire.UpdateAddHTLC) error {
+// method returns successfully, the payment is guranteeed to be in the InFlight
+// state.
+func (p *paymentControl) InitPayment(paymentHash lntypes.Hash,
+	info *PaymentCreationInfo) error {
+
+	var b bytes.Buffer
+	if err := serializePaymentCreationInfo(&b, info); err != nil {
+		return err
+	}
+	infoBytes := b.Bytes()
+
 	var takeoffErr error
 	err := p.db.Batch(func(tx *bbolt.Tx) error {
-		bucket, err := fetchPaymentBucket(tx, htlc.PaymentHash)
+		bucket, err := fetchPaymentBucket(tx, paymentHash)
 		if err != nil {
 			return err
 		}
@@ -92,41 +105,118 @@ func (p *paymentControl) ClearForTakeoff(htlc *lnwire.UpdateAddHTLC) error {
 
 		// We allow retrying failed payments.
 		case StatusFailed:
-			fallthrough
 
-		// It is safe to reattempt a payment if we know that we haven't
+		// This is a new payment that is being initialized for the
-		// left one in flight. Since this one is grounded or failed,
+		// first time.
-		// transition the payment status to InFlight to prevent others.
 		case StatusGrounded:
-			return bucket.Put(paymentStatusKey, StatusInFlight.Bytes())
 
 		// We already have an InFlight payment on the network. We will
-		// disallow any more payment until a response is received.
+		// disallow any new payments.
 		case StatusInFlight:
 			takeoffErr = ErrPaymentInFlight
+			return nil
 
 		// We've already completed a payment to this payment hash,
 		// forbid the switch from sending another.
 		case StatusCompleted:
 			takeoffErr = ErrAlreadyPaid
+			return nil
 
 		default:
 			takeoffErr = ErrUnknownPaymentStatus
+			return nil
+		}
+
+		// Obtain a new sequence number for this payment. This is used
+		// to sort the payments in order of creation, and also acts as
+		// a unique identifier for each payment.
+		sequenceNum, err := nextPaymentSequence(tx)
+		if err != nil {
+			return err
+		}
+
+		err = bucket.Put(paymentSequenceKey, sequenceNum)
+		if err != nil {
+			return err
+		}
+
+		// We'll move it into the InFlight state.
+		err = bucket.Put(paymentStatusKey, StatusInFlight.Bytes())
+		if err != nil {
+			return err
+		}
+
+		// Add the payment info to the bucket, which contains the
+		// static information for this payment
+		err = bucket.Put(paymentCreationInfoKey, infoBytes)
+		if err != nil {
+			return err
+		}
+
+		// We'll delete any lingering attempt info to start with, in
+		// case we are initializing a payment that was attempted
+		// earlier, but left in a state where we could retry.
+		err = bucket.Delete(paymentAttemptInfoKey)
+		if err != nil {
+			return err
 		}
 
 		return nil
 	})
 	if err != nil {
-		return err
+		return nil
 	}
 
 	return takeoffErr
 }
 
-// Success transitions an InFlight payment to Completed, otherwise it returns an
+// RegisterAttempt atomically records the provided PaymentAttemptInfo to the
-// error. After calling Success, ClearForTakeoff should prevent any further
+// DB.
-// attempts for the same payment hash.
+func (p *paymentControl) RegisterAttempt(paymentHash lntypes.Hash,
-func (p *paymentControl) Success(paymentHash [32]byte) error {
+	attempt *PaymentAttemptInfo) error {
+
+	// Serialize the information before opening the db transaction.
+	var a bytes.Buffer
+	if err := serializePaymentAttemptInfo(&a, attempt); err != nil {
+		return err
+	}
+	attemptBytes := a.Bytes()
+
+	var updateErr error
+	err := p.db.Batch(func(tx *bbolt.Tx) error {
+		// Reset the update error, to avoid carrying over an error
+		// from a previous execution of the batched db transaction.
+		updateErr = nil
+
+		bucket, err := fetchPaymentBucket(tx, paymentHash)
+		if err != nil {
+			return err
+		}
+
+		// We can only register attempts for payments that are
+		// in-flight.
+		if err := ensureInFlight(bucket); err != nil {
+			updateErr = err
+			return nil
+		}
+
+		// Add the payment attempt to the payments bucket.
+		return bucket.Put(paymentAttemptInfoKey, attemptBytes)
+	})
+	if err != nil {
+		return err
+	}
+
+	return updateErr
+}
+
+// Success transitions a payment into the Completed state. After invoking this
+// method, InitPayment should always return an error to prevent us from making
+// duplicate payments to the same payment hash. The provided preimage is
+// atomically saved to the DB for record keeping.
+func (p *paymentControl) Success(paymentHash lntypes.Hash,
+	preimage lntypes.Preimage) error {
+
 	var updateErr error
 	err := p.db.Batch(func(tx *bbolt.Tx) error {
 		// Reset the update error, to avoid carrying over an error
@@ -144,6 +234,13 @@ func (p *paymentControl) Success(paymentHash [32]byte) error {
 			return nil
 		}
 
+		// Record the successful payment info atomically to the
+		// payments record.
+		err = bucket.Put(paymentSettleInfoKey, preimage[:])
+		if err != nil {
+			return err
+		}
+
 		return bucket.Put(paymentStatusKey, StatusCompleted.Bytes())
 	})
 	if err != nil {
@@ -151,12 +248,13 @@ func (p *paymentControl) Success(paymentHash [32]byte) error {
 	}
 
 	return updateErr
+
 }
 
-// Fail transitions an InFlight payment to Grounded, otherwise it returns an
+// Fail transitions a payment into the Failed state. After invoking this
-// error. After calling Fail, ClearForTakeoff should fail any further attempts
+// method, InitPayment should return nil on its next call for this payment
-// for the same payment hash.
+// hash, allowing the switch to make a subsequent payment.
-func (p *paymentControl) Fail(paymentHash [32]byte) error {
+func (p *paymentControl) Fail(paymentHash lntypes.Hash) error {
 	var updateErr error
 	err := p.db.Batch(func(tx *bbolt.Tx) error {
 		// Reset the update error, to avoid carrying over an error
@@ -174,7 +272,9 @@ func (p *paymentControl) Fail(paymentHash [32]byte) error {
 			return nil
 		}
 
-		return bucket.Put(paymentStatusKey, StatusGrounded.Bytes())
+		// A failed response was received for an InFlight payment, mark
+		// it as Failed to allow subsequent attempts.
+		return bucket.Put(paymentStatusKey, StatusFailed.Bytes())
 	})
 	if err != nil {
 		return err
@@ -196,6 +296,24 @@ func fetchPaymentBucket(tx *bbolt.Tx, paymentHash lntypes.Hash) (
 	return payments.CreateBucketIfNotExists(paymentHash[:])
 }
 
+// nextPaymentSequence returns the next sequence number to store for a new
+// payment.
+func nextPaymentSequence(tx *bbolt.Tx) ([]byte, error) {
+	payments, err := tx.CreateBucketIfNotExists(paymentsRootBucket)
+	if err != nil {
+		return nil, err
+	}
+
+	seq, err := payments.NextSequence()
+	if err != nil {
+		return nil, err
+	}
+
+	b := make([]byte, 8)
+	binary.BigEndian.PutUint64(b, seq)
+	return b, nil
+}
+
 // fetchPaymentStatus fetches the payment status from the bucket.  If the
 // status isn't found, it will default to "StatusGrounded".
 func fetchPaymentStatus(bucket *bbolt.Bucket) PaymentStatus {

channeldb/control_tower_test.go

@@ -6,10 +6,11 @@ import (
 	"io"
 	"io/ioutil"
 	"testing"
+	"time"
 
 	"github.com/btcsuite/fastsha256"
 	"github.com/coreos/bbolt"
-	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/lightningnetwork/lnd/lntypes"
 )
 
 func initDB() (*DB, error) {
@@ -34,19 +35,27 @@ func genPreimage() ([32]byte, error) {
 	return preimage, nil
 }
 
-func genHtlc() (*lnwire.UpdateAddHTLC, error) {
+func genInfo() (*PaymentCreationInfo, *PaymentAttemptInfo,
+	lntypes.Preimage, error) {
+
 	preimage, err := genPreimage()
 	if err != nil {
-		return nil, fmt.Errorf("unable to generate preimage: %v", err)
+		return nil, nil, preimage, fmt.Errorf("unable to "+
+			"generate preimage: %v", err)
 	}
 
 	rhash := fastsha256.Sum256(preimage[:])
-	htlc := &lnwire.UpdateAddHTLC{
+	return &PaymentCreationInfo{
-		PaymentHash: rhash,
+			PaymentHash:    rhash,
-		Amount:      1,
+			Value:          1,
-	}
+			CreationDate:   time.Unix(time.Now().Unix(), 0),
-
+			PaymentRequest: []byte("hola"),
-	return htlc, nil
+		},
+		&PaymentAttemptInfo{
+			PaymentID:  1,
+			SessionKey: priv,
+			Route:      testRoute,
+		}, preimage, nil
 }
 
 type paymentControlTestCase func(*testing.T)
@@ -95,44 +104,47 @@ func testPaymentControlSwitchFail(t *testing.T) {
 
 	pControl := NewPaymentControl(db)
 
-	htlc, err := genHtlc()
+	info, _, preimg, err := genInfo()
 	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 {
+	err = pControl.InitPayment(info.PaymentHash, info)
+	if err != nil {
 		t.Fatalf("unable to send htlc message: %v", err)
 	}
 
-	assertPaymentStatus(t, db, htlc.PaymentHash, StatusInFlight)
+	assertPaymentStatus(t, db, info.PaymentHash, StatusInFlight)
 
 	// Fail the payment, which should moved it to Grounded.
-	if err := pControl.Fail(htlc.PaymentHash); err != nil {
+	if err := pControl.Fail(info.PaymentHash); err != nil {
 		t.Fatalf("unable to fail payment hash: %v", err)
 	}
 
-	// Verify the status is indeed Grounded.
+	// Verify the status is indeed Failed.
-	assertPaymentStatus(t, db, htlc.PaymentHash, StatusGrounded)
+	assertPaymentStatus(t, db, info.PaymentHash, StatusFailed)
 
 	// Sends the htlc again, which should succeed since the prior payment
 	// failed.
-	if err := pControl.ClearForTakeoff(htlc); err != nil {
+	err = pControl.InitPayment(info.PaymentHash, info)
+	if err != nil {
 		t.Fatalf("unable to send htlc message: %v", err)
 	}
 
-	assertPaymentStatus(t, db, htlc.PaymentHash, StatusInFlight)
+	assertPaymentStatus(t, db, info.PaymentHash, StatusInFlight)
 
 	// Verifies that status was changed to StatusCompleted.
-	if err := pControl.Success(htlc.PaymentHash); err != nil {
+	if err := pControl.Success(info.PaymentHash, preimg); err != nil {
 		t.Fatalf("error shouldn't have been received, got: %v", err)
 	}
 
-	assertPaymentStatus(t, db, htlc.PaymentHash, StatusCompleted)
+	assertPaymentStatus(t, db, info.PaymentHash, StatusCompleted)
 
 	// Attempt a final payment, which should now fail since the prior
 	// payment succeed.
-	if err := pControl.ClearForTakeoff(htlc); err != ErrAlreadyPaid {
+	err = pControl.InitPayment(info.PaymentHash, info)
+	if err != ErrAlreadyPaid {
 		t.Fatalf("unable to send htlc message: %v", err)
 	}
 }
@@ -149,23 +161,25 @@ func testPaymentControlSwitchDoubleSend(t *testing.T) {
 
 	pControl := NewPaymentControl(db)
 
-	htlc, err := genHtlc()
+	info, _, _, err := genInfo()
 	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 {
+	err = pControl.InitPayment(info.PaymentHash, info)
+	if err != nil {
 		t.Fatalf("unable to send htlc message: %v", err)
 	}
 
-	assertPaymentStatus(t, db, htlc.PaymentHash, StatusInFlight)
+	assertPaymentStatus(t, db, info.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 {
+	err = pControl.InitPayment(info.PaymentHash, info)
+	if err != ErrPaymentInFlight {
 		t.Fatalf("payment control wrong behaviour: " +
 			"double sending must trigger ErrPaymentInFlight error")
 	}
@@ -183,28 +197,31 @@ func testPaymentControlSwitchDoublePay(t *testing.T) {
 
 	pControl := NewPaymentControl(db)
 
-	htlc, err := genHtlc()
+	info, _, preimg, err := genInfo()
 	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 {
+	err = pControl.InitPayment(info.PaymentHash, info)
+	if err != nil {
 		t.Fatalf("unable to send htlc message: %v", err)
 	}
 
 	// Verify that payment is InFlight.
-	assertPaymentStatus(t, db, htlc.PaymentHash, StatusInFlight)
+	assertPaymentStatus(t, db, info.PaymentHash, StatusInFlight)
 
 	// Move payment to completed status, second payment should return error.
-	if err := pControl.Success(htlc.PaymentHash); err != nil {
+	err = pControl.Success(info.PaymentHash, preimg)
+	if err != nil {
 		t.Fatalf("error shouldn't have been received, got: %v", err)
 	}
 
 	// Verify that payment is Completed.
-	assertPaymentStatus(t, db, htlc.PaymentHash, StatusCompleted)
+	assertPaymentStatus(t, db, info.PaymentHash, StatusCompleted)
 
-	if err := pControl.ClearForTakeoff(htlc); err != ErrAlreadyPaid {
+	err = pControl.InitPayment(info.PaymentHash, info)
+	if err != ErrAlreadyPaid {
 		t.Fatalf("payment control wrong behaviour:" +
 			" double payment must trigger ErrAlreadyPaid")
 	}
@@ -222,17 +239,17 @@ func TestPaymentControlStrictSuccessesWithoutInFlight(t *testing.T) {
 
 	pControl := NewPaymentControl(db)
 
-	htlc, err := genHtlc()
+	info, _, preimg, err := genInfo()
 	if err != nil {
 		t.Fatalf("unable to generate htlc message: %v", err)
 	}
 
-	err = pControl.Success(htlc.PaymentHash)
+	err = pControl.Success(info.PaymentHash, preimg)
 	if err != ErrPaymentNotInitiated {
 		t.Fatalf("expected ErrPaymentNotInitiated, got %v", err)
 	}
 
-	assertPaymentStatus(t, db, htlc.PaymentHash, StatusGrounded)
+	assertPaymentStatus(t, db, info.PaymentHash, StatusGrounded)
 }
 
 // TestPaymentControlStrictFailsWithoutInFlight checks that a strict payment
@@ -247,17 +264,17 @@ func TestPaymentControlStrictFailsWithoutInFlight(t *testing.T) {
 
 	pControl := NewPaymentControl(db)
 
-	htlc, err := genHtlc()
+	info, _, _, err := genInfo()
 	if err != nil {
 		t.Fatalf("unable to generate htlc message: %v", err)
 	}
 
-	err = pControl.Fail(htlc.PaymentHash)
+	err = pControl.Fail(info.PaymentHash)
 	if err != ErrPaymentNotInitiated {
 		t.Fatalf("expected ErrPaymentNotInitiated, got %v", err)
 	}
 
-	assertPaymentStatus(t, db, htlc.PaymentHash, StatusGrounded)
+	assertPaymentStatus(t, db, info.PaymentHash, StatusGrounded)
 }
 
 func assertPaymentStatus(t *testing.T, db *DB,