Fix package resolution for non-Java source files (#7356)

nikita-tkachenko-datadog · web-flow · commit 4e62826c5c86 · 2024-07-26T11:46:52.000+02:00
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/PackageResolverImpl.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/PackageResolverImpl.java
@@ -8,9 +8,13 @@
 import java.nio.file.InvalidPathException;
 import java.nio.file.Path;
 import javax.annotation.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class PackageResolverImpl implements PackageResolver {
 
+  private static final Logger log = LoggerFactory.getLogger(PackageResolverImpl.class);
+
   private static final String PACKAGE_KEYWORD = "package";
   private final FileSystem fileSystem;
 
@@ -33,6 +37,7 @@ public PackageResolverImpl(FileSystem fileSystem) {
   @Nullable
   @Override
   public Path getPackage(Path sourceFile) throws IOException {
+    Language language = Language.getByFileName(sourceFile.getFileName().toString());
     Path folder = sourceFile.getParent();
     try (BufferedReader br = Files.newBufferedReader(sourceFile)) {
       String line;
@@ -42,6 +47,13 @@ public Path getPackage(Path sourceFile) throws IOException {
           continue;
         }
 
+        int charAfterPackageKeyword = packageDeclarationStart + PACKAGE_KEYWORD.length();
+        if (charAfterPackageKeyword >= line.length()
+            || !Character.isWhitespace(line.charAt(charAfterPackageKeyword))) {
+          // "package" keyword is not followed by a whitespace
+          continue;
+        }
+
         int lineLength = line.length();
         int packageNameStart = packageDeclarationStart + PACKAGE_KEYWORD.length();
         while (packageNameStart < lineLength
@@ -51,18 +63,21 @@ public Path getPackage(Path sourceFile) throws IOException {
 
         int packageNameEnd = line.indexOf(';', packageNameStart);
         if (packageNameEnd == -1) {
-          packageNameEnd = lineLength; // no ';' is possible if this is a Groovy file
+          packageNameEnd = lineLength; // possible if this is a non-Java (e.g. Groovy, Scala) file
         }
 
         String packageName = line.substring(packageNameStart, packageNameEnd);
         Path packagePath;
         try {
           packagePath = fileSystem.getPath(packageName.replace('.', File.separatorChar));
         } catch (InvalidPathException e) {
+          log.debug("Invalid package {} found for source file {}", packageName, sourceFile, e);
           continue;
         }
 
-        if (folder.endsWith(packagePath)) {
+        // we only do the "sanity check" for Java, as with the other languages
+        // it is possible to have package that does not correspond to folder
+        if (language != Language.JAVA || folder.endsWith(packagePath)) {
           return packagePath;
         }
       }
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/RepoIndexBuilder.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/RepoIndexBuilder.java
@@ -157,7 +157,7 @@ public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) {
         }
 
         Path sourceRoot =
-            language.isNonCode() ? getNonCodeSourceRoot(file) : getCodeSourceRoot(file);
+            language.isNonCode() ? getNonCodeSourceRoot(file) : getCodeSourceRoot(language, file);
         if (sourceRoot != null) {
           String relativeSourceRoot = repoRoot.relativize(sourceRoot).toString();
           int sourceRootIdx =
@@ -177,19 +177,27 @@ public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) {
       return FileVisitResult.CONTINUE;
     }
 
-    private Path getCodeSourceRoot(Path file) throws IOException {
+    private Path getCodeSourceRoot(Language language, Path file) throws IOException {
       indexingStats.sourceFilesVisited++;
       Path packagePath = packageResolver.getPackage(file);
       if (packagePath != null) {
         packageTree.add(packagePath);
 
         Path folder = file.getParent();
-        // remove package path suffix from folder path to get source root
-        return folder
-            .getRoot()
-            .resolve(folder.subpath(0, folder.getNameCount() - packagePath.getNameCount()));
+        if (folder.endsWith(packagePath)) {
+          // In non-JVM languages package names do not have to correspond to folder structure,
+          // so using package to find source root is not always possible
+          return folder
+              .getRoot()
+              .resolve(folder.subpath(0, folder.getNameCount() - packagePath.getNameCount()));
+        }
+      }
+
+      if (language != Language.JAVA) {
+        // Fallback for non-JVM languages
+        return resourceResolver.getResourceRoot(file);
       } else {
-        // assuming default package
+        // For Java assuming default package
         return file.getParent();
       }
     }
