From b9b0b5d19f4db0bbdcfe54582b703985ced37aa1 Mon Sep 17 00:00:00 2001
From: Uwe Schindler <uwe@thetaphi.de>
Date: Sun, 27 Dec 2015 19:06:40 +0100
Subject: [PATCH] Issue #91: Get order of loading from classpath right (do not
 prefer classes to check)

---
 .../de/thetaphi/forbiddenapis/Checker.java    | 23 +++++++++++++------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/src/main/java/de/thetaphi/forbiddenapis/Checker.java b/src/main/java/de/thetaphi/forbiddenapis/Checker.java
index 146a2504..bb76230a 100644
--- a/src/main/java/de/thetaphi/forbiddenapis/Checker.java
+++ b/src/main/java/de/thetaphi/forbiddenapis/Checker.java
@@ -75,7 +75,7 @@ public static enum Option {
   final java.lang.reflect.Method method_Class_getModule, method_Module_getResourceAsStream;
   final EnumSet<Option> options;
   
-  // key is the internal name (slashed):
+  // key is the binary name (dotted):
   final Map<String,ClassSignature> classesToCheck = new HashMap<String,ClassSignature>();
   // key is the binary name (dotted):
   final Map<String,ClassSignature> classpathClassCache = new HashMap<String,ClassSignature>();
@@ -265,7 +265,7 @@ private boolean isRuntimeClass(URLConnection conn) throws IOException {
     return false;
   }
   
-  /** Reads a class (binary name) from the given {@link ClassLoader}. */
+  /** Reads a class (binary name) from the given {@link ClassLoader}. If not found there, falls back to the list of classes to be checked. */
   private ClassSignature getClassFromClassLoader(final String clazz) throws ClassNotFoundException,IOException {
     final ClassSignature c;
     if (classpathClassCache.containsKey(clazz)) {
@@ -298,6 +298,12 @@ private ClassSignature getClassFromClassLoader(final String clazz) throws ClassN
           return c;
         }
       }
+      // try to get class from our list of classes we are checking:
+      c = classesToCheck.get(clazz);
+      if (c != null) {
+        classpathClassCache.put(clazz, c);
+        return c;
+      }
       // all failed => the class does not exist!
       classpathClassCache.put(clazz, null);
       throw new ClassNotFoundException(clazz);
@@ -310,10 +316,9 @@ public ClassSignature lookupRelatedClass(String internalName) {
     if (type.getSort() != Type.OBJECT) {
       return null;
     }
-    ClassSignature c = classesToCheck.get(internalName);
-    if (c == null) try {
+    try {
       // use binary name, so we need to convert:
-      c = getClassFromClassLoader(type.getClassName());
+      return getClassFromClassLoader(type.getClassName());
     } catch (ClassNotFoundException cnfe) {
       if (options.contains(Option.FAIL_ON_MISSING_CLASSES)) {
         throw new WrapperRuntimeException(cnfe);
@@ -322,11 +327,11 @@ public ClassSignature lookupRelatedClass(String internalName) {
           "The referenced class '%s' cannot be loaded. Please fix the classpath!",
           type.getClassName()
         ));
+        return null;
       }
     } catch (IOException ioe) {
       throw new WrapperRuntimeException(ioe);
     }
-    return c;
   }
   
   /** Adds the method signature to the list of disallowed methods. The Signature is checked against the given ClassLoader. */
@@ -508,7 +513,8 @@ public void addClassToCheck(final InputStream in) throws IOException {
     } finally {
       in.close();
     }
-    classesToCheck.put(reader.getClassName(), new ClassSignature(reader, false, true));
+    final String binaryName = Type.getObjectType(reader.getClassName()).getClassName();
+    classesToCheck.put(binaryName, new ClassSignature(reader, false, true));
   }
   
   /** Parses and adds a class from the given file to the list of classes to check. Does not log anything. */
@@ -588,6 +594,9 @@ public void run() throws ForbiddenApiException {
       }
     }
     
+    // Cleanup cache to get statistics right:
+    classpathClassCache.keySet().removeAll(classesToCheck.keySet());
+    
     final String message = String.format(Locale.ENGLISH, 
         "Scanned %d (and %d related) class file(s) for forbidden API invocations (in %.2fs), %d error(s).",
         classesToCheck.size(), classesToCheck.isEmpty() ? 0 : classpathClassCache.size(), (System.currentTimeMillis() - start) / 1000.0, errors);