Refactor ResourceExtractionUtil to use NIO (#2240)

Author: felixbarny

apm-agent-attach/src/main/java/co/elastic/apm/attach/ElasticApmAttacher.java

@@ -188,7 +188,7 @@ private static File getAgentJarFile() {
             if (ElasticApmAttacher.class.getResource("/elastic-apm-agent.jar") == null) {
                 return null;
             }
-            return ResourceExtractionUtil.extractResourceToTempDirectory("elastic-apm-agent.jar", "elastic-apm-agent", ".jar");
+            return ResourceExtractionUtil.extractResourceToTempDirectory("elastic-apm-agent.jar", "elastic-apm-agent", ".jar").toFile();
         }
     }
 

apm-agent-common/src/main/java/co/elastic/apm/agent/common/util/ResourceExtractionUtil.java

@@ -18,98 +18,127 @@
  */
 package co.elastic.apm.agent.common.util;
 
-import java.io.File;
-import java.io.FileInputStream;
-import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.math.BigInteger;
+import java.nio.channels.Channels;
 import java.nio.channels.FileChannel;
 import java.nio.channels.FileLock;
+import java.nio.file.FileAlreadyExistsException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.attribute.FileAttribute;
+import java.nio.file.attribute.PosixFilePermissions;
+import java.nio.file.attribute.UserPrincipal;
 import java.security.DigestInputStream;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
