channeldb+witness_beacon: use sha256 lookup+delete witness

channeldb/witness_cache.go

@@ -1,7 +1,6 @@
 package channeldb
 
 import (
-	"crypto/sha256"
 	"fmt"
 
 	"github.com/coreos/bbolt"
@@ -96,33 +95,6 @@ func (w *WitnessCache) AddSha256Witnesses(preimages ...lntypes.Preimage) error {
 	return w.addWitnessEntries(Sha256HashWitness, entries)
 }
 
-// AddWitnesses adds a batch of new witnesses of wType to the witness cache. The
-// type of the witness will be used to map each witness to the key that will be
-// used to look it up. All witnesses should be of the same WitnessType.
-//
-// TODO(roasbeef): fake closure to map instead a constructor?
-func (w *WitnessCache) AddWitnesses(wType WitnessType, witnesses ...[]byte) error {
-	// Optimistically compute the witness keys before attempting to start
-	// the db transaction.
-	entries := make([]witnessEntry, 0, len(witnesses))
-	for _, witness := range witnesses {
-		// Map each witness to its key by applying the appropriate
-		// transformation for the given witness type.
-		switch wType {
-		case Sha256HashWitness:
-			key := sha256.Sum256(witness)
-			entries = append(entries, witnessEntry{
-				key:     key[:],
-				witness: witness,
-			})
-		default:
-			return ErrUnknownWitnessType
-		}
-	}
-
-	return w.addWitnessEntries(wType, entries)
-}
-
 // addWitnessEntries inserts the witnessEntry key-value pairs into the cache,
 // using the appropriate witness type to segment the namespace of possible
 // witness types.
@@ -162,10 +134,21 @@ func (w *WitnessCache) addWitnessEntries(wType WitnessType,
 	})
 }
 
-// LookupWitness attempts to lookup a witness according to its type and also
+// LookupSha256Witness attempts to lookup the preimage for a sha256 hash. If
+// the witness isn't found, ErrNoWitnesses will be returned.
+func (w *WitnessCache) LookupSha256Witness(hash lntypes.Hash) (lntypes.Preimage, error) {
+	witness, err := w.lookupWitness(Sha256HashWitness, hash[:])
+	if err != nil {
+		return lntypes.Preimage{}, err
+	}
+
+	return lntypes.MakePreimage(witness)
+}
+
+// lookupWitness attempts to lookup a witness according to its type and also
 // its witness key. In the case that the witness isn't found, ErrNoWitnesses
 // will be returned.
