Monkeypatching Java Classes

So, remember that thing I said last post about adding some debug features to JSSE? Turns out that didn't work.

Sometimes debugging would work. Sometimes it wouldn't. I couldn't get it to work reliably.

The debugging feature in JSSE is defined by calling -Djavax.net.debug=ssl on the command line. The class that reads from the system property is sun.security.ssl.Debug and reading it explained everything:

public class Debug {

    private String prefix;
    private static String args;

    static {
        args = java.security.AccessController.doPrivileged(
            new GetPropertyAction("javax.net.debug", ""));
        args = args.toLowerCase(Locale.ENGLISH);
        if (args.equals("help")) {
            Help();
        }
    }

    public static Debug getInstance(String option)
    {
        return getInstance(option, option);
    }

    public static Debug getInstance(String option, String prefix)
    {
        if (isOn(option)) {
            Debug d = new Debug();
            d.prefix = prefix;
            return d;
        } else {
            return null;
        }
    }

    public static boolean isOn(String option)
    {
        if (args == null) {
            return false;
        } else {
            int n = 0;
            option = option.toLowerCase(Locale.ENGLISH);

            if (args.indexOf("all") != -1) {
                return true;
            } else if ((n = args.indexOf("ssl")) != -1) {
                if (args.indexOf("sslctx", n) == -1) {
                    // don't enable data and plaintext options by default
                    if (!(option.equals("data")
                        || option.equals("packet")
                        || option.equals("plaintext"))) {
                        return true;
                    }
                }
            }
            return (args.indexOf(option) != -1);
        }
    }
}

This explained why I was seeing problems with my debug code. The args field of Debug was being set in a static initialization block, when the class was first loaded into the JVM. Due to some race conditions, my code could call System.setProperty("java.net.debug", options) before Debug was loaded, but there was no way of ensuring that. And of course, the args file was marked as private static final so there was no way to modify it outside of the class after that point.

But it was even worse than that. The static methods isOn and getInstance were used by the JSSE internal classes to determine whether debug information should be logged or not, and those fields were also marked private static final, i.e. in sun.security.ssl.SSLContextImpl:

public abstract class SSLContextImpl extends SSLContextSpi {

    private static final Debug debug = Debug.getInstance("ssl");

    private X509ExtendedKeyManager chooseKeyManager(KeyManager[] kms)
            throws KeyManagementException {
        for (int i = 0; kms != null && i < kms.length; i++) {
            if (debug != null && Debug.isOn("sslctx")) {
                System.out.println(
                    "X509KeyManager passed to " +
                    "SSLContext.init():  need an " +
                    "X509ExtendedKeyManager for SSLEngine use");
            }
            return new AbstractKeyManagerWrapper((X509KeyManager)km);
        }
        return DummyX509KeyManager.INSTANCE;
    }
}

In order to turn debugging on, I not only needed to change the args field in sun.security.ssl.Debug after it had been initialized, but I also needed to change every class that used private static final Debug debug = Debug.getInstance("ssl"); so that it would no longer be null.

The first idea I came up with was to change Debug.getInstance() and Debug.isOn. Now, there are ways to change method definitions in the JVM. The package java.lang.instrument defines java agents that can swap out code at run time ("hot reloading"), there's "HotSwap" the JVM debug option, and there's the disposable classloader option. None of them are really applicable in this case – JSSE is a system level package, Play does not require a java agent, and HotSwap is… well, it would probably work fine. But we don't actually need to change the method definition. We just need to swap out private static final field references at run time.

The language tells you that you can't touch final fields and it tells you that you absolutely can't touch private fields from outside the class. It's lying. If you own the JVM, you can monkey patch any field reference.

The key is a method called sun.misc.Unsafe#putObject. It looks like this:

trait MonkeyPatcher {

  private val unsafe: sun.misc.Unsafe = {
    val field = Class.forName("sun.misc.Unsafe").getDeclaredField("theUnsafe")
    field.setAccessible(true)
    field.get(null).asInstanceOf[sun.misc.Unsafe]
  }

  /**
   * Monkeypatches any given static field.
   *
   * @param field the field to change
   * @param newObject the new object to place in the field.
   */
  def monkeyPatchField(field: Field, newObject: AnyRef) {
    val base = unsafe.staticFieldBase(field)
    val offset = unsafe.staticFieldOffset(field)
    unsafe.putObject(base, offset, newObject)
  }
}

So, we've got the means to swap out any given field. Now we need to find all the classes that have a Debug field defined. This is a little trickier, as the class loader only knows about the fields in classes that have already been loaded. We need to go through the JAR file, load all the classes in that JAR into memory, and then quiz them about whether they have that field or not.

