From 1563d3e5c8ba346bfe450e3c07dbd3d5f73e792a Mon Sep 17 00:00:00 2001
From: Ibrahim Jarif <ibrahim@dgraph.io>
Date: Thu, 23 Apr 2020 15:34:46 +0530
Subject: [PATCH 1/4] Restore: Handle incorrect encryption key

---
 worker/restore.go | 35 +++++++++++++++++++++++++++--------
 1 file changed, 27 insertions(+), 8 deletions(-)

diff --git a/worker/restore.go b/worker/restore.go
index cbc306c8c0a..c12c21bbc75 100644
--- a/worker/restore.go
+++ b/worker/restore.go
@@ -47,6 +47,33 @@ func RunRestore(pdir, location, backupId, keyfile string) LoadResult {
 		func(r io.Reader, groupId int, preds predicateSet) (uint64, error) {
 
 			dir := filepath.Join(pdir, fmt.Sprintf("p%d", groupId))
+			r, err := enc.GetReader(keyfile, r)
+			if err != nil {
+				return 0, err
+			}
+
+			gzReader, err := gzip.NewReader(r)
+			if err != nil {
+				return 0, errors.Wrap(err,
+					"Unable to read backup. Ensure encryption key is correctly set")
+			}
+			// The badger DB should be opened only after creating the backup
+			// file reader. The following sequence of events can occur if
+			// badger is opened before creating a reader -
+			// 1. Backup is created without encryption (alpha running without encryption)
+			// 2. Restore is called on an unencrypted backup with encryption set.
+			//		dgraph restore --encryption_key=xxx
+			//		If badger was opened before enc.GetReader() then the new
+			//		badger DB would have encryption set. Eventually
+			//		enc.GetReader() would crash because of invalid key but
+			//		encryption would be set in badger.
+			// 3. Re-attempt restore with encryption key not set.
+			//		dgraph restore
+			//		This command would fail with "encryption key
+			//		mismatch" because the badger DB initialized by the previous
+			//		restore command has encryption set and the current call
+			//		doesn't have the encryption key set.
+			//
 			db, err := badger.OpenManaged(badger.DefaultOptions(dir).
 				WithSyncWrites(false).
 				WithTableLoadingMode(options.MemoryMap).
@@ -60,14 +87,6 @@ func RunRestore(pdir, location, backupId, keyfile string) LoadResult {
 			if !pathExist(dir) {
 				fmt.Println("Creating new db:", dir)
 			}
-			r, err = enc.GetReader(keyfile, r)
-			if err != nil {
-				return 0, err
-			}
-			gzReader, err := gzip.NewReader(r)
-			if err != nil {
-				return 0, err
-			}
 			maxUid, err := loadFromBackup(db, gzReader, 0, preds)
 			if err != nil {
 				return 0, err

From 63c43778714c20ff017c4d9ae879e0bf2649ef72 Mon Sep 17 00:00:00 2001
From: Ibrahim Jarif <ibrahim@dgraph.io>
Date: Thu, 23 Apr 2020 15:58:48 +0530
Subject: [PATCH 2/4] fix error message

---
 worker/restore.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/worker/restore.go b/worker/restore.go
index c12c21bbc75..1d2c067cc22 100644
--- a/worker/restore.go
+++ b/worker/restore.go
@@ -55,7 +55,7 @@ func RunRestore(pdir, location, backupId, keyfile string) LoadResult {
 			gzReader, err := gzip.NewReader(r)
 			if err != nil {
 				return 0, errors.Wrap(err,
-					"Unable to read backup. Ensure encryption key is correctly set")
+					"Unable to read the backup. Ensure encryption key is correct.")
 			}
 			// The badger DB should be opened only after creating the backup
 			// file reader. The following sequence of events can occur if

From 5def236435babca03e3952d363a765d38ea2a2f9 Mon Sep 17 00:00:00 2001
From: Ibrahim Jarif <ibrahim@dgraph.io>
Date: Thu, 23 Apr 2020 22:11:06 +0530
Subject: [PATCH 3/4] Address review comments

---
 worker/restore.go | 17 +----------------
 1 file changed, 1 insertion(+), 16 deletions(-)

diff --git a/worker/restore.go b/worker/restore.go
index 1d2c067cc22..c430d6ab947 100644
--- a/worker/restore.go
+++ b/worker/restore.go
@@ -58,22 +58,7 @@ func RunRestore(pdir, location, backupId, keyfile string) LoadResult {
 					"Unable to read the backup. Ensure encryption key is correct.")
 			}
 			// The badger DB should be opened only after creating the backup
-			// file reader. The following sequence of events can occur if
-			// badger is opened before creating a reader -
-			// 1. Backup is created without encryption (alpha running without encryption)
-			// 2. Restore is called on an unencrypted backup with encryption set.
-			//		dgraph restore --encryption_key=xxx
-			//		If badger was opened before enc.GetReader() then the new
-			//		badger DB would have encryption set. Eventually
-			//		enc.GetReader() would crash because of invalid key but
-			//		encryption would be set in badger.
-			// 3. Re-attempt restore with encryption key not set.
-			//		dgraph restore
-			//		This command would fail with "encryption key
-			//		mismatch" because the badger DB initialized by the previous
-			//		restore command has encryption set and the current call
-			//		doesn't have the encryption key set.
-			//
+			// file reader and verifying the encryption in the backup file.
 			db, err := badger.OpenManaged(badger.DefaultOptions(dir).
 				WithSyncWrites(false).
 				WithTableLoadingMode(options.MemoryMap).

From 10e701baec33f672ef0925fe11ca4081022afc2f Mon Sep 17 00:00:00 2001
From: Ibrahim Jarif <ibrahim@dgraph.io>
Date: Thu, 23 Apr 2020 22:16:32 +0530
Subject: [PATCH 4/4] fixup

---
 worker/restore.go | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/worker/restore.go b/worker/restore.go
index c430d6ab947..1accaed814a 100644
--- a/worker/restore.go
+++ b/worker/restore.go
@@ -54,8 +54,12 @@ func RunRestore(pdir, location, backupId, keyfile string) LoadResult {
 
 			gzReader, err := gzip.NewReader(r)
 			if err != nil {
-				return 0, errors.Wrap(err,
-					"Unable to read the backup. Ensure encryption key is correct.")
+				if len(keyfile) != 0 {
+					err = errors.Wrap(err,
+						"Unable to read the backup. Ensure the encryption key is correct.")
+				}
+				return 0, err
+
 			}
 			// The badger DB should be opened only after creating the backup
 			// file reader and verifying the encryption in the backup file.