feat(raftwal): Add support for encryption in raftwal (#6714)

Add support for encryption of the raftwal. The data in the entries of raftwal
is encrypted if the encryption flag is set. Also add tests for raftwal encryption.
NOTE: If encryption is enabled once then it cannot be disabled later.

go.sum

@@ -6,7 +6,6 @@ contrib.go.opencensus.io/exporter/prometheus v0.1.0 h1:SByaIoWwNgMdPSgl5sMqM2KDE
 contrib.go.opencensus.io/exporter/prometheus v0.1.0/go.mod h1:cGFniUXGZlKRjzOyuZJ6mgB+PgBcCIa79kEKR8YCW+A=
 github.com/99designs/gqlgen v0.13.1-0.20200928230741-819e751c2416 h1:8qbuDq7x3pPeEUmfa2wPKuN2G5Q/+znZWAJWZJXTjDA=
 github.com/99designs/gqlgen v0.13.1-0.20200928230741-819e751c2416/go.mod h1:NV130r6f4tpRWuAI+zsrSdooO/eWUv+Gyyoi3rEfXIk=
-github.com/AndreasBriese/bbloom v0.0.0-20190306092124-e2d15f34fcf9 h1:HD8gA2tkByhMAwYaFAX9w2l7vxvBQ5NMoxDrkhqhtn4=
 github.com/AndreasBriese/bbloom v0.0.0-20190306092124-e2d15f34fcf9/go.mod h1:bOvUY6CB00SOBii9/FifXqc0awNKxLFCL/+pkDPuyl8=
 github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
@@ -73,9 +72,7 @@ github.com/coreos/go-etcd v2.0.0+incompatible/go.mod h1:Jez6KQU2B/sWsbdaef3ED8Nz
 github.com/coreos/go-semver v0.2.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
 github.com/coreos/go-systemd v0.0.0-20180511133405-39ca1b05acc7/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
 github.com/coreos/pkg v0.0.0-20160727233714-3ac0863d7acf/go.mod h1:E3G3o1h8I7cfcXa63jLwjI0eiQQMgzzUDFVpN/nH/eA=
-github.com/cpuguy83/go-md2man v1.0.10 h1:BSKMNlYxDvnunlTymqtgONjNnaRV1sTpcovwwjF22jk=
 github.com/cpuguy83/go-md2man v1.0.10/go.mod h1:SmD6nW6nTyfqj6ABTjUi3V3JVMnlJmwcJI5acqYI6dE=
-github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d h1:U+s90UTSYgptZMwQh2aRr3LuazLJIa+Pg3Kc1ylSYVY=
 github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU=
 github.com/d4l3k/messagediff v1.2.1 h1:ZcAIMYsUg0EAp9X+tt8/enBE/Q8Yd5kzPynLyKptt9U=
 github.com/d4l3k/messagediff v1.2.1/go.mod h1:Oozbb1TVXFac9FtSIxHBMnBCq2qeH/2KkEQxENCrlLo=
@@ -90,7 +87,6 @@ github.com/dgraph-io/dgo/v200 v200.0.0-20200805103119-a3544c464dd6 h1:toHzMCdCUg
 github.com/dgraph-io/dgo/v200 v200.0.0-20200805103119-a3544c464dd6/go.mod h1:rHa+h3kI4M8ASOirxyIyNeXBfHFgeskVUum2OrDMN3U=
 github.com/dgraph-io/graphql-transport-ws v0.0.0-20200916064635-48589439591b h1:PDEhlwHpkEQ5WBfOOKZCNZTXFDGyCEWTYDhxGQbyIpk=
 github.com/dgraph-io/graphql-transport-ws v0.0.0-20200916064635-48589439591b/go.mod h1:7z3c/5w0sMYYZF5bHsrh8IH4fKwG5O5Y70cPH1ZLLRQ=
-github.com/dgraph-io/ristretto v0.0.4-0.20201013194302-6d6fac64beae h1:yh5085twGpsgfuu56DXKOM3SKyZKQPskJIoMNb3jzos=
 github.com/dgraph-io/ristretto v0.0.4-0.20201013194302-6d6fac64beae/go.mod h1:bDI4cDaalvYSji3vBVDKrn9ouDZrwN974u8ZO/AhYXs=
 github.com/dgraph-io/ristretto v0.0.4-0.20201013234705-28aba7a42dfa h1:gAJJ+Ln7gBkeLEoKIuFL1p+YjbihmgQYWpOIu8XE6JM=
 github.com/dgraph-io/ristretto v0.0.4-0.20201013234705-28aba7a42dfa/go.mod h1:bDI4cDaalvYSji3vBVDKrn9ouDZrwN974u8ZO/AhYXs=

raftwal/encryption_test.go

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2020 Dgraph Labs, Inc. and Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package raftwal
+
+import (
+	"io/ioutil"
+	"math/rand"
+	"os"
+	"testing"
+
+	"github.com/dgraph-io/dgraph/x"
+	"github.com/stretchr/testify/require"
+	"go.etcd.io/etcd/raft/raftpb"
+)
+
+func TestEntryReadWrite(t *testing.T) {
+	x.WorkerConfig.EncryptionKey = []byte("badger16byteskey")
+	dir, err := ioutil.TempDir("", "raftwal")
+	require.NoError(t, err)
+	el, err := openWal(dir)
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+
+	// generate some random data
+	data := make([]byte, rand.Intn(1000))
+	rand.Read(data)
+
+	require.NoError(t, el.AddEntries([]raftpb.Entry{{Index: 1, Term: 1, Data: data}}))
+	entries := el.allEntries(0, 100, 10000)
+	require.Equal(t, 1, len(entries))
+	require.Equal(t, uint64(1), entries[0].Index)
+	require.Equal(t, uint64(1), entries[0].Term)
+	require.Equal(t, data, entries[0].Data)
+
+	// Open the wal file again.
+	el2, err := openWal(dir)
+	require.NoError(t, err)
+	entries = el2.allEntries(0, 100, 10000)
+	require.Equal(t, 1, len(entries))
+	require.Equal(t, uint64(1), entries[0].Index)
+	require.Equal(t, uint64(1), entries[0].Term)
+	require.Equal(t, data, entries[0].Data)
+
+	// Opening it with a wrong key fails.
+	x.WorkerConfig.EncryptionKey = []byte("other16byteskeys")
+	_, err = openWal(dir)
+	require.EqualError(t, err, "Encryption key mismatch")
+
+	// Opening it without encryption key fails.
+	x.WorkerConfig.EncryptionKey = nil
+	_, err = openWal(dir)
+	require.EqualError(t, err, "Logfile is encrypted but encryption key is nil")
+}

raftwal/log.go

@@ -17,14 +17,20 @@
 package raftwal
 
 import (
+	"crypto/aes"
+	cryptorand "crypto/rand"
 	"encoding/binary"
 	"fmt"
 	"os"
 	"path"
 	"sort"
 	"strconv"
 	"strings"
+	"time"
 
+	"github.com/dgraph-io/badger/v2"
+	"github.com/dgraph-io/badger/v2/pb"
+	"github.com/dgraph-io/badger/v2/y"
 	"github.com/dgraph-io/dgraph/x"
 	"github.com/dgraph-io/ristretto/z"
 	"github.com/golang/glog"
@@ -43,6 +49,9 @@ const (
 	maxNumEntries = 30000
 	// logFileOffset is offset in the log file where data is stored.
 	logFileOffset = 1 << 20 // 1MB
+	// encOffset is offset in the log file where keyID (first 8 bytes)
+	// and baseIV (remaining 8 bytes) are stored.
+	encOffset = logFileOffset - 16 // 1MB - 16B
 	// logFileSize is the initial size of the log file.
 	logFileSize = 16 << 30
 	// entrySize is the size in bytes of a single entry.
@@ -75,6 +84,10 @@ func marshalEntry(b []byte, term, index, do, typ uint64) {
 type logFile struct {
 	*z.MmapFile
 	fid int64
+
+	registry *badger.KeyRegistry
+	dataKey  *pb.DataKey
+	baseIV   []byte
 }
 
 func logFname(dir string, id int64) string {
@@ -86,19 +99,55 @@ func logFname(dir string, id int64) string {
 func openLogFile(dir string, fid int64) (*logFile, error) {
 	glog.V(3).Infof("opening log file: %d\n", fid)
 	fpath := logFname(dir, fid)
+	lf := &logFile{
+		fid: fid,
+	}
+	var err error
+	encKey := x.WorkerConfig.EncryptionKey
+	// Initialize the registry for logFile if encryption in enabled.
+	// NOTE: If encryption is enabled then there is no going back because if we disable it
+	// later then the older log files which were previously encrypted can't be opened.
+	if len(encKey) > 0 {
+		krOpt := badger.KeyRegistryOptions{
+			ReadOnly:                      false,
+			Dir:                           dir,
+			EncryptionKey:                 encKey,
+			EncryptionKeyRotationDuration: 10 * 24 * time.Hour,
+			InMemory:                      false,
+		}
+		// This won't open Badger. It would only use its key registry.
+		if lf.registry, err = badger.OpenKeyRegistry(krOpt); err != nil {
+			return nil, err
+		}
+	}
 	// Open the file in read-write mode and create it if it doesn't exist yet.
-	mf, err := z.OpenMmapFile(fpath, os.O_RDWR|os.O_CREATE, logFileSize)
+	lf.MmapFile, err = z.OpenMmapFile(fpath, os.O_RDWR|os.O_CREATE, logFileSize)
 
 	if err == z.NewFile {
 		glog.V(3).Infof("New file: %d\n", fid)
-		z.ZeroOut(mf.Data, 0, logFileOffset)
-	} else {
+		z.ZeroOut(lf.Data, 0, logFileOffset)
+		if err = lf.bootstrap(); err != nil {
+			return nil, err
+		}
+	} else if err != nil {
 		x.Check(err)
-	}
+	} else {
+		buf := lf.Data[encOffset : encOffset+16]
+		keyID := binary.BigEndian.Uint64(buf[:8])
 
-	lf := &logFile{
-		MmapFile: mf,
-		fid:      fid,
+		// If keyID is non-zero, then the opened file is encrypted.
+		if keyID != 0 {
+			// Logfile is encrypted but encryption key is not provided.
+			if encKey == nil {
+				return nil, errors.New("Logfile is encrypted but encryption key is nil")
+			}
+			// retrieve datakey from the keyID of the logfile.
+			if lf.dataKey, err = lf.registry.DataKey(keyID); err != nil {
+				return nil, err
+			}
+			lf.baseIV = buf[8:]
+			y.AssertTrue(len(lf.baseIV) == 8)
+		}
 	}
 	return lf, nil
 }
@@ -129,6 +178,16 @@ func (lf *logFile) GetRaftEntry(idx int) raftpb.Entry {
 			re.Data = append(re.Data, data...)
 		}
 	}
+	// Decrypt the data if encryption is enabled.
+	if lf.dataKey != nil && len(re.Data) > 0 {
+		// No need to worry about mmap. Because, XORBlock allocates a byte array to do the
+		// XOR. So, the given slice is not being mutated.
+		// NOTE: We can potentially use allocator for this allocation.
+		decoded, err := y.XORBlockAllocate(
+			re.Data, lf.dataKey.Data, lf.generateIV(entry.DataOffset()))
+		x.Check(err)
+		re.Data = decoded
+	}
 	return re
 }
 
@@ -242,3 +301,50 @@ func getLogFiles(dir string) ([]*logFile, error) {
 	})
 	return files, nil
 }
+
+// KeyID returns datakey's ID.
+func (lf *logFile) keyID() uint64 {
+	if lf.dataKey == nil {
+		// If there is no datakey, then we'll return 0. Which means no encryption.
+		return 0
+	}
+	return lf.dataKey.KeyId
+}
+
+// generateIV will generate IV by appending given offset with the base IV.
+func (lf *logFile) generateIV(offset uint64) []byte {
+	iv := make([]byte, aes.BlockSize)
+	// IV is of 16 bytes, in which first 8 bytes are obtained from baseIV
+	// and the remaining 8 bytes is obtained from the offset.
+	y.AssertTrue(8 == copy(iv[:8], lf.baseIV))
+	binary.BigEndian.PutUint64(iv[8:], offset)
+	return iv
+}
+
+// bootstrap will initialize the log file with key id and baseIV.
+// The below figure shows the layout of log file.
+// +----------------+------------------+------------------+
+// | keyID(8 bytes) |  baseIV(8 bytes) |	 entry...     |
+// +----------------+------------------+------------------+
+func (lf *logFile) bootstrap() error {
+	// registry is nil if we don't have encryption enabled.
+	if lf.registry == nil {
+		return nil
+	}
+	var err error
+	// generate data key for the log file.
+	if lf.dataKey, err = lf.registry.LatestDataKey(); err != nil {
+		return y.Wrapf(err, "Error while retrieving datakey in logFile.bootstrap")
+	}
+	buf := lf.Data[encOffset : encOffset+16]
+	// Put keyID in the first 8 bytes.
+	binary.BigEndian.PutUint64(buf[:8], lf.keyID())
+
+	// fill in random bytes in the last 8 bytes of buf.
+	if _, err := cryptorand.Read(buf[8:]); err != nil {
+		return y.Wrapf(err, "Error while creating base IV, while creating logfile")
+	}
+	// Initialize base IV.
+	lf.baseIV = buf[8:]
+	return nil
+}