Add db package

Add own db (storage) package with db.Storage interface, with Memory & Leveldb
implementations, and with tests.

Adapted with several changes from
https://github.com/iden3/go-iden3-core/tree/v0.0.8/db

db/db.go

@@ -0,0 +1,71 @@
+package db
+
+import (
+	"bytes"
+	"crypto/sha256"
+	"errors"
+)
+
+// ErrNotFound is used by the implementations of the interface db.Storage for
+// when a key is not found in the storage
+var ErrNotFound = errors.New("key not found")
+
+// Storage is the interface that defines the methods for the storage used in
+// the merkletree. Examples of the interface implementation can be found at
+// db/memory and db/leveldb directories.
+type Storage interface {
+	NewTx() (Tx, error)
+	WithPrefix(prefix []byte) Storage
+	Get([]byte) ([]byte, error)
+	List(int) ([]KV, error)
+	Close()
+	Info() string
+	Iterate(func([]byte, []byte) (bool, error)) error
+}
+
+// Tx is the interface that defines the methods for the db transaction used in
+// the merkletree storage. Examples of the interface implementation can be
+// found at db/memory and db/leveldb directories.
+type Tx interface {
+	Get([]byte) ([]byte, error)
+	Put(k, v []byte)
+	Add(Tx)
+	Commit() error
+	Close()
+}
+
+// KV contains a key (K) and a value (V)
+type KV struct {
+	K []byte
+	V []byte
+}
+
+// KvMap is a key-value map between a sha256 byte array hash, and a KV struct
+type KvMap map[[sha256.Size]byte]KV
+
+// Get retreives the value respective to a key from the KvMap
+func (m KvMap) Get(k []byte) ([]byte, bool) {
+	v, ok := m[sha256.Sum256(k)]
+	return v.V, ok
+}
+
+// Put stores a key and a value in the KvMap
+func (m KvMap) Put(k, v []byte) {
+	m[sha256.Sum256(k)] = KV{k, v}
+}
+
+// Concat concatenates arrays of bytes
+func Concat(vs ...[]byte) []byte {
+	var b bytes.Buffer
+	for _, v := range vs {
+		b.Write(v)
+	}
+	return b.Bytes()
+}
+
+// Clone clones a byte array into a new byte array
+func Clone(b0 []byte) []byte {
+	b1 := make([]byte, len(b0))
+	copy(b1, b0)
+	return b1
+}

db/leveldb/leveldb.go

