From 03330b3f7ca7dc06114d34fe34679eef73f821e7 Mon Sep 17 00:00:00 2001
From: Nikhil Erigila <60037808+nikhilerigila09@users.noreply.github.com>
Date: Thu, 19 Sep 2024 14:27:36 +0530
Subject: [PATCH] [fix][broker] Fix incomplete NAR file extraction which
 prevents broker from starting (#23274)

---
 .../apache/pulsar/common/nar/NarUnpacker.java | 29 ++++++++++++++-----
 .../pulsar/common/nar/NarUnpackerTest.java    | 11 +++++++
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/nar/NarUnpacker.java b/pulsar-common/src/main/java/org/apache/pulsar/common/nar/NarUnpacker.java
index e1806836d2833..ef802674b421a 100644
--- a/pulsar-common/src/main/java/org/apache/pulsar/common/nar/NarUnpacker.java
+++ b/pulsar-common/src/main/java/org/apache/pulsar/common/nar/NarUnpacker.java
@@ -32,7 +32,9 @@
 import java.io.RandomAccessFile;
 import java.nio.channels.FileChannel;
 import java.nio.channels.FileLock;
+import java.nio.file.Files;
 import java.nio.file.Path;
+import java.nio.file.StandardCopyOption;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.util.Base64;
@@ -86,19 +88,32 @@ static File doUnpackNar(final File nar, final File baseWorkingDirectory, Runnabl
             try (FileChannel channel = new RandomAccessFile(lockFile, "rw").getChannel();
                  FileLock lock = channel.lock()) {
                 File narWorkingDirectory = new File(parentDirectory, md5Sum);
-                if (narWorkingDirectory.mkdir()) {
+                if (!narWorkingDirectory.exists()) {
+                    File narExtractionTempDirectory = new File(parentDirectory, md5Sum + ".tmp");
+                    if (narExtractionTempDirectory.exists()) {
+                        FileUtils.deleteFile(narExtractionTempDirectory, true);
+                    }
+                    if (!narExtractionTempDirectory.mkdir()) {
+                        throw new IOException("Cannot create " + narExtractionTempDirectory);
+                    }
                     try {
-                        log.info("Extracting {} to {}", nar, narWorkingDirectory);
+                        log.info("Extracting {} to {}", nar, narExtractionTempDirectory);
                         if (extractCallback != null) {
                             extractCallback.run();
                         }
-                        unpack(nar, narWorkingDirectory);
+                        unpack(nar, narExtractionTempDirectory);
                     } catch (IOException e) {
                         log.error("There was a problem extracting the nar file. Deleting {} to clean up state.",
-                                narWorkingDirectory, e);
-                        FileUtils.deleteFile(narWorkingDirectory, true);
+                                narExtractionTempDirectory, e);
+                        try {
+                            FileUtils.deleteFile(narExtractionTempDirectory, true);
+                        } catch (IOException e2) {
+                            log.error("Failed to delete temporary directory {}", narExtractionTempDirectory, e2);
+                        }
                         throw e;
                     }
+                    Files.move(narExtractionTempDirectory.toPath(), narWorkingDirectory.toPath(),
+                            StandardCopyOption.ATOMIC_MOVE);
                 }
                 return narWorkingDirectory;
             }
@@ -166,7 +181,7 @@ private static void makeFile(final InputStream inputStream, final File file) thr
      * @throws IOException
      *             if cannot read file
      */
-    private static byte[] calculateMd5sum(final File file) throws IOException {
+    protected static byte[] calculateMd5sum(final File file) throws IOException {
         try (final FileInputStream inputStream = new FileInputStream(file)) {
             // codeql[java/weak-cryptographic-algorithm] - md5 is sufficient for this use case
             final MessageDigest md5 = MessageDigest.getInstance("md5");
@@ -184,4 +199,4 @@ private static byte[] calculateMd5sum(final File file) throws IOException {
             throw new IllegalArgumentException(nsae);
         }
     }
-}
+}
\ No newline at end of file
diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/nar/NarUnpackerTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/nar/NarUnpackerTest.java
index a1f915c8b7828..1c3a2c276537b 100644
--- a/pulsar-common/src/test/java/org/apache/pulsar/common/nar/NarUnpackerTest.java
+++ b/pulsar-common/src/test/java/org/apache/pulsar/common/nar/NarUnpackerTest.java
@@ -118,6 +118,17 @@ public static void main(String[] args) {
         }
     }
 
+    @Test
+    void shouldReExtractWhenUnpackedDirectoryIsMissing() throws IOException {
+        AtomicInteger extractCounter = new AtomicInteger();
+
+        File narWorkingDirectory = NarUnpacker.doUnpackNar(sampleZipFile, extractDirectory, extractCounter::incrementAndGet);
+        FileUtils.deleteFile(narWorkingDirectory, true);
+        NarUnpacker.doUnpackNar(sampleZipFile, extractDirectory, extractCounter::incrementAndGet);
+
+        assertEquals(extractCounter.get(), 2);
+    }
+
     @Test
     void shouldExtractFilesOnceInDifferentProcess() throws InterruptedException {
         int processes = 5;