Handle constructors properly within substitution chain.

raphw · Feb 8, 2023 · 94cf9a0 · 94cf9a0
1 parent abf8877
commit 94cf9a0
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 5 deletions.
diff --git a/byte-buddy-dep/src/main/java/net/bytebuddy/asm/MemberSubstitution.java b/byte-buddy-dep/src/main/java/net/bytebuddy/asm/MemberSubstitution.java
@@ -1470,7 +1470,11 @@ public StackManipulation resolve(TypeDescription targetType,
                     stackManipulations.add(resolution.getStackManipulation());
                     current = resolution.getResultType();
                 }
-                stackManipulations.add(assigner.assign(current, result, typing));
+                StackManipulation assignment = assigner.assign(current, result, typing);
+                if (!assignment.isValid()) {
+                    throw new IllegalStateException("Failed to assign " + current + " to " + result);
+                }
+                stackManipulations.add(assignment);
                 return new StackManipulation.Compound(stackManipulations);
             }
 
@@ -1556,14 +1560,24 @@ public Resolution resolve(TypeDescription targetType,
                                               TypeDescription.Generic current,
                                               Map<Integer, Integer> offsets,
                                               int freeOffset) {
-                        List<StackManipulation> stackManipulations = new ArrayList<StackManipulation>(parameters.size() + 2);
-                        stackManipulations.add(Removal.of(current));
+                        List<StackManipulation> stackManipulations;
+                        if (target instanceof MethodDescription && ((MethodDescription) target).isConstructor()) {
+                            stackManipulations = new ArrayList<StackManipulation>(parameters.size() + 4);
+                            stackManipulations.add(Removal.of(current));
+                            stackManipulations.add(TypeCreation.of(((MethodDescription) target).getDeclaringType().asErasure()));
+                            stackManipulations.add(Duplication.SINGLE);
+                        } else {
+                            stackManipulations = new ArrayList<StackManipulation>(parameters.size() + 4);
+                            stackManipulations.add(Removal.of(current));
+                        }
                         for (int index = 0; index < parameters.size(); index++) {
                             stackManipulations.add(MethodVariableAccess.of(parameters.get(index)).loadFrom(offsets.get(index)));
                         }
                         if (target instanceof MethodDescription) {
                             stackManipulations.add(MethodInvocation.invoke((MethodDescription) target));
-                            return new Simple(new StackManipulation.Compound(stackManipulations), ((MethodDescription) target).getReturnType());
+                            return new Simple(new StackManipulation.Compound(stackManipulations), ((MethodDescription) target).isConstructor()
+                                    ? ((MethodDescription) target).getDeclaringType().asGenericType()
+                                    : ((MethodDescription) target).getReturnType());
                         } else if (target instanceof FieldDescription) {
                             if (((FieldDescription) target).isStatic()) {
                                 if (parameters.isEmpty()) {
@@ -1818,7 +1832,7 @@ class ForArgumentLoading implements Step, Factory {
                     /**
                      * Creates an argument loading step.
                      *
-                     * @param index    The index of the argument to load.
+                     * @param index The index of the argument to load.
                      */
                     protected ForArgumentLoading(int index) {
                         this.index = index;

diff --git a/byte-buddy-dep/src/test/java/net/bytebuddy/asm/MemberSubstitutionTest.java b/byte-buddy-dep/src/test/java/net/bytebuddy/asm/MemberSubstitutionTest.java
@@ -776,6 +776,18 @@ public void testSubstitutionChainMethodInvocationOriginal() throws Exception {
         assertThat(type.getDeclaredField(BAR).get(instance), is((Object) FOO));
     }
 
+    @Test
+    public void testSubstitutionChainConstructionOriginal() throws Exception {
+        Class<?> type = new ByteBuddy()
+                .redefine(ConstructorSubstitutionSample.class)
+                .visit(MemberSubstitution.strict().constructor(isDeclaredBy(Object.class)).replaceWithChain(MemberSubstitution.Substitution.Chain.Step.OfOriginalExpression.INSTANCE).on(named(RUN)))
+                .make()
+                .load(ClassLoadingStrategy.BOOTSTRAP_LOADER, ClassLoadingStrategy.Default.WRAPPER)
+                .getLoaded();
+        Object instance = type.getDeclaredConstructor().newInstance();
+        assertThat(type.getDeclaredMethod(RUN).invoke(instance), instanceOf(Object.class));
+    }
+
     @Test
     public void testSubstitutionChainFieldRead() throws Exception {
         Class<?> type = new ByteBuddy()