neoneo-go/pkg/core/storage/boltdb_store.go
Roman Khimov ac5d2f94d3 storage: fix BoltDB batched Put()
It must copy both the value and the key because they can be reused for other
purposes between Put() and PutBatch(). This actually happens with values in
headers processing, leading to wrong data being written into the DB.

Extend the batch test to check for that.
2019-09-23 09:27:18 +03:00

133 lines
3.1 KiB
Go

package storage
import (
"bytes"
"fmt"
"os"
"path"
"github.com/etcd-io/bbolt"
"github.com/syndtr/goleveldb/leveldb/util"
)
// BoltDBOptions configuration for boltdb.
type BoltDBOptions struct {
FilePath string `yaml:"FilePath"`
}
// Bucket represents bucket used in boltdb to store all the data.
var Bucket = []byte("DB")
// BoltDBStore it is the storage implementation for storing and retrieving
// blockchain data.
type BoltDBStore struct {
db *bbolt.DB
}
// BoltDBBatch simple batch implementation to satisfy the Store interface.
type BoltDBBatch struct {
mem map[*[]byte][]byte
}
// Len implements the Batch interface.
func (b *BoltDBBatch) Len() int {
return len(b.mem)
}
// Put implements the Batch interface.
func (b *BoltDBBatch) Put(k, v []byte) {
vcopy := make([]byte, len(v))
copy(vcopy, v)
kcopy := make([]byte, len(k))
copy(kcopy, k)
b.mem[&kcopy] = vcopy
}
// NewBoltDBStore returns a new ready to use BoltDB storage with created bucket.
func NewBoltDBStore(cfg BoltDBOptions) (*BoltDBStore, error) {
var opts *bbolt.Options // should be exposed via BoltDBOptions if anything needed
fileMode := os.FileMode(0600) // should be exposed via BoltDBOptions if anything needed
fileName := cfg.FilePath
dir := path.Dir(fileName)
err := os.MkdirAll(dir, os.ModePerm)
if err != nil {
return nil, fmt.Errorf("could not create dir for BoltDB: %v", err)
}
db, err := bbolt.Open(fileName, fileMode, opts)
if err != nil {
return nil, err
}
err = db.Update(func(tx *bbolt.Tx) error {
_, err = tx.CreateBucketIfNotExists(Bucket)
if err != nil {
return fmt.Errorf("could not create root bucket: %v", err)
}
return nil
})
return &BoltDBStore{db: db}, nil
}
// Put implements the Store interface.
func (s *BoltDBStore) Put(key, value []byte) error {
return s.db.Update(func(tx *bbolt.Tx) error {
b := tx.Bucket(Bucket)
err := b.Put(key, value)
return err
})
}
// Get implements the Store interface.
func (s *BoltDBStore) Get(key []byte) (val []byte, err error) {
err = s.db.View(func(tx *bbolt.Tx) error {
b := tx.Bucket(Bucket)
val = b.Get(key)
return nil
})
if val == nil {
err = ErrKeyNotFound
}
return
}
// PutBatch implements the Store interface.
func (s *BoltDBStore) PutBatch(batch Batch) error {
return s.db.Batch(func(tx *bbolt.Tx) error {
b := tx.Bucket(Bucket)
for k, v := range batch.(*BoltDBBatch).mem {
err := b.Put(*k, v)
if err != nil {
return err
}
}
return nil
})
}
// Seek implements the Store interface.
func (s *BoltDBStore) Seek(key []byte, f func(k, v []byte)) {
err := s.db.View(func(tx *bbolt.Tx) error {
c := tx.Bucket(Bucket).Cursor()
prefix := util.BytesPrefix(key)
for k, v := c.Seek(prefix.Start); k != nil && bytes.Compare(k, prefix.Limit) <= 0; k, v = c.Next() {
f(k, v)
}
return nil
})
if err != nil {
fmt.Println("error while executing seek in boltDB")
}
}
// Batch implements the Batch interface and returns a boltdb
// compatible Batch.
func (s *BoltDBStore) Batch() Batch {
return &BoltDBBatch{
mem: make(map[*[]byte][]byte),
}
}
// Close releases all db resources.
func (s *BoltDBStore) Close() error {
return s.db.Close()
}