Avoid that type initializer extension adjusts exception handlers by prepending new handler blocks.

Author: raphw

byte-buddy-dep/src/main/java/net/bytebuddy/dynamic/scaffold/TypeWriter.java

@@ -70,6 +70,7 @@
 import net.bytebuddy.utility.nullability.UnknownNull;
 import net.bytebuddy.utility.privilege.GetSystemPropertyAction;
 import net.bytebuddy.utility.visitor.ContextClassVisitor;
+import net.bytebuddy.utility.visitor.ExceptionTableSensitiveMethodVisitor;
 import net.bytebuddy.utility.visitor.MetadataAwareClassVisitor;
 import org.objectweb.asm.AnnotationVisitor;
 import org.objectweb.asm.ClassReader;
@@ -4465,7 +4466,7 @@ public void complete(ClassVisitor classVisitor, Implementation.Context.Extractab
                     /**
                      * An initialization handler that appends code to a previously visited type initializer.
                      */
-                    abstract class Appending extends MethodVisitor implements InitializationHandler, TypeInitializer.Drain {
+                    abstract class Appending extends ExceptionTableSensitiveMethodVisitor implements InitializationHandler, TypeInitializer.Drain {
 
                         /**
                          * The instrumented type.
@@ -4600,6 +4601,10 @@ private static WithoutDrain withoutDrain(MethodVisitor methodVisitor,
                         public void visitCode() {
                             record.applyAttributes(mv, annotationValueFilterFactory);
                             super.visitCode();
+                        }
+
+                        @Override
+                        protected void onAfterExceptionTable() {
                             onStart();
                         }
 
@@ -4609,8 +4614,8 @@ public void visitCode() {
                         protected abstract void onStart();
 
                         @Override
-                        public void visitFrame(int type, int localVariableLength, @MaybeNull Object[] localVariable, int stackSize, @MaybeNull Object[] stack) {
-                            super.visitFrame(type, localVariableLength, localVariable, stackSize, stack);
+                        protected void onVisitFrame(int type, int localVariableLength, @MaybeNull Object[] localVariable, int stackSize, @MaybeNull Object[] stack) {
+                            super.onVisitFrame(type, localVariableLength, localVariable, stackSize, stack);
                             frameWriter.onFrame(type, localVariableLength);
                         }
 
@@ -4869,11 +4874,11 @@ protected WithActiveRecord(MethodVisitor methodVisitor,
                                 }
 
                                 @Override
-                                public void visitInsn(int opcode) {
+                                protected void onVisitInsn(int opcode) {
                                     if (opcode == Opcodes.RETURN) {
                                         mv.visitJumpInsn(Opcodes.GOTO, label);
                                     } else {
-                                        super.visitInsn(opcode);
+                                        super.onVisitInsn(opcode);
                                     }
                                 }
 
@@ -5013,11 +5018,11 @@ protected WithActiveRecord(MethodVisitor methodVisitor,
                                 }
 
                                 @Override
-                                public void visitInsn(int opcode) {
+                                protected void onVisitInsn(int opcode) {
                                     if (opcode == Opcodes.RETURN) {
                                         mv.visitJumpInsn(Opcodes.GOTO, label);
                                     } else {
-                                        super.visitInsn(opcode);
+                                        super.onVisitInsn(opcode);
                                     }
                                 }
 

byte-buddy-dep/src/test/java/net/bytebuddy/asm/AdviceTest.java

@@ -2428,6 +2428,28 @@ public void testSelfCallHandleNotAssignable() throws Exception {
         Advice.to(SelfCallHandleIllegalSample.class);
     }
 
+    @Test
+    public void testTypeInitializerExceptionHandlingRetainedOnRebase() throws Exception {
+        Class<?> type = new ByteBuddy()
+                .rebase(TypeInitializerWithExceptionHandlingSample.class)
+                .visit(Advice.to(TypeInitializerWithExceptionHandlingAdvice.class).on(ElementMatchers.isTypeInitializer()))
+                .make()
+                .load(ClassLoadingStrategy.BOOTSTRAP_LOADER, ClassLoadingStrategy.Default.WRAPPER)
+                .getLoaded();
+        assertThat(type.getDeclaredField("handled").get(null), is((Object) true));
+    }
+
+    @Test
+    public void testTypeInitializerExceptionHandlingRetainedOnRedefine() throws Exception {
+        Class<?> type = new ByteBuddy()
+                .redefine(TypeInitializerWithExceptionHandlingSample.class)
+                .visit(Advice.to(TypeInitializerWithExceptionHandlingAdvice.class).on(ElementMatchers.isTypeInitializer()))
+                .make()
+                .load(ClassLoadingStrategy.BOOTSTRAP_LOADER, ClassLoadingStrategy.Default.WRAPPER)
+                .getLoaded();
+        assertThat(type.getDeclaredField("handled").get(null), is((Object) true));
+    }
+
     @SuppressWarnings("unused")
     public static class Sample {
 
@@ -4341,4 +4363,27 @@ static void exit() {
             /* do nothing */
         }
     }
+
+    public static class TypeInitializerWithExceptionHandlingSample {
+
+        public static boolean handled = false;
+
+        static {
+            try {
+                throw new Exception();
+            } catch (Exception e) {
+                handled = true;
+            }
+        }
+    }
+
+    public static class TypeInitializerWithExceptionHandlingAdvice {
+
+        @Advice.OnMethodExit(onThrowable = Throwable.class, inline = true)
+        static void exit(
+                @Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) Object returning,
+                @Advice.Thrown(readOnly = false, typing = Assigner.Typing.DYNAMIC) Throwable throwing
+        ) {
+        }
+    }
 }