channeldb/witness_cache: create AddSha256Witnesses helper + test

channeldb/witness_cache.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 
 	"github.com/coreos/bbolt"
+	"github.com/lightningnetwork/lnd/lntypes"
 )
 
 var (
@@ -77,6 +78,24 @@ type witnessEntry struct {
 	witness []byte
 }
 
+// AddSha256Witnesses adds a batch of new sha256 preimages into the witness
+// cache. This is an alias for AddWitnesses that uses Sha256HashWitness as the
+// preimages' witness type.
+func (w *WitnessCache) AddSha256Witnesses(preimages ...lntypes.Preimage) error {
+	// Optimistically compute the preimages' hashes before attempting to
+	// start the db transaction.
+	entries := make([]witnessEntry, 0, len(preimages))
+	for i := range preimages {
+		hash := preimages[i].Hash()
+		entries = append(entries, witnessEntry{
+			key:     hash[:],
+			witness: preimages[i][:],
+		})
+	}
+
+	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.
@@ -101,6 +120,15 @@ func (w *WitnessCache) AddWitnesses(wType WitnessType, witnesses ...[]byte) erro
 		}
 	}
 
+	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.
+func (w *WitnessCache) addWitnessEntries(wType WitnessType,
+	entries []witnessEntry) error {
+
 	// Exit early if there are no witnesses to add.
 	if len(entries) == 0 {
 		return nil

channeldb/witness_cache_test.go

@@ -1,9 +1,12 @@
 package channeldb
 
 import (
+	"bytes"
 	"crypto/sha256"
 	"reflect"
 	"testing"
+
+	"github.com/lightningnetwork/lnd/lntypes"
 )
 
 // TestWitnessCacheRetrieval tests that we're able to add and lookup new
@@ -125,3 +128,85 @@ func TestWitnessCacheUnknownWitness(t *testing.T) {
 		t.Fatalf("expected ErrUnknownWitnessType, got %v", err)
 	}
 }
+
+// TestAddSha256Witnesses tests that insertion using AddSha256Witnesses behaves
+// identically to the insertion via the generalized interface.
+func TestAddSha256Witnesses(t *testing.T) {
+	cdb, cleanUp, err := makeTestDB()
+	if err != nil {
+		t.Fatalf("unable to make test database: %v", err)
+	}
+	defer cleanUp()
+
+	wCache := cdb.NewWitnessCache()
+
+	// We'll start by adding a witnesses to the cache using the generic
+	// AddWitnesses method.
+	witness1 := rev[:]
+	witness1Key := sha256.Sum256(witness1)
+
+	witness2 := key[:]
+	witness2Key := sha256.Sum256(witness2)
+
+	var (
+		preimages = []lntypes.Preimage{rev, key}
+		witnesses = [][]byte{witness1, witness2}
+		keys      = [][]byte{witness1Key[:], witness2Key[:]}
+	)
+
+	err = wCache.AddWitnesses(Sha256HashWitness, witnesses...)
+	if err != nil {
+		t.Fatalf("unable to add witness: %v", err)
+	}
+
+	for i, key := range keys {
+		witness := witnesses[i]
+
+		dbWitness, err := wCache.LookupWitness(
+			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 {
+			t.Fatalf("retrieved witness mismatch, want: %x, "+
+				"got: %x", witness, dbWitness)
+		}
+
+		// We'll now delete the witness, as we'll be reinserting it
+		// using the specialized AddSha256Witnesses method.
+		err = wCache.DeleteWitness(Sha256HashWitness, key)
+		if err != nil {
+			t.Fatalf("unable to delete witness: %v", err)
+		}
+	}
+
+	// Now, add the same witnesses using the type-safe interface for
+	// lntypes.Preimages..
+	err = wCache.AddSha256Witnesses(preimages...)
+	if err != nil {
+		t.Fatalf("unable to add sha256 preimage: %v", err)
+	}
+
+	// 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 {
+		witness := witnesses[i]
+
+		dbWitness, err := wCache.LookupWitness(
+			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 {
+			t.Fatalf("retrieved witness mismatch, want: %x, "+
+				"got: %x", witness, dbWitness)
+		}
+	}
+}