@@ -0,0 +1,186 @@
+package leveldb
+
+import (
+	"encoding/json"
+
+	log "github.com/sirupsen/logrus"
+	"github.com/syndtr/goleveldb/leveldb"
+	"github.com/syndtr/goleveldb/leveldb/errors"
+	"github.com/syndtr/goleveldb/leveldb/opt"
+	"github.com/syndtr/goleveldb/leveldb/util"
+
+	"github.com/iden3/go-merkletree/db"
+)
+
+// LevelDbStorage implements the db.Storage interface
+type LevelDbStorage struct {
+	ldb    *leveldb.DB
+	prefix []byte
+}
+
+// LevelDbStorageTx implements the db.Tx interface
+type LevelDbStorageTx struct {
+	*LevelDbStorage
+	cache db.KvMap
+}
+
+// NewLevelStorage returns a new LevelDbStorage
+func NewLevelDbStorage(path string, errorIfMissing bool) (*LevelDbStorage, error) {
+	o := &opt.Options{
+		ErrorIfMissing: errorIfMissing,
+	}
+	ldb, err := leveldb.OpenFile(path, o)
+	if err != nil {
+		return nil, err
+	}
+	return &LevelDbStorage{ldb, []byte{}}, nil
+}
+
+type storageInfo struct {
+	KeyCount   int
+	ClaimCount int
+}
+
+// Info implements the method Info of the interface db.Storage
+func (l *LevelDbStorage) Info() string {
+	snapshot, err := l.ldb.GetSnapshot()
+	if err != nil {
+		return err.Error()
+	}
+
+	keycount := 0
+	claimcount := 0
+	iter := snapshot.NewIterator(nil, nil)
+	for iter.Next() {
+		if iter.Value()[0] == byte(1) {
+			claimcount++
+		}
+
+		keycount++
+	}
+	iter.Release()
+	if err := iter.Error(); err != nil {
+		return err.Error()
+	}
+	json, _ := json.MarshalIndent(
+		storageInfo{
+			KeyCount:   keycount,
+			ClaimCount: claimcount,
+		},
+		"", "  ",
+	)
+	return string(json)
+}
+
+// WithPrefix implements the method WithPrefix of the interface db.Storage
+func (l *LevelDbStorage) WithPrefix(prefix []byte) db.Storage {
+	return &LevelDbStorage{l.ldb, db.Concat(l.prefix, prefix)}
+}
+
+// NewTx implements the method NewTx of the interface db.Storage
+func (l *LevelDbStorage) NewTx() (db.Tx, error) {
+	return &LevelDbStorageTx{l, make(db.KvMap)}, nil
+}
+
+// Get retreives a value from a key in the db.Storage
+func (l *LevelDbStorage) Get(key []byte) ([]byte, error) {
+	v, err := l.ldb.Get(db.Concat(l.prefix, key[:]), nil)
+	if err == errors.ErrNotFound {
+		return nil, db.ErrNotFound
+	}
+	return v, err
+}
+
+// Iterate implements the method Iterate of the interface db.Storage
+func (l *LevelDbStorage) Iterate(f func([]byte, []byte) (bool, error)) error {
+	// FIXME: Use the prefix!
+	snapshot, err := l.ldb.GetSnapshot()
+	if err != nil {
+		return err
+	}
+	iter := snapshot.NewIterator(util.BytesPrefix(l.prefix), nil)
+	defer iter.Release()
+	for iter.Next() {
+		localKey := iter.Key()[len(l.prefix):]
+		if cont, err := f(localKey, iter.Value()); err != nil {
+			return err
+		} else if !cont {
+			break
+		}
+	}
+	iter.Release()
+	return iter.Error()
+}
+
+// Get retreives a value from a key in the interface db.Tx
+func (l *LevelDbStorageTx) Get(key []byte) ([]byte, error) {
+	var err error
+
+	fullkey := db.Concat(l.prefix, key)
+
+	if value, ok := l.cache.Get(fullkey); ok {
+		return value, nil
+	}
+
+	value, err := l.ldb.Get(fullkey, nil)
+	if err == errors.ErrNotFound {
+		return nil, db.ErrNotFound
+	}
+
+	return value, err
+}
+
+// Insert saves a key:value into the db.Storage
+func (tx *LevelDbStorageTx) Put(k, v []byte) {
+	tx.cache.Put(db.Concat(tx.prefix, k[:]), v)
+}
+
+// Add implements the method Add of the interface db.Tx
+func (tx *LevelDbStorageTx) Add(atx db.Tx) {
+	ldbtx := atx.(*LevelDbStorageTx)
+	for _, v := range ldbtx.cache {
+		tx.cache.Put(v.K, v.V)
+	}
+}
+
+// Commit implements the method Commit of the interface db.Tx
+func (l *LevelDbStorageTx) Commit() error {
+	var batch leveldb.Batch
+	for _, v := range l.cache {
+		batch.Put(v.K, v.V)
+	}
+
+	l.cache = nil
+	return l.ldb.Write(&batch, nil)
+}
+
+// Close implements the method Close of the interface db.Tx
+func (l *LevelDbStorageTx) Close() {
+	l.cache = nil
+}
+
+// Close implements the method Close of the interface db.Storage
+func (l *LevelDbStorage) Close() {
+	if err := l.ldb.Close(); err != nil {
+		panic(err)
+	}
+	log.Info("Database closed")
+}
+
+// LevelDB is an extra method that returns the *leveldb.DB
+func (l *LevelDbStorage) LevelDB() *leveldb.DB {
+	return l.ldb
+}
+
+// List implements the method List of the interface db.Storage
+func (l *LevelDbStorage) List(limit int) ([]db.KV, error) {
+	ret := []db.KV{}
+	err := l.Iterate(func(key []byte, value []byte) (bool, error) {
+		ret = append(ret, db.KV{K: db.Clone(key), V: db.Clone(value)})
+		if len(ret) == limit {
+			return false, nil
+		}
+		return true, nil
+	})
+	return ret, err
+}