So, we add a class finder trait:

trait ClassFinder {
  def logger: org.slf4j.Logger
  def initialResource: String
  def isValidClass(className: String): Boolean

  def findClasses: Set[Class[_]] = {
    logger.debug(s"findClasses: using initialResource = ${initialResource}")

    val classSet = scala.collection.mutable.Set[Class[_]]()
    val classLoader: ClassLoader = Thread.currentThread.getContextClassLoader

    val urlToSource: URL = this.getClass.getResource(initialResource)
    logger.debug(s"findClasses: urlToSource = ${urlToSource}")

    val parts: Array[String] = urlToSource.toString.split("!")
    val jarURLString: String = parts(0).replace("jar:", "")

    logger.debug(s"findClasses: Loading from ${jarURLString}")

    val jar: URL = new URL(jarURLString)
    val jarConnection: URLConnection = jar.openConnection
    val jis: JarInputStream = new JarInputStream(jarConnection.getInputStream)
    try {
      var je: JarEntry = jis.getNextJarEntry
      while (je != null) {
        if (!je.isDirectory) {
          var className: String = je.getName.substring(0, je.getName.length - 6)
          className = className.replace('/', '.')
          if (isValidClass(className)) {
            logger.debug(s"findClasses: adding valid class ${className}")

            val c: Class[_] = classLoader.loadClass(className)
            classSet.add(c)
          }
        }
        je = jis.getNextJarEntry
      }
    } finally {
      jis.close()
    }
    classSet.toSet
  }
}

And then we're going to set up something that will do Debug and args swapping specifically. There are a couple of wrinkles to this: we need to have AccessController give us privileged conditions, and between 1.6 and 1.7 the package name of Debug changed from com.sun.net.ssl.internal.ssl.Debug to sun.security.ssl.Debug, so we have to get the class we want at runtime as well.

abstract class FixLoggingAction extends PrivilegedExceptionAction[Unit] with MonkeyPatcher with ClassFinder {

  def newOptions: String

  def isValidField(field: Field, definedType: Class[_]): Boolean = {
    import java.lang.reflect.Modifier._
    val modifiers: Int = field.getModifiers
    field.getType == definedType && isStatic(modifiers) && isFinal(modifiers)
  }
}

object FixInternalDebugLogging {
  class MonkeyPatchInternalSslDebugAction(val newOptions: String) extends FixLoggingAction {

    def initialResource = "/javax/net/ssl/SSLContext.class"

    def isValidClass(className: String): Boolean = {
      className.startsWith("com.sun.net.ssl.internal.ssl") || className.startsWith("sun.security.ssl")
    }

    def run() {
      val debugType: Class[_] = {
        val debugClassName = foldVersion(
          run16 = "com.sun.net.ssl.internal.ssl.Debug",
          runHigher = "sun.security.ssl.Debug"
        )
        Thread.currentThread().getContextClassLoader.loadClass(debugClassName)
      }

      val newDebug: AnyRef = debugType.newInstance().asInstanceOf[AnyRef]
      
      // Switch out all the classes with a static final field
      for (debugClass <- findClasses) {
        for (debugField <- debugClass.getDeclaredFields) {
          if (isValidField(debugField, debugType)) {
            monkeyPatchField(debugField, newDebug)
          }
        }
      }

      // Switch out the Debug.args field.
      val argsField = debugType.getDeclaredField("args")
      monkeyPatchField(argsField, newOptions)
    }
  }

  def apply(newOptions: String) {
    try {
      val action = new MonkeyPatchInternalSslDebugAction(newOptions)
      AccessController.doPrivileged(action)
    } catch {
      case NonFatal(e) =>
        throw new IllegalStateException("InternalDebug configuration error", e)
    }
  }
}

That's it! Just call FixInternalDebugLogging("ssl") and you can turn on debug information in TLS dynamically, at runtime.

I tried a couple of different things to swap out the Debug class to use a logger instead of System.out.println – this workd in some cases where debug.println was used, but not in others, where System.out.println is used explicitly. There was a patch at one point to fix this, but as of 2014, the only way to do this reliably is to change System.out to a logger. Still, this is tremendously useful when you are in the Play console or running a couple of tests to a specific client.

This is a technique that works best on initialization code where the original codebase just needs to be tweaked a bit, but monkey patching can be applied to any accessible field. It's obviously unsafe, but it's effective. Evil, sure. But effective.

Mua ha ha.

Comments