Implement basic in-memory ReplayLog for use in tests.

Jim Posen · Jim Posen · commit d667f7a1e135 · 2018-03-26T11:39:01.000-07:00
diff --git a/replaylog.go b/replaylog.go
@@ -2,6 +2,7 @@ package sphinx
 
 import (
 	"crypto/sha256"
+	"errors"
 )
 
 const (
@@ -15,6 +16,16 @@ const (
 // of a Hash256, and is used to detect duplicate sphinx packets.
 type HashPrefix [HashPrefixSize]byte
 
+// errReplayLogAlreadyStarted is an error returned when Start() is called on a
+// ReplayLog after it is started and before it is stopped.
+var errReplayLogAlreadyStarted error = errors.New(
+	"Replay log has already been started")
+
+// errReplayLogNotStarted is an error returned when methods other than Start()
+// are called on a ReplayLog before it is started or after it is stopped.
+var errReplayLogNotStarted error = errors.New(
+	"Replay log has not been started")
+
 // hashSharedSecret Sha-256 hashes the shared secret and returns the first
 // HashPrefixSize bytes of the hash.
 func hashSharedSecret(sharedSecret *Hash256) *HashPrefix {
@@ -58,3 +69,119 @@ type ReplayLog interface {
 	// that are replays and an error if one occurs.
 	PutBatch(*Batch) (*ReplaySet, error)
 }
+
+// MemoryReplayLog is a simple ReplayLog implementation that stores all added
+// sphinx packets and processed batches in memory with no persistence.
+//
+// This is designed for use just in testing.
+type MemoryReplayLog struct {
+	batches map[string]*ReplaySet
+	entries map[HashPrefix]uint32
+}
+
+// NewMemoryReplayLog constructs a new MemoryReplayLog.
+func NewMemoryReplayLog() *MemoryReplayLog {
+	return &MemoryReplayLog{}
+}
+
+// Start initializes the log and must be called before any other methods.
+func (rl *MemoryReplayLog) Start() error {
+	rl.batches = make(map[string]*ReplaySet)
+	rl.entries = make(map[HashPrefix]uint32)
+	return nil
+}
+
+// Stop wipes the state of the log.
+func (rl *MemoryReplayLog) Stop() error {
+	if rl.entries == nil || rl.batches == nil {
+		return errReplayLogNotStarted
+	}
+
+	rl.batches = nil
+	rl.entries = nil
+	return nil
+}
+
+// Get retrieves an entry from the log given its hash prefix. It returns the
+// value stored and an error if one occurs. It returns ErrLogEntryNotFound
+// if the entry is not in the log.
+func (rl *MemoryReplayLog) Get(hash *HashPrefix) (uint32, error) {
+	if rl.entries == nil || rl.batches == nil {
+		return 0, errReplayLogNotStarted
+	}
+
+	cltv, exists := rl.entries[*hash]
+	if !exists {
+		return 0, ErrLogEntryNotFound
+	}
+
+	return cltv, nil
+}
+
+// Put stores an entry into the log given its hash prefix and an accompanying
+// purposefully general type. It returns ErrReplayedPacket if the provided hash
+// prefix already exists in the log.
+func (rl *MemoryReplayLog) Put(hash *HashPrefix, cltv uint32) error {
+	if rl.entries == nil || rl.batches == nil {
+		return errReplayLogNotStarted
+	}
+
+	_, exists := rl.entries[*hash]
+	if exists {
+		return ErrReplayedPacket
+	}
+
+	rl.entries[*hash] = cltv
+	return nil
+}
+
+// Delete deletes an entry from the log given its hash prefix.
+func (rl *MemoryReplayLog) Delete(hash *HashPrefix) error {
+	if rl.entries == nil || rl.batches == nil {
+		return errReplayLogNotStarted
+	}
+
+	delete(rl.entries, *hash)
+	return nil
+}
+
+// PutBatch stores a batch of sphinx packets into the log given their hash
+// prefixes and accompanying values. Returns the set of entries in the batch
+// that are replays and an error if one occurs.
+func (rl *MemoryReplayLog) PutBatch(batch *Batch) (*ReplaySet, error) {
+	if rl.entries == nil || rl.batches == nil {
+		return nil, errReplayLogNotStarted
+	}
+
+	// Return the result when the batch was first processed to provide
+	// idempotence.
+	replays, exists := rl.batches[string(batch.id)]
+
+	if !exists {
+		replays = NewReplaySet()
+		for seqNum, entry := range batch.entries {
+			err := rl.Put(&entry.hashPrefix, entry.cltv)
+			if err == ErrReplayedPacket {
+				replays.Add(seqNum)
+				continue
+			}
+
+			// This would be bad because we have already updated the entries
+			// map, but no errors other than ErrReplayedPacket should occur.
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		replays.Merge(batch.replaySet)
+		rl.batches[string(batch.id)] = replays
+	}
+
+	batch.replaySet = replays
+	batch.isCommitted = true
+
+	return replays, nil
+}
+
+// A compile time asserting *MemoryReplayLog implements the RelayLog interface.
+var _ ReplayLog = (*MemoryReplayLog)(nil)
diff --git a/replaylog_test.go b/replaylog_test.go
@@ -0,0 +1,132 @@
+package sphinx
+
+import (
+	"testing"
+)
+
+// TestMemoryReplayLogStorageAndRetrieval tests that the non-batch methods on
+// MemoryReplayLog work as expected.
+func TestMemoryReplayLogStorageAndRetrieval(t *testing.T) {
+	rl := NewMemoryReplayLog()
+	rl.Start()
+	defer rl.Stop()
+
+	var hashPrefix HashPrefix
+	hashPrefix[0] = 1
+
+	var cltv1 uint32 = 1
+
+	// Attempt to lookup unknown sphinx packet.
+	_, err := rl.Get(&hashPrefix)
+	if err == nil {
+		t.Fatalf("Expected ErrLogEntryNotFound")
+	}
+	if err != ErrLogEntryNotFound {
+		t.Fatalf("Get failed - received unexpected error upon Get: %v", err)
+	}
+
+	// Log incoming sphinx packet.
+	err = rl.Put(&hashPrefix, cltv1)
+	if err != nil {
+		t.Fatalf("Put failed - received unexpected error upon Put: %v", err)
+	}
+
+	// Attempt to replay sphinx packet.
+	err = rl.Put(&hashPrefix, cltv1)
+	if err == nil {
+		t.Fatalf("Expected ErrReplayedPacket")
+	}
+	if err != ErrReplayedPacket {
+		t.Fatalf("Put failed - received unexpected error upon Put: %v", err)
+	}
+
+	// Lookup logged sphinx packet.
+	cltv, err := rl.Get(&hashPrefix)
+	if err != nil {
+		t.Fatalf("Get failed - received unexpected error upon Get: %v", err)
+	}
+	if cltv != cltv1 {
+		t.Fatalf("Get returned wrong value: expected %v, got %v", cltv1, cltv)
+	}
+
+	// Delete sphinx packet from log.
+	err = rl.Delete(&hashPrefix)
+	if err != nil {
+		t.Fatalf("Delete failed - received unexpected error upon Delete: %v", err)
+	}
+
+	// Attempt to lookup deleted sphinx packet.
+	_, err = rl.Get(&hashPrefix)
+	if err == nil {
+		t.Fatalf("Expected ErrLogEntryNotFound")
+	}
+	if err != ErrLogEntryNotFound {
+		t.Fatalf("Get failed - received unexpected error upon Get: %v", err)
+	}
+
+	// Reinsert incoming sphinx packet into the log.
+	var cltv2 uint32 = 2
+	err = rl.Put(&hashPrefix, cltv2)
+	if err != nil {
+		t.Fatalf("Put failed - received unexpected error upon Put: %v", err)
+	}
+
+	// Lookup logged sphinx packet.
+	cltv, err = rl.Get(&hashPrefix)
+	if err != nil {
+		t.Fatalf("Get failed - received unexpected error upon Get: %v", err)
+	}
+	if cltv != cltv2 {
+		t.Fatalf("Get returned wrong value: expected %v, got %v", cltv2, cltv)
+	}
+}
+
+// TestMemoryReplayLogPutBatch tests that the batch adding of packets to a log
+// works as expected.
+func TestMemoryReplayLogPutBatch(t *testing.T) {
+	rl := NewMemoryReplayLog()
+	rl.Start()
+	defer rl.Stop()
+
+	var hashPrefix1, hashPrefix2 HashPrefix
+	hashPrefix1[0] = 1
+	hashPrefix2[0] = 2
+
+	// Create a batch with a duplicated packet.
+	batch1 := NewBatch([]byte{1})
+	err := batch1.Put(1, &hashPrefix1, 1)
+	if err != nil {
+		t.Fatalf("Unexpected error adding entry to batch: %v", err)
+	}
+	err = batch1.Put(1, &hashPrefix1, 1)
+	if err != nil {
+		t.Fatalf("Unexpected error adding entry to batch: %v", err)
+	}
+
+	replays, err := rl.PutBatch(batch1)
+	if replays.Size() != 1 || !replays.Contains(1) {
+		t.Fatalf("Unexpected replay set after adding batch 1 to log: %v", err)
+	}
+
+	// Create a batch with one replayed packet and one valid one.
+	batch2 := NewBatch([]byte{2})
+	err = batch2.Put(1, &hashPrefix1, 1)
+	if err != nil {
+		t.Fatalf("Unexpected error adding entry to batch: %v", err)
+	}
+	err = batch2.Put(2, &hashPrefix2, 2)
+	if err != nil {
+		t.Fatalf("Unexpected error adding entry to batch: %v", err)
+	}
+
+	replays, err = rl.PutBatch(batch2)
+	if replays.Size() != 1 || !replays.Contains(1) {
+		t.Fatalf("Unexpected replay set after adding batch 2 to log: %v", err)
+	}
+
+	// Reprocess batch 2, which should be idempotent.
+	replays, err = rl.PutBatch(batch2)
+	if replays.Size() != 1 || !replays.Contains(1) {
+		t.Fatalf("Unexpected replay set after adding batch 2 to log: %v", err)
+	}
+}
