From 09f6c9ef81cf88647961e5667543d93186354002 Mon Sep 17 00:00:00 2001
From: Craig Raw <craigraw@gmail.com>
Date: Mon, 21 Jun 2021 11:50:58 +0200
Subject: [PATCH] ensure cleanup of migrated wallets when importing

---
 .../com/sparrowwallet/sparrow/io/Sparrow.java | 23 +++++++++++++++----
 .../com/sparrowwallet/sparrow/io/Storage.java |  8 +++++--
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/src/main/java/com/sparrowwallet/sparrow/io/Sparrow.java b/src/main/java/com/sparrowwallet/sparrow/io/Sparrow.java
index a28a06db..aeea0641 100644
--- a/src/main/java/com/sparrowwallet/sparrow/io/Sparrow.java
+++ b/src/main/java/com/sparrowwallet/sparrow/io/Sparrow.java
@@ -83,23 +83,38 @@ public class Sparrow implements WalletImport, WalletExport {
 
     @Override
     public Wallet importWallet(InputStream inputStream, String password) throws ImportException {
+        Storage storage = null;
+        Wallet wallet = null;
         File tempFile = null;
         try {
             tempFile = File.createTempFile("sparrow", null);
             java.nio.file.Files.copy(inputStream, tempFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
-            Storage storage = new Storage(PersistenceType.JSON, tempFile);
+            storage = new Storage(PersistenceType.JSON, tempFile);
             if(!isEncrypted(tempFile)) {
-                return storage.loadUnencryptedWallet().getWallet();
+                wallet = storage.loadUnencryptedWallet().getWallet();
             } else {
                 WalletBackupAndKey walletBackupAndKey = storage.loadEncryptedWallet(password);
-                walletBackupAndKey.getWallet().decrypt(walletBackupAndKey.getKey());
-                return walletBackupAndKey.getWallet();
+                wallet = walletBackupAndKey.getWallet();
+                wallet.decrypt(walletBackupAndKey.getKey());
             }
+
+            return wallet;
         } catch(IOException | StorageException e) {
             log.error("Error importing Sparrow wallet", e);
             throw new ImportException("Error importing Sparrow wallet", e);
         } finally {
+            if(storage != null) {
+                storage.close();
+            }
+
             if(tempFile != null) {
+                if(wallet != null) {
+                    File migratedWalletFile = Storage.getExistingWallet(tempFile.getParentFile(), wallet.getName());
+                    if(migratedWalletFile != null) {
+                        migratedWalletFile.delete();
+                    }
+                }
+
                 tempFile.delete();
             }
         }
diff --git a/src/main/java/com/sparrowwallet/sparrow/io/Storage.java b/src/main/java/com/sparrowwallet/sparrow/io/Storage.java
index 79ad56fa..43c58ce2 100644
--- a/src/main/java/com/sparrowwallet/sparrow/io/Storage.java
+++ b/src/main/java/com/sparrowwallet/sparrow/io/Storage.java
@@ -289,13 +289,17 @@ public class Storage {
     }
 
     public static File getExistingWallet(String walletName) {
-        File encrypted = new File(getWalletsDir(), walletName.trim());
+        return getExistingWallet(getWalletsDir(), walletName);
+    }
+
+    public static File getExistingWallet(File walletsDir, String walletName) {
+        File encrypted = new File(walletsDir, walletName.trim());
         if(encrypted.exists()) {
             return encrypted;
         }
 
         for(PersistenceType persistenceType : PersistenceType.values()) {
-            File unencrypted = new File(getWalletsDir(), walletName.trim() + "." + persistenceType.getExtension());
+            File unencrypted = new File(walletsDir, walletName.trim() + "." + persistenceType.getExtension());
             if(unencrypted.exists()) {
                 return unencrypted;
             }