db/leveldb/leveldb_test.go

@@ -0,0 +1,57 @@
+package leveldb
+
+import (
+	"io/ioutil"
+	"os"
+	"testing"
+
+	"github.com/iden3/go-merkletree/db"
+	"github.com/iden3/go-merkletree/db/test"
+	"github.com/stretchr/testify/require"
+)
+
+var rmDirs []string
+
+func levelDbStorage(t *testing.T) db.Storage {
+	dir, err := ioutil.TempDir("", "db")
+	rmDirs = append(rmDirs, dir)
+	if err != nil {
+		t.Fatal(err)
+		return nil
+	}
+	sto, err := NewLevelDbStorage(dir, false)
+	if err != nil {
+		t.Fatal(err)
+		return nil
+	}
+	return sto
+}
+
+func TestLevelDb(t *testing.T) {
+	test.TestReturnKnownErrIfNotExists(t, levelDbStorage(t))
+	test.TestStorageInsertGet(t, levelDbStorage(t))
+	test.TestStorageWithPrefix(t, levelDbStorage(t))
+	test.TestConcatTx(t, levelDbStorage(t))
+	test.TestList(t, levelDbStorage(t))
+	test.TestIterate(t, levelDbStorage(t))
+}
+
+func TestLevelDbInterface(t *testing.T) {
+	var db db.Storage //nolint:gosimple
+
+	dir, err := ioutil.TempDir("", "db")
+	require.Nil(t, err)
+	rmDirs = append(rmDirs, dir)
+	sto, err := NewLevelDbStorage(dir, false)
+	require.Nil(t, err)
+	db = sto
+	require.NotNil(t, db)
+}
+
+func TestMain(m *testing.M) {
+	result := m.Run()
+	for _, dir := range rmDirs {
+		os.RemoveAll(dir)
+	}
+	os.Exit(result)
+}

db/memory/memory.go

