From 5799c5c14e229d6b766d938d9cf864e1ddfe4fa8 Mon Sep 17 00:00:00 2001
From: Nikita Tkachenko
 <121111529+nikita-tkachenko-datadog@users.noreply.github.com>
Date: Mon, 13 Nov 2023 12:55:23 +0100
Subject: [PATCH] Fix error when trying to get class stream for Mockito mocks
 (#6183)

---
 .../agent-ci-visibility/build.gradle          |  1 +
 .../trace/civisibility/source/Utils.java      | 23 +++++++++++++++++--
 .../civisibility/source/index/RepoIndex.java  | 15 ++++--------
 .../RepoIndexSourcePathResolverTest.groovy    | 12 ++++++++++
 4 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/dd-java-agent/agent-ci-visibility/build.gradle b/dd-java-agent/agent-ci-visibility/build.gradle
index 9917f17cbdb..d3d0ad668cb 100644
--- a/dd-java-agent/agent-ci-visibility/build.gradle
+++ b/dd-java-agent/agent-ci-visibility/build.gradle
@@ -73,6 +73,7 @@ excludedClassesCoverage += [
   "datadog.trace.civisibility.source.index.RepoIndexBuilder.RepoIndexingFileVisitor",
   "datadog.trace.civisibility.source.index.RepoIndexFetcher",
   "datadog.trace.civisibility.source.index.RepoIndexSourcePathResolver",
+  "datadog.trace.civisibility.source.Utils",
   "datadog.trace.civisibility.utils.ShellCommandExecutor",
   "datadog.trace.civisibility.utils.ShellCommandExecutor.OutputParser",
   "datadog.trace.civisibility.utils.SpanUtils"
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/Utils.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/Utils.java
index f9023ab5b60..8c826776656 100644
--- a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/Utils.java
+++ b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/Utils.java
@@ -7,7 +7,26 @@ public abstract class Utils {
 
   public static InputStream getClassStream(Class<?> clazz) throws IOException {
     String className = clazz.getName();
-    String classPath = "/" + className.replace('.', '/') + ".class";
-    return clazz.getResourceAsStream(classPath);
+    InputStream classStream = clazz.getResourceAsStream(toResourceName(className));
+    if (classStream != null) {
+      return classStream;
+    } else {
+      // might be auto-generated inner class (e.g. Mockito mock)
+      String topLevelClassName = stripNestedClassNames(clazz.getName());
+      return clazz.getResourceAsStream(toResourceName(topLevelClassName));
+    }
+  }
+
+  private static String toResourceName(String className) {
+    return "/" + className.replace('.', '/') + ".class";
+  }
+
+  public static String stripNestedClassNames(String className) {
+    int innerClassNameIdx = className.indexOf('$');
+    if (innerClassNameIdx >= 0) {
+      return className.substring(0, innerClassNameIdx);
+    } else {
+      return className;
+    }
   }
 }
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/RepoIndex.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/RepoIndex.java
index d8010aeb072..54537b18810 100644
--- a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/RepoIndex.java
+++ b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/source/index/RepoIndex.java
@@ -50,7 +50,7 @@ public List<String> getRootPackages() {
 
   @Nullable
   public String getSourcePath(@Nonnull Class<?> c) {
-    String topLevelClassName = stripNestedClassNames(c.getName());
+    String topLevelClassName = Utils.stripNestedClassNames(c.getName());
     SourceType sourceType = detectSourceType(c);
     String extension = sourceType.getExtension();
     String classNameWithExtension = topLevelClassName + extension;
@@ -107,15 +107,6 @@ private SourceType detectSourceType(Class<?> c) {
     return SourceType.JAVA;
   }
 
-  private String stripNestedClassNames(String className) {
-    int innerClassNameIdx = className.indexOf('$');
-    if (innerClassNameIdx >= 0) {
-      return className.substring(0, innerClassNameIdx);
-    } else {
-      return className;
-    }
-  }
-
   /**
    * Names of package-private classes do not have to correspond to the names of their source code
    * files. For such classes filename is extracted from SourceFile attribute that is available in
@@ -126,6 +117,10 @@ private String getSourcePathForPackagePrivateOrNonJavaClass(Class<?> c) {
       SourceFileAttributeVisitor sourceFileAttributeVisitor = new SourceFileAttributeVisitor();
 
       try (InputStream classStream = Utils.getClassStream(c)) {
+        if (classStream == null) {
+          log.debug("Could not get input stream for class {}", c.getName());
+          return null;
+        }
         ClassReader classReader = new ClassReader(classStream);
         classReader.accept(
             sourceFileAttributeVisitor, ClassReader.SKIP_CODE | ClassReader.SKIP_FRAMES);
diff --git a/dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/source/index/RepoIndexSourcePathResolverTest.groovy b/dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/source/index/RepoIndexSourcePathResolverTest.groovy
index 59a29662254..57b6b0ad229 100644
--- a/dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/source/index/RepoIndexSourcePathResolverTest.groovy
+++ b/dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/source/index/RepoIndexSourcePathResolverTest.groovy
@@ -72,6 +72,17 @@ class RepoIndexSourcePathResolverTest extends Specification {
     sourcePathResolver.getSourcePath(PackagePrivateClass) == expectedSourcePath
   }
 
+  def "test source path resolution for class nested into package-private class"() {
+    setup:
+    def expectedSourcePath = givenSourceFile(RepoIndexSourcePathResolverTest, repoRoot + "/src")
+
+    when:
+    def sourcePathResolver = new RepoIndexSourcePathResolver(repoRoot, packageResolver, fileSystem)
+
+    then:
+    sourcePathResolver.getSourcePath(PackagePrivateClass.NestedIntoPackagePrivateClass) == expectedSourcePath
+  }
+
   def "test source path resolution for non-java class whose file name is different from class name"() {
     setup:
     def expectedSourcePath = givenSourceFile(RepoIndexSourcePathResolverTest, repoRoot + "/src")
@@ -173,6 +184,7 @@ class RepoIndexSourcePathResolverTest extends Specification {
 
 @PackageScope
 class PackagePrivateClass {
+  class NestedIntoPackagePrivateClass {}
 }
 
 class PublicClassWhoseNameDoesNotCorrespondToFileName {