Skip to content

Commit

Permalink
Add method invocation step for member substitution chain.
Browse files Browse the repository at this point in the history
  • Loading branch information
@raphw
raphw committed Feb 4, 2023
1 parent db80456 commit 58a88fa
Show file tree
Hide file tree
Showing 2 changed files with 166 additions and 10 deletions.
140 changes: 130 additions & 10 deletions byte-buddy-dep/src/main/java/net/bytebuddy/asm/MemberSubstitution.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,7 @@
import net.bytebuddy.dynamic.ClassFileLocator;
import net.bytebuddy.dynamic.scaffold.MethodGraph;
import net.bytebuddy.implementation.Implementation;
import net.bytebuddy.implementation.bytecode.Duplication;
import net.bytebuddy.implementation.bytecode.Removal;
import net.bytebuddy.implementation.bytecode.StackManipulation;
import net.bytebuddy.implementation.bytecode.StackSize;
import net.bytebuddy.implementation.bytecode.*;
import net.bytebuddy.implementation.bytecode.assign.Assigner;
import net.bytebuddy.implementation.bytecode.constant.*;
import net.bytebuddy.implementation.bytecode.member.FieldAccess;
Expand Down Expand Up @@ -1936,13 +1933,8 @@ public Resolution resolve(TypeDescription targetType,
if (index >= parameters.size()) {
throw new IllegalStateException(target + " has not " + index + " arguments");
}
StackManipulation stackManipulation = assigner.assign(parameters.get(index), current, typing);
if (!stackManipulation.isValid()) {
throw new IllegalStateException("Cannot assign " + parameters.get(index) + " to " + current);
}
return new Simple(new StackManipulation.Compound(Removal.of(current),
MethodVariableAccess.of(parameters.get(index)).loadFrom(offsets.get(index)),
stackManipulation), current);
MethodVariableAccess.of(parameters.get(index)).loadFrom(offsets.get(index))), parameters.get(index));
}


Expand Down Expand Up @@ -1972,6 +1964,134 @@ public Step make(Assigner assigner,
}
}
}

/**
* A step for invoking a method or constructor. If non-static, a method is invoked upon a the current stack argument of the chain.
* Arguments are loaded from the intercepted byte code element with a possibility of substitution.
*/
@HashCodeAndEqualsPlugin.Enhance
class ForInvocation implements Step {

/**
* The invoked method or constructor.
*/
private final MethodDescription methodDescription;

/**
* A mapping of substituted indices.
*/
private final Map<Integer, Integer> substitutions;

/**
* The assigner to use.
*/
private final Assigner assigner;

/**
* The typing to use when assigning.
*/
private final Assigner.Typing typing;

/**
* Creates a new step of an invocation.
*
* @param methodDescription The invoked method or constructor.
* @param substitutions A mapping of substituted indices.
* @param assigner The assigner to use.
* @param typing The typing to use when assigning.
*/
public ForInvocation(MethodDescription methodDescription, Map<Integer, Integer> substitutions, Assigner assigner, Assigner.Typing typing) {
this.methodDescription = methodDescription;
this.substitutions = substitutions;
this.assigner = assigner;
this.typing = typing;
}

/**
* {@inheritDoc}
*/
public Resolution resolve(TypeDescription targetType,
ByteCodeElement target,
TypeList.Generic parameters,
TypeDescription.Generic current,
Map<Integer, Integer> offsets,
int freeOffset) {
List<StackManipulation> stackManipulations = new ArrayList<StackManipulation>(3 + parameters.size() * 2);
if (methodDescription.isStatic()) {
stackManipulations.add(Removal.of(current));
} else if (methodDescription.isConstructor()) {
stackManipulations.add(Removal.of(current));
stackManipulations.add(TypeCreation.of(methodDescription.getDeclaringType().asErasure()));
} else {
StackManipulation assignment = assigner.assign(current, methodDescription.getDeclaringType().asGenericType(), typing);
if (!assignment.isValid()) {
throw new IllegalStateException("Cannot assign " + current + " to " + methodDescription.getDeclaringType());
}
stackManipulations.add(assignment);
}
for (int index = 0; index < methodDescription.getParameters().size(); index++) {
Integer substitution = substitutions.getOrDefault(index, index + (methodDescription.isStatic() ? 0 : 1));
if (substitution >= parameters.size()) {
throw new IllegalStateException(target + " does not support an index " + substitution);
}
stackManipulations.add(MethodVariableAccess.of(parameters.get(substitution)).loadFrom(offsets.get(substitution)));
StackManipulation assignment = assigner.assign(parameters.get(substitution), methodDescription.getParameters().get(index).getType(), typing);
if (!assignment.isValid()) {
throw new IllegalStateException("Cannot assign parameter with " + index + " of type " + parameters.get(substitution) + " to " + methodDescription);
}
stackManipulations.add(assignment);
}
stackManipulations.add(MethodInvocation.invoke(methodDescription));
return new Simple(new StackManipulation.Compound(stackManipulations), methodDescription.getReturnType());
}

/**
* A factory to create a step for a method invocation.
*/
@HashCodeAndEqualsPlugin.Enhance
public static class Factory implements Step.Factory {

/**
* The invoked method or constructor.
*/
private final MethodDescription methodDescription;

/**
* A mapping of substituted parameter indices.
*/
private final Map<Integer, Integer> substitutions;

/**
* Creates a factory for a method invocation without parameter substitutions.
*
* @param methodDescription The invoked method or constructor.
*/
public Factory(MethodDescription methodDescription) {
this(methodDescription, Collections.<Integer, Integer>emptyMap());
}

/**
* Creates a factory for a method invocation.
*
* @param methodDescription The invoked method or constructor.
* @param substitutions A mapping of substituted parameter indices.
*/
public Factory(MethodDescription methodDescription, Map<Integer, Integer> substitutions) {
this.methodDescription = methodDescription;
this.substitutions = substitutions;
}

/**
* {@inheritDoc}
*/
public Step make(Assigner assigner,
Assigner.Typing typing,
TypeDescription instrumentedType,
MethodDescription instrumentedMethod) {
return new ForInvocation(methodDescription, substitutions, assigner, typing);
}
}
}
}

