Added support for advice on constructors.

raphw · Apr 6, 2016 · 96935da · 96935da
1 parent 5163122
commit 96935da
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 9 deletions.
diff --git a/byte-buddy-dep/src/main/java/net/bytebuddy/asm/Advice.java b/byte-buddy-dep/src/main/java/net/bytebuddy/asm/Advice.java
@@ -244,7 +244,7 @@ protected interface MetaDataHandler {
 
         void injectHandlerFrame(MethodVisitor methodVisitor);
 
-        void injectCompletionFrame(MethodVisitor methodVisitor, boolean secondaryCompletion);
+        void injectCompletionFrame(MethodVisitor methodVisitor, boolean secondary);
 
         interface ForInstrumentedMethod extends MetaDataHandler {
 
@@ -311,7 +311,7 @@ public void injectHandlerFrame(MethodVisitor methodVisitor) {
             }
 
             @Override
-            public void injectCompletionFrame(MethodVisitor methodVisitor, boolean secondaryCompletion) {
+            public void injectCompletionFrame(MethodVisitor methodVisitor, boolean secondary) {
                 /* do nothing */
             }
         }
@@ -490,9 +490,9 @@ public void injectHandlerFrame(MethodVisitor methodVisitor) {
             }
 
             @Override
-            public void injectCompletionFrame(MethodVisitor methodVisitor, boolean secondaryCompletion) {
-                if (!expandFrames && currentFrameDivergence == 0) {
-                    if (secondaryCompletion) {
+            public void injectCompletionFrame(MethodVisitor methodVisitor, boolean secondary) {
+                if (!expandFrames && currentFrameDivergence == 0 && (secondary || !instrumentedMethod.isConstructor())) {
+                    if (secondary) {
                         methodVisitor.visitFrame(Opcodes.F_SAME, 0, EMPTY, 0, EMPTY);
                     } else {
                         Object[] local = instrumentedMethod.getReturnType().represents(void.class)
@@ -540,6 +540,7 @@ protected int copy(MethodDescription.InDefinedShape instrumentedMethod,
                         return length;
                     }
                 }, ENTRY {
+
                     @Override
                     protected int copy(MethodDescription.InDefinedShape instrumentedMethod,
                                        MethodDescription.InDefinedShape methodDescription,
@@ -555,6 +556,7 @@ protected int copy(MethodDescription.InDefinedShape instrumentedMethod,
                             translated[index++] = toFrame(typeDescription);
                         }
                         return index;
+
                     }
                 }, EXIT {
                     @Override
@@ -646,9 +648,9 @@ public void injectHandlerFrame(MethodVisitor methodVisitor) {
                 }
 
                 @Override
-                public void injectCompletionFrame(MethodVisitor methodVisitor, boolean secondaryCompletion) {
-                    if (!expandFrames && currentFrameDivergence == 0 && yieldedTypes.size() < 4) {
-                        if (secondaryCompletion || yieldedTypes.isEmpty()) {
+                public void injectCompletionFrame(MethodVisitor methodVisitor, boolean secondary) {
+                    if ((!expandFrames && currentFrameDivergence == 0 && yieldedTypes.size() < 4)) {
+                        if (secondary || yieldedTypes.isEmpty()) {
                             methodVisitor.visitFrame(Opcodes.F_SAME, 0, EMPTY, 0, EMPTY);
                         } else {
                             Object[] local = new Object[yieldedTypes.size()];
@@ -3446,7 +3448,8 @@ public String toString() {
     public @interface OnMethodExit {
 
         /**
-         * Indicates that the advise method should also be called when a method terminates exceptionally.
+         * Indicates that the advise method should also be called when a method terminates exceptionally. This property must
+         * not be set to {@code true} when defining advice for a constructor.
          *
          * @return {@code true} if the advise method should be invoked when a method terminates exceptionally.
          */

diff --git a/byte-buddy-dep/src/test/java/net/bytebuddy/asm/AdviceTest.java b/byte-buddy-dep/src/test/java/net/bytebuddy/asm/AdviceTest.java
@@ -21,6 +21,7 @@
 import java.util.Iterator;
 
 import static junit.framework.TestCase.fail;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static org.hamcrest.CoreMatchers.*;
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -202,6 +203,19 @@ public void testEmptyAdviceExitWithExceptionHandlingAndSuppression() throws Exce
         assertThat(type.getDeclaredMethod(FOO).invoke(type.newInstance()), nullValue(Object.class));
     }
 
+    @Test
+    public void testAdviceOnConstructor() throws Exception {
+        Class<?> type = new ByteBuddy() // TODO: Exception when constructor for certain types and throwable
+                .redefine(Sample.class)
+                .visit(Advice.to(TrivialAdviceSkipException.class).on(isConstructor()))
+                .make()
+                .load(null, ClassLoadingStrategy.Default.WRAPPER)
+                .getLoaded();
+        assertThat(type.newInstance(), notNullValue(Object.class));
+        assertThat(type.getDeclaredField(ENTER).get(null), is((Object) 1));
+        assertThat(type.getDeclaredField(EXIT).get(null), is((Object) 1));
+    }
+
     @Test
     public void testAdviceWithImplicitArgument() throws Exception {
         Class<?> type = new ByteBuddy()