From 5b051d9ab39674f35f4e9c0ae2dae837edd12fc6 Mon Sep 17 00:00:00 2001
From: Alex Sharov <AskAlexSharov@gmail.com>
Date: Wed, 17 Jun 2020 16:57:48 +0700
Subject: [PATCH] Lmdb: avoid empty key values in put and seek (#673)

* avoid 0 length key/values and seek arguments

* less abstractions in .get method, and copy memory from C space without unsafe conversion into go memory

* fix test
---
 eth/stagedsync/stage_execute.go |  3 +-
 ethdb/database_test.go          |  6 ++--
 ethdb/kv_lmdb.go                | 55 +++++++++++++++------------------
 3 files changed, 30 insertions(+), 34 deletions(-)

diff --git a/eth/stagedsync/stage_execute.go b/eth/stagedsync/stage_execute.go
index 529f6378a..b408589d6 100644
--- a/eth/stagedsync/stage_execute.go
+++ b/eth/stagedsync/stage_execute.go
@@ -169,10 +169,11 @@ func SpawnExecuteBlocksStage(s *StageState, stateDB ethdb.Database, blockchain B
 				return err
 			}
 			start := time.Now()
+			sz := batch.BatchSize()
 			if _, err = batch.Commit(); err != nil {
 				return err
 			}
-			log.Info("Batch committed", "in", time.Since(start), "size", common.StorageSize(batch.BatchSize()))
+			log.Info("Batch committed", "in", time.Since(start), "size", common.StorageSize(sz))
 		}
 
 		if prof {
diff --git a/ethdb/database_test.go b/ethdb/database_test.go
index 0fb2b56b9..ba1d94068 100644
--- a/ethdb/database_test.go
+++ b/ethdb/database_test.go
@@ -95,7 +95,7 @@ func testPutGet(db MinDatabase, t *testing.T) {
 	t.Parallel()
 
 	for _, k := range testValues {
-		err := db.Put(testBucket, []byte(k), []byte{})
+		err := db.Put(testBucket, []byte(k), []byte{1})
 		if err != nil {
 			t.Fatalf("put failed: %v", err)
 		}
@@ -106,8 +106,8 @@ func testPutGet(db MinDatabase, t *testing.T) {
 		if err != nil {
 			t.Fatalf("get failed: %v", err)
 		}
-		if len(data) != 0 {
-			t.Fatalf("get returned wrong result, got %q expected nil", string(data))
+		if len(data) != 1 {
+			t.Fatalf("get returned wrong result, got %q expected 1", string(data))
 		}
 	}
 
diff --git a/ethdb/kv_lmdb.go b/ethdb/kv_lmdb.go
index f51ab575f..b1d17219d 100644
--- a/ethdb/kv_lmdb.go
+++ b/ethdb/kv_lmdb.go
@@ -422,6 +422,13 @@ func (b *lmdbBucket) Put(key []byte, value []byte) error {
 	default:
 	}
 
+	if len(key) == 0 {
+		return fmt.Errorf("lmdb doesn't support empty keys. bucket: %s", dbutils.Buckets[b.id])
+	}
+	if len(value) == 0 {
+		return fmt.Errorf("lmdb doesn't support empty values. bucket: %s", dbutils.Buckets[b.id])
+	}
+
 	err := b.tx.tx.Put(b.dbi, key, value, 0)
 	if err != nil {
 		return fmt.Errorf("failed LmdbKV.Put: %w", err)
@@ -517,7 +524,7 @@ func (c *lmdbCursor) Seek(seek []byte) (k, v []byte, err error) {
 		}
 	}
 
-	if seek == nil {
+	if len(seek) == 0 {
 		k, v, err = c.cursor.Get(nil, nil, lmdb.First)
 	} else {
 		k, v, err = c.cursor.Get(seek, nil, lmdb.SetRange)
@@ -591,6 +598,12 @@ func (c *lmdbCursor) Put(key []byte, value []byte) error {
 	default:
 	}
 
+	if len(key) == 0 {
+		return fmt.Errorf("lmdb doesn't support empty keys. bucket: %s", dbutils.Buckets[c.bucket.id])
+	}
+	if len(value) == 0 {
+		return fmt.Errorf("lmdb doesn't support empty values. bucket: %s", dbutils.Buckets[c.bucket.id])
+	}
 	if c.cursor == nil {
 		if err := c.initCursor(); err != nil {
 			return err
@@ -637,51 +650,33 @@ func (c *lmdbNoValuesCursor) Walk(walker func(k []byte, vSize uint32) (bool, err
 }
 
 func (c *lmdbNoValuesCursor) First() (k []byte, v uint32, err error) {
-	if c.cursor == nil {
-		if err := c.initCursor(); err != nil {
-			return []byte{}, 0, err
-		}
-	}
-
-	var val []byte
-	if len(c.prefix) == 0 {
-		k, val, err = c.cursor.Get(nil, nil, lmdb.First)
-	} else {
-		k, val, err = c.cursor.Get(c.prefix, nil, lmdb.SetKey)
-	}
-	if err != nil {
-		if lmdb.IsNotFound(err) {
-			return []byte{}, uint32(len(val)), nil
-		}
-		return []byte{}, 0, err
-	}
-
-	if c.prefix != nil && !bytes.HasPrefix(k, c.prefix) {
-		k, val = nil, nil
-	}
-	return k, uint32(len(val)), err
+	return c.Seek(c.prefix)
 }
 
-func (c *lmdbNoValuesCursor) Seek(seek []byte) (k []byte, v uint32, err error) {
+func (c *lmdbNoValuesCursor) Seek(seek []byte) (k []byte, vSize uint32, err error) {
 	if c.cursor == nil {
 		if err := c.initCursor(); err != nil {
 			return []byte{}, 0, err
 		}
 	}
 
-	var val []byte
-	k, val, err = c.cursor.Get(seek, nil, lmdb.SetKey)
+	var v []byte
+	if len(seek) == 0 {
+		k, v, err = c.cursor.Get(nil, nil, lmdb.First)
+	} else {
+		k, v, err = c.cursor.Get(seek, nil, lmdb.SetRange)
+	}
 	if err != nil {
 		if lmdb.IsNotFound(err) {
-			return []byte{}, uint32(len(val)), nil
+			return []byte{}, uint32(len(v)), nil
 		}
 		return []byte{}, 0, err
 	}
 	if c.prefix != nil && !bytes.HasPrefix(k, c.prefix) {
-		k, val = nil, nil
+		k, v = nil, nil
 	}
 
-	return k, uint32(len(val)), err
+	return k, uint32(len(v)), err
 }
 
 func (c *lmdbNoValuesCursor) SeekTo(seek []byte) ([]byte, uint32, error) {