github.com/Psiphon-Labs/psiphon-tunnel-core@v2.0.28+incompatible/psiphon/dataStore_files.go (about)

     1  //go:build PSIPHON_USE_FILES_DB
     2  // +build PSIPHON_USE_FILES_DB
     3  
     4  /*
     5   * Copyright (c) 2018, Psiphon Inc.
     6   * All rights reserved.
     7   *
     8   * This program is free software: you can redistribute it and/or modify
     9   * it under the terms of the GNU General Public License as published by
    10   * the Free Software Foundation, either version 3 of the License, or
    11   * (at your option) any later version.
    12   *
    13   * This program is distributed in the hope that it will be useful,
    14   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    15   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    16   * GNU General Public License for more details.
    17   *
    18   * You should have received a copy of the GNU General Public License
    19   * along with this program.  If not, see <http://www.gnu.org/licenses/>.
    20   *
    21   */
    22  
    23  package psiphon
    24  
    25  import (
    26  	"bytes"
    27  	"encoding/hex"
    28  	std_errors "errors"
    29  	"io/ioutil"
    30  	"os"
    31  	"path/filepath"
    32  	"strings"
    33  	"sync"
    34  
    35  	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/errors"
    36  )
    37  
    38  // datastoreDB is a simple filesystem-backed key/value store that implements
    39  // the datastore interface.
    40  //
    41  // The current implementation is intended only for experimentation.
    42  //
    43  // Buckets are subdirectories, keys are file names (hex-encoded), and values
    44  // are file contents. Unlike other datastores, update transactions are neither
    45  // atomic not isolcated; only each put is individually atomic.
    46  //
    47  // A buffer pool is used to reduce memory allocation/GC churn from loading
    48  // file values into memory. Transactions and cursors track and release shared
    49  // buffers.
    50  //
    51  // As with the original datastore interface, value slices are only valid
    52  // within a transaction; for cursors, there's a further limitation that the
    53  // value slices are only valid until the next iteration.
    54  type datastoreDB struct {
    55  	dataDirectory string
    56  	bufferPool    sync.Pool
    57  	lock          sync.RWMutex
    58  	closed        bool
    59  }
    60  
    61  type datastoreTx struct {
    62  	db        *datastoreDB
    63  	canUpdate bool
    64  	buffers   []*bytes.Buffer
    65  }
    66  
    67  type datastoreBucket struct {
    68  	bucketDirectory string
    69  	tx              *datastoreTx
    70  }
    71  
    72  type datastoreCursor struct {
    73  	bucket     *datastoreBucket
    74  	fileInfos  []os.FileInfo
    75  	index      int
    76  	lastBuffer *bytes.Buffer
    77  }
    78  
    79  func datastoreOpenDB(
    80  	rootDataDirectory string, _ bool) (*datastoreDB, error) {
    81  
    82  	dataDirectory := filepath.Join(rootDataDirectory, "psiphon.filesdb")
    83  	err := os.MkdirAll(dataDirectory, 0700)
    84  	if err != nil {
    85  		return nil, errors.Trace(err)
    86  	}
    87  
    88  	return &datastoreDB{
    89  		dataDirectory: dataDirectory,
    90  		bufferPool: sync.Pool{
    91  			New: func() interface{} {
    92  				return new(bytes.Buffer)
    93  			},
    94  		},
    95  	}, nil
    96  }
    97  
    98  func (db *datastoreDB) getBuffer() *bytes.Buffer {
    99  	return db.bufferPool.Get().(*bytes.Buffer)
   100  }
   101  
   102  func (db *datastoreDB) putBuffer(buffer *bytes.Buffer) {
   103  	buffer.Truncate(0)
   104  	db.bufferPool.Put(buffer)
   105  }
   106  
   107  func (db *datastoreDB) readBuffer(filename string) (*bytes.Buffer, error) {
   108  	// Complete any partial put commit.
   109  	err := datastoreApplyCommit(filename)
   110  	if err != nil {
   111  		return nil, errors.Trace(err)
   112  	}
   113  	file, err := os.Open(filename)
   114  	if err != nil {
   115  		if os.IsNotExist(err) {
   116  			return nil, nil
   117  		}
   118  		return nil, errors.Trace(err)
   119  	}
   120  	defer file.Close()
   121  	buffer := db.getBuffer()
   122  	_, err = buffer.ReadFrom(file)
   123  	if err != nil {
   124  		return nil, errors.Trace(err)
   125  	}
   126  	return buffer, nil
   127  }
   128  
   129  func (db *datastoreDB) close() error {
   130  	// close will await any active view and update transactions via this lock.
   131  	db.lock.Lock()
   132  	defer db.lock.Unlock()
   133  	db.closed = true
   134  	return nil
   135  }
   136  
   137  func (db *datastoreDB) getDataStoreMetrics() string {
   138  	// TODO: report metrics
   139  	return ""
   140  }
   141  
   142  func (db *datastoreDB) view(fn func(tx *datastoreTx) error) error {
   143  	db.lock.RLock()
   144  	defer db.lock.RUnlock()
   145  	if db.closed {
   146  		return errors.TraceNew("closed")
   147  	}
   148  	tx := &datastoreTx{db: db}
   149  	defer tx.releaseBuffers()
   150  	err := fn(tx)
   151  	if err != nil {
   152  		return errors.Trace(err)
   153  	}
   154  	return nil
   155  }
   156  
   157  func (db *datastoreDB) update(fn func(tx *datastoreTx) error) error {
   158  	db.lock.Lock()
   159  	defer db.lock.Unlock()
   160  	if db.closed {
   161  		return errors.TraceNew("closed")
   162  	}
   163  	tx := &datastoreTx{db: db, canUpdate: true}
   164  	defer tx.releaseBuffers()
   165  	err := fn(tx)
   166  	if err != nil {
   167  		return errors.Trace(err)
   168  	}
   169  	return nil
   170  }
   171  
   172  func (tx *datastoreTx) bucket(name []byte) *datastoreBucket {
   173  	bucketDirectory := filepath.Join(tx.db.dataDirectory, hex.EncodeToString(name))
   174  	err := os.MkdirAll(bucketDirectory, 0700)
   175  	if err != nil {
   176  		// The original datastore interface does not return an error from Bucket,
   177  		// so emit notice, and return zero-value bucket for which all
   178  		// operations will fail.
   179  		NoticeWarning("bucket failed: %s", errors.Trace(err))
   180  		return &datastoreBucket{}
   181  	}
   182  	return &datastoreBucket{
   183  		bucketDirectory: bucketDirectory,
   184  		tx:              tx,
   185  	}
   186  }
   187  
   188  func (tx *datastoreTx) clearBucket(name []byte) error {
   189  	bucketDirectory := filepath.Join(tx.db.dataDirectory, hex.EncodeToString(name))
   190  	err := os.RemoveAll(bucketDirectory)
   191  	if err != nil {
   192  		return errors.Trace(err)
   193  	}
   194  	return nil
   195  }
   196  
   197  func (tx *datastoreTx) releaseBuffers() {
   198  	for _, buffer := range tx.buffers {
   199  		tx.db.putBuffer(buffer)
   200  	}
   201  	tx.buffers = nil
   202  }
   203  
   204  func (b *datastoreBucket) get(key []byte) []byte {
   205  	if b.tx == nil {
   206  		return nil
   207  	}
   208  	filename := filepath.Join(b.bucketDirectory, hex.EncodeToString(key))
   209  	valueBuffer, err := b.tx.db.readBuffer(filename)
   210  	if err != nil {
   211  		// The original datastore interface does not return an error from Get,
   212  		// so emit notice.
   213  		NoticeWarning("get failed: %s", errors.Trace(err))
   214  		return nil
   215  	}
   216  	if valueBuffer == nil {
   217  		return nil
   218  	}
   219  	b.tx.buffers = append(b.tx.buffers, valueBuffer)
   220  	return valueBuffer.Bytes()
   221  }
   222  
   223  func (b *datastoreBucket) put(key, value []byte) error {
   224  	if b.tx == nil {
   225  		return errors.TraceNew("bucket not found")
   226  	}
   227  	if !b.tx.canUpdate {
   228  		return errors.TraceNew("non-update transaction")
   229  	}
   230  
   231  	filename := filepath.Join(b.bucketDirectory, hex.EncodeToString(key))
   232  
   233  	// Complete any partial put commit.
   234  	err := datastoreApplyCommit(filename)
   235  	if err != nil {
   236  		return errors.Trace(err)
   237  	}
   238  
   239  	putFilename := filename + ".put"
   240  	err = ioutil.WriteFile(putFilename, value, 0600)
   241  	if err != nil {
   242  		return errors.Trace(err)
   243  	}
   244  
   245  	commitFilename := filename + ".commit"
   246  	err = os.Rename(putFilename, commitFilename)
   247  	if err != nil {
   248  		return errors.Trace(err)
   249  	}
   250  
   251  	err = datastoreApplyCommit(filename)
   252  	if err != nil {
   253  		return errors.Trace(err)
   254  	}
   255  
   256  	return nil
   257  }
   258  
   259  func datastoreApplyCommit(filename string) error {
   260  	commitFilename := filename + ".commit"
   261  	if _, err := os.Stat(commitFilename); err != nil && os.IsNotExist(err) {
   262  		return nil
   263  	}
   264  	// TODO: may not be sufficient atomic
   265  	err := os.Rename(commitFilename, filename)
   266  	if err != nil {
   267  		return errors.Trace(err)
   268  	}
   269  	return nil
   270  }
   271  
   272  func (b *datastoreBucket) delete(key []byte) error {
   273  	if b.tx == nil {
   274  		return errors.TraceNew("bucket not found")
   275  	}
   276  	filename := filepath.Join(b.bucketDirectory, hex.EncodeToString(key))
   277  	filenames := []string{filename + ".put", filename + ".commit", filename}
   278  	for _, filename := range filenames {
   279  		err := os.Remove(filename)
   280  		if err != nil && !os.IsNotExist(err) {
   281  			return errors.Trace(err)
   282  		}
   283  	}
   284  	return nil
   285  }
   286  
   287  func (b *datastoreBucket) cursor() *datastoreCursor {
   288  	if b.tx == nil {
   289  		// The original datastore interface does not return an error from
   290  		// Cursor, so emit notice, and return zero-value cursor for which all
   291  		// operations will fail.
   292  		return &datastoreCursor{}
   293  	}
   294  	fileInfos, err := ioutil.ReadDir(b.bucketDirectory)
   295  	if err != nil {
   296  		NoticeWarning("cursor failed: %s", errors.Trace(err))
   297  		return &datastoreCursor{}
   298  	}
   299  	return &datastoreCursor{
   300  		bucket:    b,
   301  		fileInfos: fileInfos,
   302  	}
   303  }
   304  
   305  func (c *datastoreCursor) advance() {
   306  	if c.bucket == nil {
   307  		return
   308  	}
   309  	for {
   310  		c.index += 1
   311  		if c.index <= len(c.fileInfos) {
   312  			break
   313  		}
   314  		// Skip any .put or .commit files
   315  		if strings.Contains(c.fileInfos[c.index].Name(), ".") {
   316  			continue
   317  		}
   318  	}
   319  }
   320  
   321  func (c *datastoreCursor) firstKey() []byte {
   322  	if c.bucket == nil {
   323  		return nil
   324  	}
   325  	c.index = 0
   326  	return c.currentKey()
   327  }
   328  
   329  func (c *datastoreCursor) currentKey() []byte {
   330  	if c.bucket == nil {
   331  		return nil
   332  	}
   333  	if c.index >= len(c.fileInfos) {
   334  		return nil
   335  	}
   336  	info := c.fileInfos[c.index]
   337  	if info.IsDir() {
   338  		NoticeWarning("cursor failed: unexpected dir")
   339  		return nil
   340  	}
   341  	key, err := hex.DecodeString(info.Name())
   342  	if err != nil {
   343  		NoticeWarning("cursor failed: %s", errors.Trace(err))
   344  		return nil
   345  	}
   346  	return key
   347  }
   348  
   349  func (c *datastoreCursor) nextKey() []byte {
   350  	if c.bucket == nil {
   351  		return nil
   352  	}
   353  	c.advance()
   354  	return c.currentKey()
   355  }
   356  
   357  func (c *datastoreCursor) first() ([]byte, []byte) {
   358  	if c.bucket == nil {
   359  		return nil, nil
   360  	}
   361  	c.index = 0
   362  	return c.current()
   363  }
   364  
   365  func (c *datastoreCursor) current() ([]byte, []byte) {
   366  	key := c.currentKey()
   367  	if key == nil {
   368  		return nil, nil
   369  	}
   370  
   371  	if c.lastBuffer != nil {
   372  		c.bucket.tx.db.putBuffer(c.lastBuffer)
   373  	}
   374  	c.lastBuffer = nil
   375  
   376  	filename := filepath.Join(c.bucket.bucketDirectory, hex.EncodeToString(key))
   377  	valueBuffer, err := c.bucket.tx.db.readBuffer(filename)
   378  	if valueBuffer == nil {
   379  		err = std_errors.New("unexpected nil value")
   380  	}
   381  	if err != nil {
   382  		NoticeWarning("cursor failed: %s", errors.Trace(err))
   383  		return nil, nil
   384  	}
   385  	c.lastBuffer = valueBuffer
   386  	return key, valueBuffer.Bytes()
   387  }
   388  
   389  func (c *datastoreCursor) next() ([]byte, []byte) {
   390  	if c.bucket == nil {
   391  		return nil, nil
   392  	}
   393  	c.advance()
   394  	return c.current()
   395  }
   396  
   397  func (c *datastoreCursor) close() {
   398  	if c.lastBuffer != nil {
   399  		c.bucket.tx.db.putBuffer(c.lastBuffer)
   400  		c.lastBuffer = nil
   401  	}
   402  }