From f75e121d8ceca11a43a9025cf521d9a0d9b898f0 Mon Sep 17 00:00:00 2001
From: Jeff Mitchell <jeffrey.mitchell@gmail.com>
Date: Thu, 4 Feb 2016 09:40:35 -0500
Subject: [PATCH] Introduce a locking inmem storage for unit tests that are
 doing concurrent things

---
 builtin/logical/transit/backend_test.go |  2 +-
 logical/locking_inmem_storage.go        | 58 +++++++++++++++++++++++++
 logical/locking_inmem_storage_test.go   | 13 ++++++
 3 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 logical/locking_inmem_storage.go
 create mode 100644 logical/locking_inmem_storage_test.go

diff --git a/builtin/logical/transit/backend_test.go b/builtin/logical/transit/backend_test.go
index 039f32834e..4a1660db1c 100644
--- a/builtin/logical/transit/backend_test.go
+++ b/builtin/logical/transit/backend_test.go
@@ -549,7 +549,7 @@ func TestKeyUpgrade(t *testing.T) {
 func TestPolicyFuzzing(t *testing.T) {
 	be := Backend()
 
-	storage := &logical.InmemStorage{}
+	storage := &logical.LockingInmemStorage{}
 	wg := sync.WaitGroup{}
 
 	funcs := []string{"encrypt", "decrypt", "rotate", "change_min_version"}
diff --git a/logical/locking_inmem_storage.go b/logical/locking_inmem_storage.go
new file mode 100644
index 0000000000..9ed360f078
--- /dev/null
+++ b/logical/locking_inmem_storage.go
@@ -0,0 +1,58 @@
+package logical
+
+import (
+	"strings"
+	"sync"
+)
+
+// LockingInmemStorage implements Storage and stores all data in memory.
+type LockingInmemStorage struct {
+	sync.RWMutex
+
+	Data map[string]*StorageEntry
+
+	once sync.Once
+}
+
+func (s *LockingInmemStorage) List(prefix string) ([]string, error) {
+	s.once.Do(s.init)
+
+	s.RLock()
+	defer s.RUnlock()
+
+	var result []string
+	for k, _ := range s.Data {
+		if strings.HasPrefix(k, prefix) {
+			result = append(result, k)
+		}
+	}
+
+	return result, nil
+}
+
+func (s *LockingInmemStorage) Get(key string) (*StorageEntry, error) {
+	s.once.Do(s.init)
+	s.RLock()
+	defer s.RUnlock()
+	return s.Data[key], nil
+}
+
+func (s *LockingInmemStorage) Put(entry *StorageEntry) error {
+	s.once.Do(s.init)
+	s.Lock()
+	defer s.Unlock()
+	s.Data[entry.Key] = entry
+	return nil
+}
+
+func (s *LockingInmemStorage) Delete(k string) error {
+	s.once.Do(s.init)
+	s.Lock()
+	defer s.Unlock()
+	delete(s.Data, k)
+	return nil
+}
+
+func (s *LockingInmemStorage) init() {
+	s.Data = make(map[string]*StorageEntry)
+}
diff --git a/logical/locking_inmem_storage_test.go b/logical/locking_inmem_storage_test.go
new file mode 100644
index 0000000000..221fed4d35
--- /dev/null
+++ b/logical/locking_inmem_storage_test.go
@@ -0,0 +1,13 @@
+package logical
+
+import (
+	"testing"
+)
+
+// Note: This uses the normal TestStorage, but the best way to exercise this is
+// to run transit's unit tests, which spawn 1000 goroutines to hammer the
+// backend for 10 seconds with this as the storage.
+
+func TestLockingInmemStorage(t *testing.T) {
+	TestStorage(t, new(LockingInmemStorage))
+}