re-encrypt restore from encrypted backups (#5140)

As part of the encrypted backup/restore feature, currently, the restore tool restores the "p" dir unencrypted whether the backup is encrypted or not. 

This story allow restores to be encrypted with the same key file used to decrypt the backup. If backup was unencrypted, the restore will remain unencrypted.

ee/backup/run.go

@@ -105,8 +105,8 @@ $ dgraph restore -p . -l /var/backups/dgraph -z localhost:5080
 	flag.StringVarP(&opt.zero, "zero", "z", "", "gRPC address for Dgraph zero. ex: localhost:5080")
 	flag.StringVarP(&opt.backupId, "backup_id", "", "", "The ID of the backup series to "+
 		"restore. If empty, it will restore the latest series.")
-	flag.StringVarP(&opt.keyfile, "keyfile", "k", "",
-		"Key file to decrypt the backup")
+	flag.StringVarP(&opt.keyfile, "keyfile", "k", "", "Key file to decrypt the backup. "+
+		"The same key is also used to re-encrypt the restored data.")
 	_ = Restore.Cmd.MarkFlagRequired("postings")
 	_ = Restore.Cmd.MarkFlagRequired("location")
 }

systest/backup/encryption/backup_test.go

@@ -271,7 +271,9 @@ func runRestore(t *testing.T, lastDir string, commitTs uint64) map[string]string
 	require.NoError(t, os.RemoveAll(restoreDir))
 
 	t.Logf("--- Restoring from: %q", localBackupDst)
-	argv := []string{"dgraph", "restore", "-l", localBackupDst, "-p", "data/restore", "-k", "../../../ee/enc/enc-key"}
+	testutil.KeyFile = "../../../ee/enc/enc-key"
+	argv := []string{"dgraph", "restore", "-l", localBackupDst, "-p", "data/restore",
+		"-k", testutil.KeyFile}
 	cwd, err := os.Getwd()
 	require.NoError(t, err)
 	err = testutil.ExecWithOpts(argv, testutil.CmdOpts{Dir: cwd})

testutil/backup.go

@@ -21,16 +21,21 @@ import (
 
 	"github.com/dgraph-io/badger/v2"
 	"github.com/dgraph-io/badger/v2/options"
+	"github.com/dgraph-io/dgraph/ee/enc"
 	"github.com/dgraph-io/dgraph/posting"
 	"github.com/dgraph-io/dgraph/protos/pb"
 	"github.com/dgraph-io/dgraph/types"
 	"github.com/dgraph-io/dgraph/x"
 )
 
+// KEYFILE is set to the path of the file containing the key. Used for testing purposes only.
+var KeyFile string
+
 func openDgraph(pdir string) (*badger.DB, error) {
 	opt := badger.DefaultOptions(pdir).WithTableLoadingMode(options.MemoryMap).
 		// TOOD(Ibrahim): Remove compression level once badger is updated.
-		WithReadOnly(true).WithZSTDCompressionLevel(1)
+		WithReadOnly(true).WithZSTDCompressionLevel(1).
+		WithEncryptionKey(enc.ReadEncryptionKeyFile(KeyFile))
 	return badger.OpenManaged(opt)
 }
 

worker/restore.go

@@ -51,7 +51,8 @@ func RunRestore(pdir, location, backupId, keyfile string) LoadResult {
 				WithSyncWrites(false).
 				WithTableLoadingMode(options.MemoryMap).
 				WithValueThreshold(1 << 10).
-				WithNumVersionsToKeep(math.MaxInt32))
+				WithNumVersionsToKeep(math.MaxInt32).
+				WithEncryptionKey(enc.ReadEncryptionKeyFile(keyfile)))
 			if err != nil {
 				return 0, err
 			}