@@ -0,0 +1,129 @@
+package memory
+
+import (
+	"bytes"
+	"sort"
+
+	"github.com/iden3/go-merkletree/db"
+)
+
+// MemoryStorage implements the db.Storage interface
+type MemoryStorage struct {
+	prefix []byte
+	kv     db.KvMap
+}
+
+// MemoryStorageTx implements the db.Tx interface
+type MemoryStorageTx struct {
+	s  *MemoryStorage
+	kv db.KvMap
+}
+
+// NewMemoryStorage returns a new MemoryStorage
+func NewMemoryStorage() *MemoryStorage {
+	kvmap := make(db.KvMap)
+	return &MemoryStorage{[]byte{}, kvmap}
+}
+
+// Info implements the method Info of the interface db.Storage
+func (l *MemoryStorage) Info() string {
+	return "in-memory"
+}
+
+// WithPrefix implements the method WithPrefix of the interface db.Storage
+func (m *MemoryStorage) WithPrefix(prefix []byte) db.Storage {
+	return &MemoryStorage{db.Concat(m.prefix, prefix), m.kv}
+}
+
+// NewTx implements the method NewTx of the interface db.Storage
+func (m *MemoryStorage) NewTx() (db.Tx, error) {
+	return &MemoryStorageTx{m, make(db.KvMap)}, nil
+}
+
+// Get retreives a value from a key in the db.Storage
+func (l *MemoryStorage) Get(key []byte) ([]byte, error) {
+	if v, ok := l.kv.Get(db.Concat(l.prefix, key[:])); ok {
+		return v, nil
+	}
+	return nil, db.ErrNotFound
+}
+
+// Iterate implements the method Iterate of the interface db.Storage
+func (l *MemoryStorage) Iterate(f func([]byte, []byte) (bool, error)) error {
+	kvs := make([]db.KV, 0)
+	for _, v := range l.kv {
+		if len(v.K) < len(l.prefix) || !bytes.Equal(v.K[:len(l.prefix)], l.prefix) {
+			continue
+		}
+		localkey := v.K[len(l.prefix):]
+		kvs = append(kvs, db.KV{K: localkey, V: v.V})
+
+	}
+	sort.SliceStable(kvs, func(i, j int) bool { return bytes.Compare(kvs[i].K, kvs[j].K) < 0 })
+
+	for _, kv := range kvs {
+		if cont, err := f(kv.K, kv.V); err != nil {
+			return err
+		} else if !cont {
+			break
+		}
+	}
+	return nil
+}
+
+// Get implements the method Get of the interface db.Tx
+func (tx *MemoryStorageTx) Get(key []byte) ([]byte, error) {
+
+	if v, ok := tx.kv.Get(db.Concat(tx.s.prefix, key)); ok {
+		return v, nil
+	}
+	if v, ok := tx.s.kv.Get(db.Concat(tx.s.prefix, key)); ok {
+		return v, nil
+	}
+
+	return nil, db.ErrNotFound
+}
+
+// Put implements the method Put of the interface db.Tx
+func (tx *MemoryStorageTx) Put(k, v []byte) {
+	tx.kv.Put(db.Concat(tx.s.prefix, k), v)
+}
+
+// Commit implements the method Commit of the interface db.Tx
+func (tx *MemoryStorageTx) Commit() error {
+	for _, v := range tx.kv {
+		tx.s.kv.Put(v.K, v.V)
+	}
+	tx.kv = nil
+	return nil
+}
+
+// Add implements the method Add of the interface db.Tx
+func (tx *MemoryStorageTx) Add(atx db.Tx) {
+	mstx := atx.(*MemoryStorageTx)
+	for _, v := range mstx.kv {
+		tx.kv.Put(v.K, v.V)
+	}
+}
+
+// Close implements the method Close of the interface db.Tx
+func (tx *MemoryStorageTx) Close() {
+	tx.kv = nil
+}
+
+// Close implements the method Close of the interface db.Storage
+func (m *MemoryStorage) Close() {
+}
+
+// List implements the method List of the interface db.Storage
+func (l *MemoryStorage) List(limit int) ([]db.KV, error) {
+	ret := []db.KV{}
+	err := l.Iterate(func(key []byte, value []byte) (bool, error) {
+		ret = append(ret, db.KV{K: db.Clone(key), V: db.Clone(value)})
+		if len(ret) == limit {
+			return false, nil
+		}
+		return true, nil
+	})
+	return ret, err
+}

db/memory/memory_test.go

