Merge pull request #97 from bmatsuo/bmatsuo/fix-cursor-op-set-with-by…

…tes-buffer-keys

Cursor.Get handles key return values specicially for Set

lmdb/cursor.go

@@ -135,6 +135,11 @@ func (c *Cursor) DBI() DBI {
 // returned by Get reference readonly sections of memory that must not be
 // accessed after the transaction has terminated.
 //
+// In a Txn with RawRead set to true the Set op causes the returned key to
+// share its memory with setkey (making it writable memory). In a Txn with
+// RawRead set to false the Set op returns key values with memory distinct from
+// setkey, as is always the case when using RawRead.
+//
 // Get ignores setval if setkey is empty.
 //
 // See mdb_cursor_get.
@@ -152,11 +157,30 @@ func (c *Cursor) Get(setkey, setval []byte, op uint) (key, val []byte, err error
 		*c.txn.val = C.MDB_val{}
 		return nil, nil, err
 	}
-	k := c.txn.bytes(c.txn.key)
-	v := c.txn.bytes(c.txn.val)
+
+	// When MDB_SET is passed to mdb_cursor_get its first argument will be
+	// returned unchanged.  Unfortunately, the normal slice copy/extraction
+	// routines will be bad for the Go runtime when operating on Go memory
+	// (panic or potentially garbage memory reference).
+	if op == Set {
+		if c.txn.RawRead {
+			key = setkey
+		} else {
+			p := make([]byte, len(setkey))
+			copy(p, setkey)
+			key = p
+		}
+	} else {
+		key = c.txn.bytes(c.txn.key)
+	}
+	val = c.txn.bytes(c.txn.val)
+
+	// Clear transaction storage record storage area for future use and to
+	// prevent dangling references.
 	*c.txn.key = C.MDB_val{}
 	*c.txn.val = C.MDB_val{}
-	return k, v, nil
+
+	return key, val, nil
 }
 
 // getVal0 retrieves items from the database without using given key or value

lmdb/cursor_test.go

@@ -250,7 +250,7 @@ func TestCursor_Get_KV(t *testing.T) {
 			t.Errorf("unexpected key: %q (not %q)", k, "key")
 		}
 		if string(v) != "1" {
-			t.Errorf("unexpected key: %q (not %q)", k, "1")
+			t.Errorf("unexpected value: %q (not %q)", k, "1")
 		}
 
 		_, _, err = cur.Get([]byte("key"), []byte("1"), GetBoth)
@@ -261,6 +261,82 @@ func TestCursor_Get_KV(t *testing.T) {
 	}
 }
 
+func TestCursor_Get_op_Set_bytesBuffer(t *testing.T) {
+	env := setup(t)
+	defer clean(env, t)
+
+	var dbi DBI
+	err := env.Update(func(txn *Txn) (err error) {
+		dbi, err = txn.OpenDBI("testdb", Create|DupSort)
+		return err
+	})
+	if err != nil {
+		t.Errorf("%s", err)
+		return
+	}
+
+	err = env.Update(func(txn *Txn) (err error) {
+		put := func(k, v []byte) {
+			if err == nil {
+				err = txn.Put(dbi, k, v, 0)
+			}
+		}
+		put([]byte("k1"), []byte("v11"))
+		put([]byte("k1"), []byte("v12"))
+		put([]byte("k1"), []byte("v13"))
+		put([]byte("k2"), []byte("v21"))
+		put([]byte("k2"), []byte("v22"))
+		put([]byte("k2"), []byte("v23"))
+		return err
+	})
+	if err != nil {
+		t.Errorf("%s", err)
+	}
+
+	err = env.View(func(txn *Txn) (err error) {
+		cur, err := txn.OpenCursor(dbi)
+		if err != nil {
+			return err
+		}
+		defer cur.Close()
+
+		// Create bytes.Buffer values containing a amount of bytes.  Byte
+		// slices returned from buf.Bytes() have a history of tricking the cgo
+		// argument checker.
+		var kbuf bytes.Buffer
+		kbuf.WriteString("k2")
+
+		k, _, err := cur.Get(kbuf.Bytes(), nil, Set)
+		if err != nil {
+			return err
+		}
+		if string(k) != kbuf.String() {
+			t.Errorf("unexpected key: %q (not %q)", k, kbuf.String())
+		}
+
+		// No guarantee is made about the return value of mdb_cursor_get when
+		// MDB_SET is the op, so its value is not checked as part of this test.
+		// That said, it is important that Cursor.Get not panic if given a
+		// short buffer as an input value for a Set op (despite that not really
+		// having any significance)
+		var vbuf bytes.Buffer
+		vbuf.WriteString("v22")
+
+		k, _, err = cur.Get(kbuf.Bytes(), vbuf.Bytes(), Set)
+		if err != nil {
+			return err
+		}
+		if string(k) != kbuf.String() {
+			t.Errorf("unexpected key: %q (not %q)", k, kbuf.String())
+		}
+
+		return nil
+	})
+	if err != nil {
+		t.Errorf("%s", err)
+	}
+}
+
 func TestCursor_Get_DupFixed(t *testing.T) {
 	env := setup(t)
 	defer clean(env, t)