/**
Expand Down
36 changes: 36 additions & 0 deletions byte-buddy-dep/src/test/java/net/bytebuddy/asm/MemberSubstitutionTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
import net.bytebuddy.implementation.FieldAccessor;
import net.bytebuddy.implementation.bytecode.constant.NullConstant;
import net.bytebuddy.implementation.bytecode.member.FieldAccess;
import net.bytebuddy.pool.TypePool;
import net.bytebuddy.test.packaging.MemberSubstitutionTestHelper;
import org.junit.Test;
Expand Down Expand Up @@ -774,6 +775,41 @@ public void testSubstitutionChainMethodInvocationOriginal() throws Exception {
assertThat(type.getDeclaredField(BAR).get(instance), is((Object) FOO));
}

@Test
public void testSubstitutionChainMethodInvocation() throws Exception {
Class<?> type = new ByteBuddy()
.redefine(FieldAccessSample.class)
.visit(MemberSubstitution.strict().field(named(FOO)).replaceWithChain(
new MemberSubstitution.Substitution.Chain.Step.ForArgumentLoading.Factory(0),
new MemberSubstitution.Substitution.Chain.Step.ForInvocation.Factory(new MethodDescription.ForLoadedMethod(FieldAccessSample.class.getDeclaredMethod("baz")))).on(named(RUN)))
.make()
.load(ClassLoadingStrategy.BOOTSTRAP_LOADER, ClassLoadingStrategy.Default.WRAPPER)
.getLoaded();
Object instance = type.getDeclaredConstructor().newInstance();
assertThat(type.getDeclaredField(FOO).get(instance), is((Object) FOO));
assertThat(type.getDeclaredField(BAR).get(instance), is((Object) BAR));
assertThat(type.getDeclaredMethod(RUN).invoke(instance), nullValue(Object.class));
assertThat(type.getDeclaredField(FOO).get(instance), is((Object) FOO));
assertThat(type.getDeclaredField(BAR).get(instance), is((Object) BAZ));
}

@Test
public void testSubstitutionChainStaticMethodInvocation() throws Exception {
Class<?> type = new ByteBuddy()
.redefine(StaticFieldAccessSample.class)
.visit(MemberSubstitution.strict().field(named(FOO)).replaceWithChain(
new MemberSubstitution.Substitution.Chain.Step.ForInvocation.Factory(new MethodDescription.ForLoadedMethod(StaticFieldAccessSample.class.getDeclaredMethod("baz")))).on(named(RUN)))
.make()
.load(ClassLoadingStrategy.BOOTSTRAP_LOADER, ClassLoadingStrategy.Default.WRAPPER)
.getLoaded();
Object instance = type.getDeclaredConstructor().newInstance();
assertThat(type.getDeclaredField(FOO).get(instance), is((Object) FOO));
assertThat(type.getDeclaredField(BAR).get(instance), is((Object) BAR));
assertThat(type.getDeclaredMethod(RUN).invoke(instance), nullValue(Object.class));
assertThat(type.getDeclaredField(FOO).get(instance), is((Object) FOO));
assertThat(type.getDeclaredField(BAR).get(instance), is((Object) BAZ));
}

@Test(expected = IllegalStateException.class)
public void testFieldNotAccessible() throws Exception {
new ByteBuddy()
Expand Down

0 comments on commit 58a88fa

Please sign in to comment.