+import java.util.EnumSet;
+
+import static java.nio.file.LinkOption.NOFOLLOW_LINKS;
+import static java.nio.file.StandardOpenOption.CREATE_NEW;
+import static java.nio.file.StandardOpenOption.READ;
+import static java.nio.file.StandardOpenOption.WRITE;
+import static java.nio.file.attribute.PosixFilePermission.OWNER_READ;
+import static java.nio.file.attribute.PosixFilePermission.OWNER_WRITE;
 
 public class ResourceExtractionUtil {
 
     /**
      * Extracts a classpath resource to {@code ${System.getProperty("java.io.tmpdir")}/$prefix-$hash.$suffix}.
      * If the file has already been extracted it will not be extracted again.
      *
-     * @param resource       The classpath resource to extract.
-     * @param prefix         The prefix of the extracted file.
-     * @param suffix         The suffix of the extracted file.
+     * @param resource The classpath resource to extract.
+     * @param prefix   The prefix of the extracted file.
+     * @param suffix   The suffix of the extracted file.
      * @return the extracted file.
      */
-    public static synchronized File extractResourceToTempDirectory(String resource, String prefix, String suffix) {
-        return extractResourceToDirectory(resource, prefix, suffix, System.getProperty("java.io.tmpdir"));
+    public static synchronized Path extractResourceToTempDirectory(String resource, String prefix, String suffix) {
+        return extractResourceToDirectory(resource, prefix, suffix, Paths.get(System.getProperty("java.io.tmpdir")));
     }
 
     /**
      * Extracts a classpath resource to {@code $directory/$prefix-$userHash-$hash.$suffix}.
      * If the file has already been extracted it will not be extracted again.
      *
-     * @param resource       The classpath resource to extract.
-     * @param prefix         The prefix of the extracted file.
-     * @param suffix         The suffix of the extracted file.
-     * @param directory      The directory in which the file is to be created, or null if the default temporary-file directory is to be used.
+     * @param resource  The classpath resource to extract.
+     * @param prefix    The prefix of the extracted file.
+     * @param suffix    The suffix of the extracted file.
+     * @param directory The directory in which the file is to be created, or null if the default temporary-file directory is to be used.
      * @return the extracted file.
      */
     /*
      * Why it's synchronized : if the same JVM try to lock file, we got an java.nio.channels.OverlappingFileLockException.
      * So we need to block until the file is totally written.
      */
-    public static synchronized File extractResourceToDirectory(String resource, String prefix, String suffix, String directory) {
+    public static synchronized Path extractResourceToDirectory(String resource, String prefix, String suffix, Path directory) {
         try (InputStream resourceStream = ResourceExtractionUtil.class.getResourceAsStream("/" + resource)) {
             if (resourceStream == null) {
                 throw new IllegalStateException(resource + " not found");
             }
-            String userHash = "";
-            if (System.getProperties().contains("user.name")) {
-                // we have to include current user name as multiple copies of the same agent could be attached
-                // to multiple JVMs, each running under a different user. Also, we have to make it path-friendly.
-                userHash = md5Hash(System.getProperty("user.name"));
-                userHash += "-";
-            }
-            String resourceHash = md5Hash(ResourceExtractionUtil.class.getResourceAsStream("/" + resource));
+            UserPrincipal currentUserPrincipal = getCurrentUserPrincipal();
+            // we have to include current user name as multiple copies of the same agent could be attached
+            // to multiple JVMs, each running under a different user. Hashing makes the name path-friendly.
+            String userHash = hash(currentUserPrincipal.getName());
+            // to guard against re-using previous versions
+            String resourceHash = hash(ResourceExtractionUtil.class.getResourceAsStream("/" + resource));
 
-            File tempFile = new File(directory, prefix + "-" + userHash + resourceHash + suffix);
-            if (!tempFile.exists()) {
-                try (FileOutputStream out = new FileOutputStream(tempFile)) {
-                    FileChannel channel = out.getChannel();
-                    // If multiple JVM start on same compute, they can write in same file
-                    // and this file will be corrupted.
-                    try (FileLock ignored = channel.lock()) {
-                        if (tempFile.length() == 0) {
-                            byte[] buffer = new byte[1024];
-                            for (int length; (length = resourceStream.read(buffer)) != -1; ) {
-                                out.write(buffer, 0, length);
-                            }
+            Path tempFile = directory.resolve(prefix + "-" + userHash.substring(0, 32) + "-" + resourceHash.substring(0, 32) + suffix);
+            try {
+                FileAttribute<?>[] attr;
+                if (tempFile.getFileSystem().supportedFileAttributeViews().contains("posix")) {
+                    attr = new FileAttribute[]{PosixFilePermissions.asFileAttribute(EnumSet.of(OWNER_WRITE, OWNER_READ))};
+                } else {
+                    attr = new FileAttribute[0];
+                }
+                try (FileChannel channel = FileChannel.open(tempFile, EnumSet.of(CREATE_NEW, WRITE), attr)) {
+                    // make other JVM instances wait until fully written
+                    try (FileLock writeLock = channel.lock()) {
+                        channel.transferFrom(Channels.newChannel(resourceStream), 0, Long.MAX_VALUE);
+                    }
+                }
+            } catch (FileAlreadyExistsException e) {
+                try (FileChannel channel = FileChannel.open(tempFile, READ, NOFOLLOW_LINKS)) {
+                    // wait until other JVM instances have fully written the file
+                    // multiple JVMs can read the file at the same time
+                    try (FileLock readLock = channel.lock(0, Long.MAX_VALUE, true)) {
+                        if (!hash(Files.newInputStream(tempFile)).equals(resourceHash)) {
+                            throw new IllegalStateException("Invalid checksum of " + tempFile + ". Please delete this file.");
+                        } else if (!Files.getOwner(tempFile).equals(currentUserPrincipal)) {
+                            throw new IllegalStateException("File " + tempFile + " is not owned by '" + currentUserPrincipal.getName() + "'. Please delete this file.");
                         }
                     }
                 }
-            } else if (!md5Hash(new FileInputStream(tempFile)).equals(resourceHash)) {
-                throw new IllegalStateException("Invalid MD5 checksum of " + tempFile + ". Please delete this file.");
             }
-            return tempFile;
+            return tempFile.toAbsolutePath();
         } catch (NoSuchAlgorithmException | IOException e) {
             throw new IllegalStateException(e);
         }
     }
 
-    private static String md5Hash(InputStream resourceAsStream) throws IOException, NoSuchAlgorithmException {
+    private static UserPrincipal getCurrentUserPrincipal() throws IOException {
+        Path whoami = Files.createTempFile("whoami", ".tmp");
+        try {
+            return Files.getOwner(whoami);
+        } finally {
+            Files.delete(whoami);
+        }
+    }
+
+    private static String hash(InputStream resourceAsStream) throws IOException, NoSuchAlgorithmException {
         try (InputStream is = resourceAsStream) {
-            MessageDigest md = MessageDigest.getInstance("MD5");
+            MessageDigest md = MessageDigest.getInstance("SHA-256");
             byte[] buffer = new byte[1024];
             DigestInputStream dis = new DigestInputStream(is, md);
             while (dis.read(buffer) != -1) {}
-            return String.format("%032x", new BigInteger(1, md.digest()));
+            return new BigInteger(1, md.digest()).toString(16);
         }
     }
 
-    private static String md5Hash(String s) throws NoSuchAlgorithmException {
-        MessageDigest md = MessageDigest.getInstance("MD5");
+    private static String hash(String s) throws NoSuchAlgorithmException {
+        MessageDigest md = MessageDigest.getInstance("SHA-256");
         md.update(s.getBytes());
-        return String.format("%032x", new BigInteger(1, md.digest()));
+        return new BigInteger(1, md.digest()).toString(16);
     }
 }

apm-agent-common/src/test/java/co/elastic/apm/agent/common/util/ResourceExtractionUtilTest.java

@@ -19,11 +19,13 @@
 package co.elastic.apm.agent.common.util;
 
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
 
-import java.io.File;
 import java.net.URISyntaxException;
+import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.nio.file.attribute.FileTime;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.UUID;
@@ -40,54 +42,57 @@
 class ResourceExtractionUtilTest {
 
     @Test
-    void exportResourceToDirectory() throws URISyntaxException {
-        File tmp = ResourceExtractionUtil.extractResourceToTempDirectory("test.txt", UUID.randomUUID().toString(), ".tmp");
-        tmp.deleteOnExit();
+    void exportResourceToDirectory(@TempDir Path tmpDir) throws URISyntaxException {
+        Path tmp = ResourceExtractionUtil.extractResourceToDirectory("test.txt", "test", ".tmp", tmpDir);
 
         Path referenceFile = Paths.get(ResourceExtractionUtil.class.getResource("/test.txt").toURI());
 
-        assertThat(tmp)
-            .hasSameTextualContentAs(referenceFile.toFile());
+        assertThat(tmp).hasSameTextualContentAs(referenceFile);
     }
 
     @Test
-    void exportResourceToDirectoryIdempotence() throws InterruptedException {
-        String destination = UUID.randomUUID().toString();
-        File tmp = ResourceExtractionUtil.extractResourceToTempDirectory("test.txt", destination, ".tmp");
-        tmp.deleteOnExit();
-        long actual = tmp.lastModified();
+    void exportResourceToDirectoryIdempotence(@TempDir Path tmpDir) throws Exception {
+        Path tmp = ResourceExtractionUtil.extractResourceToDirectory("test.txt", "test", ".tmp", tmpDir);
+        FileTime created = Files.getLastModifiedTime(tmp);
         Thread.sleep(1000);
-        File after = ResourceExtractionUtil.extractResourceToTempDirectory("test.txt", destination, ".tmp");
-        assertThat(actual).isEqualTo(after.lastModified());
+        Path after = ResourceExtractionUtil.extractResourceToDirectory("test.txt", "test", ".tmp", tmpDir);
+        assertThat(created).isEqualTo(Files.getLastModifiedTime(after));
     }
 
     @Test
-    void exportResourceToDirectory_throwExceptionIfNotFound() {
-        assertThatThrownBy(() -> ResourceExtractionUtil.extractResourceToTempDirectory("nonexist", UUID.randomUUID().toString(), ".tmp")).hasMessage("nonexist not found");
+    void testContentDoesNotMatch(@TempDir Path tmpDir) throws Exception {
+        Path tmp = ResourceExtractionUtil.extractResourceToDirectory("test.txt", "test", ".tmp", tmpDir);
+        Files.writeString(tmp, "changed");
+        assertThatThrownBy(() -> ResourceExtractionUtil.extractResourceToDirectory("test.txt", "test", ".tmp", tmpDir))
+            .isInstanceOf(IllegalStateException.class);
     }
 
     @Test
-    void exportResourceToDirectoryInMultipleThreads() throws InterruptedException, ExecutionException {
+    void exportResourceToDirectory_throwExceptionIfNotFound(@TempDir Path tmpDir) {
+        assertThatThrownBy(() -> ResourceExtractionUtil.extractResourceToDirectory("nonexist", "nonexist", ".tmp", tmpDir))
+            .hasMessage("nonexist not found");
+    }
+
+    @Test
+    void exportResourceToDirectoryInMultipleThreads(@TempDir Path tmpDir) throws InterruptedException, ExecutionException {
         final int nbThreads = 10;
         final ExecutorService executorService = Executors.newFixedThreadPool(nbThreads);
         final CountDownLatch countDownLatch = new CountDownLatch(nbThreads);
-        final List<Future<File>> futureList = new ArrayList<>(nbThreads);
+        final List<Future<Path>> futureList = new ArrayList<>(nbThreads);
         final String tempFileNamePrefix = UUID.randomUUID().toString();
 
         for (int i = 0; i < nbThreads; i++) {
             futureList.add(executorService.submit(() -> {
                 countDownLatch.countDown();
                 countDownLatch.await();
-                File file = ResourceExtractionUtil.extractResourceToTempDirectory("test.txt", tempFileNamePrefix, ".tmp");
-                file.deleteOnExit();
-                return file;
+                return ResourceExtractionUtil.extractResourceToDirectory("test.txt", tempFileNamePrefix, ".tmp", tmpDir);
             }));
         }
 
         executorService.shutdown();
         executorService.awaitTermination(1, TimeUnit.SECONDS);
 
-        for (Future<File> future : futureList) {
+        for (Future<Path> future : futureList) {
             assertThat(future.get()).isNotNull();
             assertThat(future.get()).exists();
         }

apm-agent-plugins/apm-profiling-plugin/src/main/java/co/elastic/apm/agent/profiler/asyncprofiler/AsyncProfiler.java

@@ -39,6 +39,8 @@
 import javax.annotation.Nullable;
 import java.io.File;
 import java.io.IOException;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 
 /**
  * Java API for in-process profiling. Serves as a wrapper around
@@ -101,8 +103,8 @@ static void reset() {
 
     private static void loadNativeLibrary(String libraryDirectory) {
         String libraryName = getLibraryFileName();
-        File file = ResourceExtractionUtil.extractResourceToDirectory("asyncprofiler/" + libraryName + ".so", libraryName, ".so", libraryDirectory);
-        System.load(file.getAbsolutePath());
+        Path file = ResourceExtractionUtil.extractResourceToDirectory("asyncprofiler/" + libraryName + ".so", libraryName, ".so", Paths.get(libraryDirectory));
+        System.load(file.toString());
     }
 
     static String getLibraryFileName() {