@@ -0,0 +1,25 @@
+package memory
+
+import (
+	"testing"
+
+	"github.com/iden3/go-merkletree/db"
+	"github.com/iden3/go-merkletree/db/test"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMemoryStorageInterface(t *testing.T) {
+	var db db.Storage //nolint:gosimple
+
+	db = NewMemoryStorage()
+	require.NotNil(t, db)
+}
+
+func TestMemory(t *testing.T) {
+	test.TestReturnKnownErrIfNotExists(t, NewMemoryStorage())
+	test.TestStorageInsertGet(t, NewMemoryStorage())
+	test.TestStorageWithPrefix(t, NewMemoryStorage())
+	test.TestConcatTx(t, NewMemoryStorage())
+	test.TestList(t, NewMemoryStorage())
+	test.TestIterate(t, NewMemoryStorage())
+}

db/test/test.go

@@ -0,0 +1,193 @@
+package test
+
+import (
+	"testing"
+
+	"github.com/iden3/go-merkletree/db"
+	"github.com/stretchr/testify/assert"
+)
+
+// TestReturnKnownErrIfNotExists checks that the implementation of the
+// db.Storage interface returns the expected error in the case that the value
+// is not found
+func TestReturnKnownErrIfNotExists(t *testing.T, sto db.Storage) {
+	k := []byte("key")
+
+	tx, err := sto.NewTx()
+	assert.Nil(t, err)
+	_, err = tx.Get(k)
+	assert.EqualError(t, err, db.ErrNotFound.Error())
+}
+
+// TestStorageInsertGet checks that the implementation of the db.Storage
+// interface behaves as expected
+func TestStorageInsertGet(t *testing.T, sto db.Storage) {
+	key := []byte("key")
+	value := []byte("data")
+
+	tx, err := sto.NewTx()
+	assert.Nil(t, err)
+	tx.Put(key, value)
+	v, err := tx.Get(key)
+	assert.Nil(t, err)
+	assert.Equal(t, value, v)
+	assert.Nil(t, tx.Commit())
+
+	tx, err = sto.NewTx()
+	assert.Nil(t, err)
+	v, err = tx.Get(key)
+	assert.Nil(t, err)
+	assert.Equal(t, value, v)
+}
+
+// TestStorageWithPrefix checks that the implementation of the db.Storage
+// interface behaves as expected for the WithPrefix method
+func TestStorageWithPrefix(t *testing.T, sto db.Storage) {
+	k := []byte{9}
+
+	sto1 := sto.WithPrefix([]byte{1})
+	sto2 := sto.WithPrefix([]byte{2})
+
+	// check within tx
+
+	sto1tx, err := sto1.NewTx()
+	assert.Nil(t, err)
+	sto1tx.Put(k, []byte{4, 5, 6})
+	v1, err := sto1tx.Get(k)
+	assert.Nil(t, err)
+	assert.Equal(t, v1, []byte{4, 5, 6})
+	assert.Nil(t, sto1tx.Commit())
+
+	sto2tx, err := sto2.NewTx()
+	assert.Nil(t, err)
+	sto2tx.Put(k, []byte{8, 9})
+	v2, err := sto2tx.Get(k)
+	assert.Nil(t, err)
+	assert.Equal(t, v2, []byte{8, 9})
+	assert.Nil(t, sto2tx.Commit())
+
+	// check outside tx
+
+	v1, err = sto1.Get(k)
+	assert.Nil(t, err)
+	assert.Equal(t, v1, []byte{4, 5, 6})
+
+	v2, err = sto2.Get(k)
+	assert.Nil(t, err)
+	assert.Equal(t, v2, []byte{8, 9})
+}
+
+// TestIterate checks that the implementation of the db.Storage interface
+// behaves as expected for the Iterate method
+func TestIterate(t *testing.T, sto db.Storage) {
+	r := []db.KV{}
+	lister := func(k []byte, v []byte) (bool, error) {
+		r = append(r, db.KV{K: db.Clone(k), V: db.Clone(v)})
+		return true, nil
+	}
+
+	sto1 := sto.WithPrefix([]byte{1})
+	err := sto1.Iterate(lister)
+	assert.Nil(t, err)
+	assert.Equal(t, 0, len(r))
+
+	sto1tx, _ := sto1.NewTx()
+	sto1tx.Put([]byte{1}, []byte{4})
+	sto1tx.Put([]byte{2}, []byte{5})
+	sto1tx.Put([]byte{3}, []byte{6})
+	assert.Nil(t, sto1tx.Commit())
+
+	sto2 := sto.WithPrefix([]byte{2})
+	sto2tx, _ := sto2.NewTx()
+	sto2tx.Put([]byte{1}, []byte{7})
+	sto2tx.Put([]byte{2}, []byte{8})
+	sto2tx.Put([]byte{3}, []byte{9})
+	assert.Nil(t, sto2tx.Commit())
+
+	r = []db.KV{}
+	err = sto1.Iterate(lister)
+	assert.Nil(t, err)
+	assert.Equal(t, 3, len(r))
+	assert.Equal(t, db.KV{[]byte{1}, []byte{4}}, r[0])
+	assert.Equal(t, db.KV{[]byte{2}, []byte{5}}, r[1])
+	assert.Equal(t, db.KV{[]byte{3}, []byte{6}}, r[2])
+
+	r = []db.KV{}
+	err = sto2.Iterate(lister)
+	assert.Nil(t, err)
+	assert.Equal(t, 3, len(r))
+	assert.Equal(t, db.KV{[]byte{1}, []byte{7}}, r[0])
+	assert.Equal(t, db.KV{[]byte{2}, []byte{8}}, r[1])
+	assert.Equal(t, db.KV{[]byte{3}, []byte{9}}, r[2])
+}
+
+// TestConcatTx checks that the implementation of the db.Storage interface
+// behaves as expected
+func TestConcatTx(t *testing.T, sto db.Storage) {
+	k := []byte{9}
+
+	sto1 := sto.WithPrefix([]byte{1})
+	sto2 := sto.WithPrefix([]byte{2})
+
+	// check within tx
+
+	sto1tx, err := sto1.NewTx()
+	if err != nil {
+		panic(err)
+	}
+	sto1tx.Put(k, []byte{4, 5, 6})
+	sto2tx, err := sto2.NewTx()
+	if err != nil {
+		panic(err)
+	}
+	sto2tx.Put(k, []byte{8, 9})
+
+	sto1tx.Add(sto2tx)
+	assert.Nil(t, sto1tx.Commit())
+
+	// check outside tx
+
+	v1, err := sto1.Get(k)
+	assert.Nil(t, err)
+	assert.Equal(t, v1, []byte{4, 5, 6})
+
+	v2, err := sto2.Get(k)
+	assert.Nil(t, err)
+	assert.Equal(t, v2, []byte{8, 9})
+}
+
+// TestList checks that the implementation of the db.Storage interface behaves
+// as expected
+func TestList(t *testing.T, sto db.Storage) {
+	sto1 := sto.WithPrefix([]byte{1})
+	r1, err := sto1.List(100)
+	assert.Nil(t, err)
+	assert.Equal(t, 0, len(r1))
+
+	sto1tx, _ := sto1.NewTx()
+	sto1tx.Put([]byte{1}, []byte{4})
+	sto1tx.Put([]byte{2}, []byte{5})
+	sto1tx.Put([]byte{3}, []byte{6})
+	assert.Nil(t, sto1tx.Commit())
+
+	sto2 := sto.WithPrefix([]byte{2})
+	sto2tx, _ := sto2.NewTx()
+	sto2tx.Put([]byte{1}, []byte{7})
+	sto2tx.Put([]byte{2}, []byte{8})
+	sto2tx.Put([]byte{3}, []byte{9})
+	assert.Nil(t, sto2tx.Commit())
+
+	r, err := sto1.List(100)
+	assert.Nil(t, err)
+	assert.Equal(t, 3, len(r))
+	assert.Equal(t, r[0], db.KV{[]byte{1}, []byte{4}})
+	assert.Equal(t, r[1], db.KV{[]byte{2}, []byte{5}})
+	assert.Equal(t, r[2], db.KV{[]byte{3}, []byte{6}})
+
+	r, err = sto1.List(2)
+	assert.Nil(t, err)
+	assert.Equal(t, 2, len(r))
+	assert.Equal(t, r[0], db.KV{[]byte{1}, []byte{4}})
+	assert.Equal(t, r[1], db.KV{[]byte{2}, []byte{5}})
+
+}