-func (w *WitnessCache) LookupWitness(wType WitnessType, witnessKey []byte) ([]byte, error) {
+func (w *WitnessCache) lookupWitness(wType WitnessType, witnessKey []byte) ([]byte, error) {
 	var witness []byte
 	err := w.db.View(func(tx *bbolt.Tx) error {
 		witnessBucket := tx.Bucket(witnessBucketKey)
@@ -199,8 +182,13 @@ func (w *WitnessCache) LookupWitness(wType WitnessType, witnessKey []byte) ([]by
 	return witness, nil
 }
 
-// DeleteWitness attempts to delete a particular witness from the database.
+// DeleteSha256Witness attempts to delete a sha256 preimage identified by hash.
-func (w *WitnessCache) DeleteWitness(wType WitnessType, witnessKey []byte) error {
+func (w *WitnessCache) DeleteSha256Witness(hash lntypes.Hash) error {
+	return w.deleteWitness(Sha256HashWitness, hash[:])
+}
+
+// deleteWitness attempts to delete a particular witness from the database.
+func (w *WitnessCache) deleteWitness(wType WitnessType, witnessKey []byte) error {
 	return w.db.Batch(func(tx *bbolt.Tx) error {
 		witnessBucket, err := tx.CreateBucketIfNotExists(witnessBucketKey)
 		if err != nil {

channeldb/witness_cache_test.go

@@ -1,17 +1,15 @@
 package channeldb
 
 import (
-	"bytes"
 	"crypto/sha256"
-	"reflect"
 	"testing"
 
 	"github.com/lightningnetwork/lnd/lntypes"
 )
 
-// TestWitnessCacheRetrieval tests that we're able to add and lookup new
+// TestWitnessCacheSha256Retrieval tests that we're able to add and lookup new
-// witnesses to the witness cache.
+// sha256 preimages to the witness cache.
-func TestWitnessCacheRetrieval(t *testing.T) {
+func TestWitnessCacheSha256Retrieval(t *testing.T) {
 	t.Parallel()
 
 	cdb, cleanUp, err := makeTestDB()
@@ -22,44 +20,41 @@ func TestWitnessCacheRetrieval(t *testing.T) {
 
 	wCache := cdb.NewWitnessCache()
 
-	// We'll be attempting to add then lookup two simple hash witnesses
+	// We'll be attempting to add then lookup two simple sha256 preimages
 	// within this test.
-	witness1 := rev[:]
+	preimage1 := lntypes.Preimage(rev)
-	witness1Key := sha256.Sum256(witness1)
+	preimage2 := lntypes.Preimage(key)
 
-	witness2 := key[:]
+	preimages := []lntypes.Preimage{preimage1, preimage2}
-	witness2Key := sha256.Sum256(witness2)
+	hashes := []lntypes.Hash{preimage1.Hash(), preimage2.Hash()}
 
-	witnesses := [][]byte{witness1, witness2}
+	// First, we'll attempt to add the preimages to the database.
-	keys := [][]byte{witness1Key[:], witness2Key[:]}
+	err = wCache.AddSha256Witnesses(preimages...)
-
-	// First, we'll attempt to add the witnesses to the database.
-	err = wCache.AddWitnesses(Sha256HashWitness, witnesses...)
 	if err != nil {
 		t.Fatalf("unable to add witness: %v", err)
 	}
 
-	// With the witnesses stored, we'll now attempt to look them up.
+	// With the preimages stored, we'll now attempt to look them up.
-	for i, key := range keys {
+	for i, hash := range hashes {
-		witness := witnesses[i]
+		preimage := preimages[i]
 
-		// We should get back the *exact* same witness as we originally
+		// We should get back the *exact* same preimage as we originally
 		// stored.
-		dbWitness, err := wCache.LookupWitness(Sha256HashWitness, key)
+		dbPreimage, err := wCache.LookupSha256Witness(hash)
 		if err != nil {
 			t.Fatalf("unable to look up witness: %v", err)
 		}
 
-		if !reflect.DeepEqual(witness, dbWitness[:]) {
+		if preimage != dbPreimage {
 			t.Fatalf("witnesses don't match: expected %x, got %x",
-				witness[:], dbWitness[:])
+				preimage[:], dbPreimage[:])
 		}
 	}
 }
 
-// TestWitnessCacheDeletion tests that we're able to delete a single witness,
+// TestWitnessCacheSha256Deletion tests that we're able to delete a single
-// and also a class of witnesses from the cache.
+// sha256 preimage, and also a class of witnesses from the cache.
-func TestWitnessCacheDeletion(t *testing.T) {
+func TestWitnessCacheSha256Deletion(t *testing.T) {
 	t.Parallel()
 
 	cdb, cleanUp, err := makeTestDB()
@@ -70,39 +65,39 @@ func TestWitnessCacheDeletion(t *testing.T) {
 
 	wCache := cdb.NewWitnessCache()
 
-	// We'll start by adding two witnesses to the cache.
+	// We'll start by adding two preimages to the cache.
-	witness1 := rev[:]
+	preimage1 := lntypes.Preimage(key)
-	witness1Key := sha256.Sum256(witness1)
+	hash1 := preimage1.Hash()
 
-	if err := wCache.AddWitnesses(Sha256HashWitness, witness1); err != nil {
+	preimage2 := lntypes.Preimage(rev)
+	hash2 := preimage2.Hash()
+
+	if err := wCache.AddSha256Witnesses(preimage1); err != nil {
 		t.Fatalf("unable to add witness: %v", err)
 	}
 
-	witness2 := key[:]
+	if err := wCache.AddSha256Witnesses(preimage2); err != nil {
-	witness2Key := sha256.Sum256(witness2)
-
-	if err := wCache.AddWitnesses(Sha256HashWitness, witness2); err != nil {
 		t.Fatalf("unable to add witness: %v", err)
 	}
 
-	// We'll now delete the first witness. If we attempt to look it up, we
+	// We'll now delete the first preimage. If we attempt to look it up, we
 	// should get ErrNoWitnesses.
-	err = wCache.DeleteWitness(Sha256HashWitness, witness1Key[:])
+	err = wCache.DeleteSha256Witness(hash1)
 	if err != nil {
 		t.Fatalf("unable to delete witness: %v", err)
 	}
-	_, err = wCache.LookupWitness(Sha256HashWitness, witness1Key[:])
+	_, err = wCache.LookupSha256Witness(hash1)
 	if err != ErrNoWitnesses {
 		t.Fatalf("expected ErrNoWitnesses instead got: %v", err)
 	}
 
 	// Next, we'll attempt to delete the entire witness class itself. When
-	// we try to lookup the second witness, we should again get
+	// we try to lookup the second preimage, we should again get
 	// ErrNoWitnesses.
 	if err := wCache.DeleteWitnessClass(Sha256HashWitness); err != nil {
 		t.Fatalf("unable to delete witness class: %v", err)
 	}
-	_, err = wCache.LookupWitness(Sha256HashWitness, witness2Key[:])
+	_, err = wCache.LookupSha256Witness(hash2)
 	if err != ErrNoWitnesses {
 		t.Fatalf("expected ErrNoWitnesses instead got: %v", err)
 	}
@@ -123,7 +118,7 @@ func TestWitnessCacheUnknownWitness(t *testing.T) {
 
 	// We'll attempt to add a new, undefined witness type to the database.
 	// We should get an error.
-	err = wCache.AddWitnesses(234, key[:])
+	err = wCache.legacyAddWitnesses(234, key[:])
 	if err != ErrUnknownWitnessType {
 		t.Fatalf("expected ErrUnknownWitnessType, got %v", err)
 	}
@@ -143,41 +138,41 @@ func TestAddSha256Witnesses(t *testing.T) {
 	// We'll start by adding a witnesses to the cache using the generic
 	// AddWitnesses method.
 	witness1 := rev[:]
-	witness1Key := sha256.Sum256(witness1)
+	preimage1 := lntypes.Preimage(rev)
+	hash1 := preimage1.Hash()
 
 	witness2 := key[:]
-	witness2Key := sha256.Sum256(witness2)
+	preimage2 := lntypes.Preimage(key)
+	hash2 := preimage2.Hash()
 
 	var (
-		preimages = []lntypes.Preimage{rev, key}
 		witnesses = [][]byte{witness1, witness2}
-		keys      = [][]byte{witness1Key[:], witness2Key[:]}
+		preimages = []lntypes.Preimage{preimage1, preimage2}
+		hashes    = []lntypes.Hash{hash1, hash2}
 	)
 
-	err = wCache.AddWitnesses(Sha256HashWitness, witnesses...)
+	err = wCache.legacyAddWitnesses(Sha256HashWitness, witnesses...)
 	if err != nil {
 		t.Fatalf("unable to add witness: %v", err)
 	}
 
-	for i, key := range keys {
+	for i, hash := range hashes {
-		witness := witnesses[i]
+		preimage := preimages[i]
 
-		dbWitness, err := wCache.LookupWitness(
+		dbPreimage, err := wCache.LookupSha256Witness(hash)
-			Sha256HashWitness, key,
-		)
 		if err != nil {
 			t.Fatalf("unable to lookup witness: %v", err)
 		}
 
 		// Assert that the retrieved witness matches the original.
-		if bytes.Compare(dbWitness, witness) != 0 {
+		if dbPreimage != preimage {
 			t.Fatalf("retrieved witness mismatch, want: %x, "+
-				"got: %x", witness, dbWitness)
+				"got: %x", preimage, dbPreimage)
 		}
 
 		// We'll now delete the witness, as we'll be reinserting it
 		// using the specialized AddSha256Witnesses method.
-		err = wCache.DeleteWitness(Sha256HashWitness, key)
+		err = wCache.DeleteSha256Witness(hash)
 		if err != nil {
 			t.Fatalf("unable to delete witness: %v", err)
 		}
@@ -193,20 +188,51 @@ func TestAddSha256Witnesses(t *testing.T) {
 	// Finally, iterate over the keys and assert that the returned witnesses
 	// match the original witnesses. This asserts that the specialized
 	// insertion method behaves identically to the generalized interface.
-	for i, key := range keys {
+	for i, hash := range hashes {
-		witness := witnesses[i]
+		preimage := preimages[i]
 
-		dbWitness, err := wCache.LookupWitness(
+		dbPreimage, err := wCache.LookupSha256Witness(hash)
-			Sha256HashWitness, key,
-		)
 		if err != nil {
 			t.Fatalf("unable to lookup witness: %v", err)
 		}
 
 		// Assert that the retrieved witness matches the original.
-		if bytes.Compare(dbWitness, witness) != 0 {
+		if dbPreimage != preimage {
 			t.Fatalf("retrieved witness mismatch, want: %x, "+
-				"got: %x", witness, dbWitness)
+				"got: %x", preimage, dbPreimage)
 		}
 	}
 }
+
+// legacyAddWitnesses adds a batch of new witnesses of wType to the witness
+// cache. The type of the witness will be used to map each witness to the key
+// that will be used to look it up. All witnesses should be of the same
+// WitnessType.
+//
+// NOTE: Previously this method exposed a generic interface for adding
+// witnesses, which has since been deprecated in favor of a strongly typed
+// interface for each witness class. We keep this method around to assert the
+// correctness of specialized witness adding methods.
+func (w *WitnessCache) legacyAddWitnesses(wType WitnessType,
+	witnesses ...[]byte) error {
+
+	// Optimistically compute the witness keys before attempting to start
+	// the db transaction.
+	entries := make([]witnessEntry, 0, len(witnesses))
+	for _, witness := range witnesses {
+		// Map each witness to its key by applying the appropriate
+		// transformation for the given witness type.
+		switch wType {
+		case Sha256HashWitness:
+			key := sha256.Sum256(witness)
+			entries = append(entries, witnessEntry{
+				key:     key[:],
+				witness: witness,
+			})
+		default:
+			return ErrUnknownWitnessType
+		}
+	}
+
+	return w.addWitnessEntries(wType, entries)
+}

witness_beacon.go

@@ -90,20 +90,12 @@ func (p *preimageBeacon) LookupPreimage(
 	}
 
 	// Otherwise, we'll perform a final check using the witness cache.
-	preimageBytes, err := p.wCache.LookupWitness(
+	preimage, err := p.wCache.LookupSha256Witness(payHash)
-		channeldb.Sha256HashWitness, payHash[:],
-	)
 	if err != nil {
 		ltndLog.Errorf("Unable to lookup witness: %v", err)
 		return lntypes.Preimage{}, false
 	}
 
-	preimage, err := lntypes.MakePreimage(preimageBytes)
-	if err != nil {
-		ltndLog.Errorf("Unable to build witness: %v", err)
-		return lntypes.Preimage{}, false
-	}
-
 	return preimage, true
 }