diff --git a/channeldb/witness_cache.go b/channeldb/witness_cache.go index d21d9373..c9c7a9e0 100644 --- a/channeldb/witness_cache.go +++ b/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 diff --git a/channeldb/witness_cache_test.go b/channeldb/witness_cache_test.go index 073e7d61..668317bb 100644 --- a/channeldb/witness_cache_test.go +++ b/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) + } + } +}