diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/AgentJarIndex.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/AgentJarIndex.java index 9035465dda3..5350933897c 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/AgentJarIndex.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/AgentJarIndex.java @@ -7,18 +7,17 @@ import java.io.DataOutputStream; import java.io.File; import java.io.IOException; -import java.nio.file.FileVisitResult; import java.nio.file.Files; +import java.nio.file.LinkOption; import java.nio.file.Path; import java.nio.file.Paths; -import java.nio.file.SimpleFileVisitor; -import java.nio.file.attribute.BasicFileAttributes; import java.util.ArrayList; import java.util.Arrays; import java.util.HashSet; import java.util.List; import java.util.Set; import java.util.jar.JarFile; +import java.util.stream.Stream; import java.util.zip.ZipEntry; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -90,7 +89,7 @@ public static AgentJarIndex readIndex(JarFile agentJar) { * Generates an index from the contents of the 'build/resources' directory that makes up the agent * jar. */ - static class IndexGenerator extends SimpleFileVisitor { + static class IndexGenerator { private static final Set ignoredFileNames = new HashSet<>(Arrays.asList("MANIFEST.MF", "NOTICE", "LICENSE.renamed")); @@ -102,9 +101,6 @@ static class IndexGenerator extends SimpleFileVisitor { private final List collectedEntryKeys = new ArrayList<>(); private final List collectedPrefixIds = new ArrayList<>(); - private Path prefixRoot; - private int prefixId; - IndexGenerator(Path resourcesDir) { this.resourcesDir = resourcesDir; @@ -112,7 +108,24 @@ static class IndexGenerator extends SimpleFileVisitor { } void buildIndex() throws IOException { - Files.walkFileTree(resourcesDir, this); + Set seen = new HashSet<>(); + try (Stream paths = Files.walk(resourcesDir)) { + paths + .filter(path -> Files.isRegularFile(path, LinkOption.NOFOLLOW_LINKS)) + .map(resourcesDir::relativize) + .sorted() + .filter(entry -> entry.getNameCount() >= 2) + .forEach( + entry -> { + String prefix = entry.getName(0) + "/"; + int prefixId = prefixIdFor(prefix); + String entryKey = computeEntryKey(entry.subpath(1, entry.getNameCount())); + if (null != entryKey && seen.add(prefixId + "\0" + entryKey)) { + collectedEntryKeys.add(entryKey); + collectedPrefixIds.add(prefixId); + } + }); + } for (int i = 0; i < collectedEntryKeys.size(); i++) { prefixTrie.put(collectedEntryKeys.get(i), collectedPrefixIds.get(i)); @@ -152,42 +165,13 @@ private String getPrefix(int prefixId) { return prefixes.get(prefixId - 1); } - @Override - public FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs) { - if (dir.getParent().equals(resourcesDir)) { - String prefix = dir.getFileName() + "/"; - prefixRoot = dir; - prefixId = 1 + prefixes.indexOf(prefix); - if (prefixId < 1) { - prefixes.add(prefix); - prefixId = prefixes.size(); - } - } - return FileVisitResult.CONTINUE; - } - - @Override - public FileVisitResult postVisitDirectory(Path dir, IOException exc) { - if (dir.equals(prefixRoot)) { - prefixRoot = null; - } - return FileVisitResult.CONTINUE; - } - - @Override - public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) { - if (null != prefixRoot) { - String entryKey = computeEntryKey(prefixRoot.relativize(file)); - if (null != entryKey) { - collectedEntryKeys.add(entryKey); - collectedPrefixIds.add(prefixId); - if (entryKey.endsWith("*")) { - // optimization: wildcard will match everything under here so can skip - return FileVisitResult.SKIP_SIBLINGS; - } - } + private int prefixIdFor(String prefix) { + int prefixId = 1 + prefixes.indexOf(prefix); + if (prefixId < 1) { + prefixes.add(prefix); + prefixId = prefixes.size(); } - return FileVisitResult.CONTINUE; + return prefixId; } static String computeEntryKey(Path path) { diff --git a/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/AgentJarIndexTest.java b/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/AgentJarIndexTest.java index 1af386972e2..7ff3028b9cf 100644 --- a/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/AgentJarIndexTest.java +++ b/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/AgentJarIndexTest.java @@ -1,13 +1,18 @@ package datadog.trace.bootstrap; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import java.io.BufferedInputStream; +import java.io.DataInputStream; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; +import java.util.Arrays; +import java.util.List; import java.util.jar.JarFile; import java.util.zip.ZipEntry; import java.util.zip.ZipOutputStream; @@ -103,6 +108,25 @@ private static void createFile(Path root, String relativePath) throws IOExceptio Files.createFile(file); } + private static Path writeIndex(Path resourcesDir, Path indexFile) throws IOException { + AgentJarIndex.IndexGenerator generator = new AgentJarIndex.IndexGenerator(resourcesDir); + generator.buildIndex(); + generator.writeIndex(indexFile); + return indexFile; + } + + private static List readPrefixes(Path indexFile) throws IOException { + try (DataInputStream in = + new DataInputStream(new BufferedInputStream(Files.newInputStream(indexFile)))) { + int prefixCount = in.readInt(); + String[] prefixes = new String[prefixCount]; + for (int i = 0; i < prefixCount; i++) { + prefixes[i] = in.readUTF(); + } + return Arrays.asList(prefixes); + } + } + @Test void classEntryNameResolvesBootstrapClassToTopLevel(@TempDir Path tempDir) throws Exception { Path resources = tempDir.resolve("resources"); @@ -199,6 +223,52 @@ void multiplePrefixesAreIndexedIndependently(@TempDir Path tempDir) throws Excep index.classEntryName("com.datadoghq.stats.StatsClient")); } + @Test + void buildIndexWritesPrefixesInSortedOrder(@TempDir Path tempDir) throws Exception { + Path resources = tempDir.resolve("resources"); + createFile(resources, "metrics/com/datadoghq/stats/StatsClient.classdata"); + createFile(resources, "inst/datadog/trace/instrumentation/servlet/ServletAdvice.classdata"); + createFile(resources, "appsec/com/datadog/appsec/Event.classdata"); + + Path indexFile = writeIndex(resources, tempDir.resolve("dd-java-agent.index")); + + assertEquals(Arrays.asList("appsec/", "inst/", "metrics/"), readPrefixes(indexFile)); + } + + @Test + void buildIndexIgnoresTopLevelFilesWhenCollectingPrefixes(@TempDir Path tempDir) + throws Exception { + Path resources = tempDir.resolve("resources"); + createFile(resources, "root-resource.properties"); + createFile(resources, "inst/datadog/trace/instrumentation/servlet/ServletAdvice.classdata"); + + Path indexFile = writeIndex(resources, tempDir.resolve("dd-java-agent.index")); + + assertEquals(Arrays.asList("inst/"), readPrefixes(indexFile)); + } + + @Test + void buildIndexIsIndependentOfDirectoryCreationOrder(@TempDir Path tempDir) throws Exception { + Path buildResources = tempDir.resolve("build-resources"); + createFile(buildResources, "appsec/com/datadog/appsec/Event.classdata"); + createFile(buildResources, "ci-visibility/com/datadog/ci/Visibility.classdata"); + createFile(buildResources, "cws-tls/com/datadog/cws/Tls.classdata"); + createFile( + buildResources, "inst/datadog/trace/instrumentation/servlet/ServletAdvice.classdata"); + + Path deployResources = tempDir.resolve("deploy-resources"); + createFile(deployResources, "cws-tls/com/datadog/cws/Tls.classdata"); + createFile(deployResources, "ci-visibility/com/datadog/ci/Visibility.classdata"); + createFile(deployResources, "appsec/com/datadog/appsec/Event.classdata"); + createFile( + deployResources, "inst/datadog/trace/instrumentation/servlet/ServletAdvice.classdata"); + + Path buildIndex = writeIndex(buildResources, tempDir.resolve("build-dd-java-agent.index")); + Path deployIndex = writeIndex(deployResources, tempDir.resolve("deploy-dd-java-agent.index")); + + assertArrayEquals(Files.readAllBytes(buildIndex), Files.readAllBytes(deployIndex)); + } + @Test void buildIndexWithEmptyResourcesDirProducesWorkingIndex(@TempDir Path tempDir) throws Exception { Path resources = tempDir.resolve("resources");