Make substitutable class file transformers unwrappable and unwrap bef…

…ore applying a differential matcher.
raphw · Aug 13, 2023 · 0f1d715 · 0f1d715
1 parent 5e3f7fc
commit 0f1d715
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 12 deletions.
diff --git a/byte-buddy-dep/src/main/java/net/bytebuddy/agent/builder/AgentBuilder.java b/byte-buddy-dep/src/main/java/net/bytebuddy/agent/builder/AgentBuilder.java
@@ -11337,7 +11337,9 @@ public ResettableClassFileTransformer patchOn(Instrumentation instrumentation, R
  * {@inheritDoc}
  */
  public ResettableClassFileTransformer patchOn(Instrumentation instrumentation, ResettableClassFileTransformer classFileTransformer, PatchMode patchMode) {
- return doInstall(instrumentation, new Transformation.DifferentialMatcher(ignoreMatcher, transformations, classFileTransformer), patchMode.toHandler(classFileTransformer));
+ return doInstall(instrumentation, new Transformation.DifferentialMatcher(ignoreMatcher, transformations, classFileTransformer instanceof ResettableClassFileTransformer.Substitutable
+ ? ((ResettableClassFileTransformer.Substitutable) classFileTransformer).unwrap()
+ : classFileTransformer), patchMode.toHandler(classFileTransformer));
  }
 
  /**

diff --git a/byte-buddy-dep/src/main/java/net/bytebuddy/agent/builder/ResettableClassFileTransformer.java b/byte-buddy-dep/src/main/java/net/bytebuddy/agent/builder/ResettableClassFileTransformer.java
@@ -298,6 +298,13 @@ interface Substitutable extends ResettableClassFileTransformer {
  * @param classFileTransformer The class file transformer to use.
  */
  void substitute(ResettableClassFileTransformer classFileTransformer);
+
+ /**
+ * Returns the underlying non-substitutable class file transformer.
+ *
+ * @return The underlying non-substitutable class file transformer.
+ */
+ ResettableClassFileTransformer unwrap();
  }
 
  /**
@@ -489,27 +496,25 @@ private static <T> T doPrivileged(PrivilegedAction<T> action) {
  * @param classFileTransformer The class file transformer to wrap.
  * @return A substitutable version of the supplied class file transformer.
  */
- public static ResettableClassFileTransformer.Substitutable of(ResettableClassFileTransformer classFileTransformer) {
+ public static Substitutable of(ResettableClassFileTransformer classFileTransformer) {
  return DISPATCHER.make(classFileTransformer);
  }
 
  /**
- * Returns the delegate class file transformer.
- *
- * @return The delegate class file transformer.
+ * {@inheritDoc}
  */
- protected ResettableClassFileTransformer getClassFileTransformer() {
- return classFileTransformer;
+ public void substitute(ResettableClassFileTransformer classFileTransformer) {
+ while (classFileTransformer instanceof Substitutable) {
+ classFileTransformer = ((Substitutable) classFileTransformer).unwrap();
+ }
+ this.classFileTransformer = classFileTransformer;
  }
 
  /**
  * {@inheritDoc}
  */
- public void substitute(ResettableClassFileTransformer classFileTransformer) {
- while (classFileTransformer instanceof WithDelegation.Substitutable) {
- classFileTransformer = ((WithDelegation.Substitutable) classFileTransformer).getClassFileTransformer();
- }
- this.classFileTransformer = classFileTransformer;
+ public ResettableClassFileTransformer unwrap() {
+ return classFileTransformer;
  }
 
  /**