Merge 5295a6b into 12fb082

badger/cmd/bank.go

@@ -70,6 +70,7 @@ var mmap bool
 var checkStream bool
 var checkSubscriber bool
 var verbose bool
+var encryptionEnabled bool
 
 const keyPrefix = "account:"
 
@@ -98,6 +99,8 @@ func init() {
 
 	bankDisect.Flags().IntVarP(&numPrevious, "previous", "p", 12,
 		"Starting from the violation txn, how many previous versions to retrieve.")
+	bankDisect.Flags().BoolVarP(&encryptionEnabled, "encryption", "e", false,
+		"If it is true, badger will encrypt all the data stored on the disk.")
 }
 
 func key(account int) []byte {
@@ -338,6 +341,14 @@ func runTest(cmd *cobra.Command, args []string) error {
 	}
 	log.Printf("Opening DB with options: %+v\n", opts)
 
+	if encryptionEnabled {
+		key := make([]byte, 32)
+		_, err := rand.Read(key)
+		if err != nil {
+			return err
+		}
+		opts = opts.WithEncryptionKey(key)
+	}
 	db, err := badger.Open(opts)
 	if err != nil {
 		return err

badger/cmd/info.go

@@ -321,7 +321,7 @@ func printInfo(dir, valueDir string) error {
 		})
 		for _, tableID := range tableIDs {
 			tableFile := table.IDToFilename(tableID)
-			tm, ok1 := manifest.Tables[tableID]
+			_, ok1 := manifest.Tables[tableID]
 			file, ok2 := fileinfoByName[tableFile]
 			if ok1 && ok2 {
 				fileinfoMarked[tableFile] = true
@@ -333,8 +333,8 @@ func printInfo(dir, valueDir string) error {
 				}
 				levelSizes[level] += fileSize
 				// (Put level on every line to make easier to process with sed/perl.)
-				fmt.Printf("[%25s] %-12s %6s L%d %x%s\n", dur(baseTime, file.ModTime()),
-					tableFile, hbytes(fileSize), level, tm.Checksum, emptyString)
+				fmt.Printf("[%25s] %-12s %6s L%d %s\n", dur(baseTime, file.ModTime()),
+					tableFile, hbytes(fileSize), level, emptyString)
 			} else {
 				fmt.Printf("%s [MISSING]\n", tableFile)
 				numMissing++

db.go

@@ -88,7 +88,8 @@ type DB struct {
 
 	orc *oracle
 
-	pub *publisher
+	pub      *publisher
+	registry *KeyRegistry
 }
 
 const (
@@ -289,6 +290,18 @@ func Open(opt Options) (db *DB, err error) {
 		pub:           newPublisher(),
 	}
 
+	krOpt := KeyRegistryOptions{
+		ReadOnly:                      opt.ReadOnly,
+		Dir:                           opt.Dir,
+		EncryptionKey:                 opt.EncryptionKey,
+		EncryptionKeyRotationDuration: opt.EncryptionKeyRotationDuration,
+	}
+
+	kr, err := OpenKeyRegistry(krOpt)
+	if err != nil {
+		return nil, err
+	}
+	db.registry = kr
 	// Calculate initial size.
 	db.calculateSize()
 	db.closers.updateSize = y.NewCloser(1)
@@ -460,6 +473,9 @@ func (db *DB) close() (err error) {
 	if manifestErr := db.manifest.close(); err == nil {
 		err = errors.Wrap(manifestErr, "DB.Close")
 	}
+	if registryErr := db.registry.Close(); err == nil {
+		err = errors.Wrap(registryErr, "DB.Close")
+	}
 
 	// Fsync directories to ensure that lock file, and any other removed files whose directory
 	// we haven't specifically fsynced, are guaranteed to have their directory entry removal
@@ -891,15 +907,20 @@ func (db *DB) handleFlushTask(ft flushTask) error {
 	headTs := y.KeyWithTs(head, db.orc.nextTs())
 	ft.mt.Put(headTs, y.ValueStruct{Value: val})
 
+	dk, err := db.registry.latestDataKey()
+	if err != nil {
+		return y.Wrapf(err, "failed to get datakey in db.handleFlushTask")
+	}
 	bopts := table.Options{
 		BlockSize:          db.opt.BlockSize,
 		BloomFalsePositive: db.opt.BloomFalsePositive,
+		DataKey:            dk,
 	}
 	tableData := buildL0Table(ft, bopts)
 
 	fileID := db.lc.reserveFileID()
 	if db.opt.KeepL0InMemory {
-		tbl, err := table.OpenInMemoryTable(tableData, fileID)
+		tbl, err := table.OpenInMemoryTable(tableData, fileID, dk)
 		if err != nil {
 			return errors.Wrapf(err, "failed to open table in memory")
 		}
@@ -928,6 +949,7 @@ func (db *DB) handleFlushTask(ft flushTask) error {
 	opts := table.Options{
 		LoadingMode: db.opt.TableLoadingMode,
 		ChkMode:     db.opt.ChecksumVerificationMode,
+		DataKey:     dk,
 	}
 	tbl, err := table.OpenTable(fd, opts)
 	if err != nil {
@@ -1543,3 +1565,8 @@ func (db *DB) Subscribe(ctx context.Context, cb func(kv *KVList), prefixes ...[]
 		}
 	}
 }
+
+// shouldEncrypt returns bool, which tells whether to encrypt or not.
+func (db *DB) shouldEncrypt() bool {
+	return len(db.opt.EncryptionKey) > 0
+}

db2_test.go

@@ -542,3 +542,54 @@ func createTableWithRange(t *testing.T, db *DB, start, end int) *table.Table {
 	require.NoError(t, err)
 	return tab
 }
+
+func TestReadSameVlog(t *testing.T) {
+	key := func(i int) []byte {
+		return []byte(fmt.Sprintf("%d%10d", i, i))
+	}
+	testReadingSameKey := func(t *testing.T, opt Options) {
+		// Forcing to read all values from vlog.
+		opt.ValueThreshold = 1
+		db, err := Open(opt)
+		require.NoError(t, err)
+		for i := 0; i < 50; i++ {
+			err = db.Update(func(txn *Txn) error {
+				return txn.Set(key(i), key(i))
+			})
+		}
+		// reading it again several times
+		for i := 0; i < 50; i++ {
+			for j := 0; j < 10; j++ {
+				err = db.View(func(txn *Txn) error {
+					item, err := txn.Get(key(i))
+					require.NoError(t, err)
+					require.Equal(t, key(i), getItemValue(t, item))
+					return nil
+				})
+			}
+		}
+	}
+
+	t.Run("Test Read Again Plain Text", func(t *testing.T) {
+		dir, err := ioutil.TempDir("", "badger-test")
+		require.NoError(t, err)
+		defer os.RemoveAll(dir)
+
+		opt := getTestOptions(dir)
+		testReadingSameKey(t, opt)
+	})
+
+	t.Run("Test Read Again Encryption", func(t *testing.T) {
+		dir, err := ioutil.TempDir("", "badger-test")
+		require.NoError(t, err)
+		defer os.RemoveAll(dir)
+		opt := getTestOptions(dir)
+		// Generate encryption key.
+		eKey := make([]byte, 32)
+		_, err = rand.Read(eKey)
+		require.NoError(t, err)
+		opt.EncryptionKey = eKey
+		// Test reading with encryption enabled.
+		testReadingSameKey(t, opt)
+	})
+}