Skip to content

Commit

Permalink
Add ability to allow list instance methods on the script class (#76045)…
Browse files Browse the repository at this point in the history
… (#76047)

This change adds support to allow list instance methods on the script class in Painless and super 
classes of the script class. It works the same way as allow listing works now for other classes.
  • Loading branch information
elasticsearchmachine authored Aug 3, 2021
1 parent 362a7d4 commit 34829e8
Show file tree
Hide file tree
Showing 7 changed files with 206 additions and 31 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,7 @@
import org.elasticsearch.painless.symbol.IRDecorations.IRDSize;
import org.elasticsearch.painless.symbol.IRDecorations.IRDStoreType;
import org.elasticsearch.painless.symbol.IRDecorations.IRDSymbol;
import org.elasticsearch.painless.symbol.IRDecorations.IRDThisMethod;
import org.elasticsearch.painless.symbol.IRDecorations.IRDTypeParameters;
import org.elasticsearch.painless.symbol.IRDecorations.IRDUnaryType;
import org.elasticsearch.painless.symbol.IRDecorations.IRDValue;
Expand Down Expand Up @@ -1650,6 +1651,7 @@ public void visitInvokeCallMember(InvokeCallMemberNode irInvokeCallMemberNode, W
methodWriter.writeDebugInfo(irInvokeCallMemberNode.getLocation());

LocalFunction localFunction = irInvokeCallMemberNode.getDecorationValue(IRDFunction.class);
PainlessMethod thisMethod = irInvokeCallMemberNode.getDecorationValue(IRDThisMethod.class);
PainlessMethod importedMethod = irInvokeCallMemberNode.getDecorationValue(IRDMethod.class);
PainlessClassBinding classBinding = irInvokeCallMemberNode.getDecorationValue(IRDClassBinding.class);
PainlessInstanceBinding instanceBinding = irInvokeCallMemberNode.getDecorationValue(IRDInstanceBinding.class);
Expand All @@ -1669,6 +1671,16 @@ public void visitInvokeCallMember(InvokeCallMemberNode irInvokeCallMemberNode, W
} else {
methodWriter.invokeVirtual(CLASS_TYPE, localFunction.getAsmMethod());
}
} else if (thisMethod != null) {
methodWriter.loadThis();

for (ExpressionNode irArgumentNode : irArgumentNodes) {
visit(irArgumentNode, writeScope);
}

Method asmMethod = new Method(thisMethod.javaMethod.getName(),
thisMethod.methodType.dropParameterTypes(0, 1).toMethodDescriptorString());
methodWriter.invokeVirtual(CLASS_TYPE, asmMethod);
} else if (importedMethod != null) {
for (ExpressionNode irArgumentNode : irArgumentNodes) {
visit(irArgumentNode, writeScope);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,7 @@
import org.elasticsearch.painless.symbol.Decorations.StandardPainlessMethod;
import org.elasticsearch.painless.symbol.Decorations.StaticType;
import org.elasticsearch.painless.symbol.Decorations.TargetType;
import org.elasticsearch.painless.symbol.Decorations.ThisPainlessMethod;
import org.elasticsearch.painless.symbol.Decorations.TypeParameters;
import org.elasticsearch.painless.symbol.Decorations.UnaryType;
import org.elasticsearch.painless.symbol.Decorations.UpcastPainlessCast;
Expand Down Expand Up @@ -1714,6 +1715,7 @@ public void visitCallLocal(ECallLocal userCallLocalNode, SemanticScope semanticS
ScriptScope scriptScope = semanticScope.getScriptScope();

FunctionTable.LocalFunction localFunction = null;
PainlessMethod thisMethod = null;
PainlessMethod importedMethod = null;
PainlessClassBinding classBinding = null;
int classBindingOffset = 0;
Expand All @@ -1728,44 +1730,47 @@ public void visitCallLocal(ECallLocal userCallLocalNode, SemanticScope semanticS
localFunction = null;
}

if (localFunction != null) {
semanticScope.setUsesInstanceMethod();
} else {
importedMethod = scriptScope.getPainlessLookup().lookupImportedPainlessMethod(methodName, userArgumentsSize);
if (localFunction == null) {
thisMethod = scriptScope.getPainlessLookup().lookupPainlessMethod(
scriptScope.getScriptClassInfo().getBaseClass(), false, methodName, userArgumentsSize);

if (importedMethod == null) {
classBinding = scriptScope.getPainlessLookup().lookupPainlessClassBinding(methodName, userArgumentsSize);
if (thisMethod == null) {
importedMethod = scriptScope.getPainlessLookup().lookupImportedPainlessMethod(methodName, userArgumentsSize);

// check to see if this class binding requires an implicit this reference
if (classBinding != null && classBinding.typeParameters.isEmpty() == false &&
classBinding.typeParameters.get(0) == scriptScope.getScriptClassInfo().getBaseClass()) {
classBinding = null;
}
if (importedMethod == null) {
classBinding = scriptScope.getPainlessLookup().lookupPainlessClassBinding(methodName, userArgumentsSize);

if (classBinding == null) {
// This extra check looks for a possible match where the class binding requires an implicit this
// reference. This is a temporary solution to allow the class binding access to data from the
// base script class without need for a user to add additional arguments. A long term solution
// will likely involve adding a class instance binding where any instance can have a class binding
// as part of its API. However, the situation at run-time is difficult and will modifications that
// are a substantial change if even possible to do.
classBinding = scriptScope.getPainlessLookup().lookupPainlessClassBinding(methodName, userArgumentsSize + 1);

if (classBinding != null) {
if (classBinding.typeParameters.isEmpty() == false &&
classBinding.typeParameters.get(0) == scriptScope.getScriptClassInfo().getBaseClass()) {
classBindingOffset = 1;
} else {
classBinding = null;
}
// check to see if this class binding requires an implicit this reference
if (classBinding != null && classBinding.typeParameters.isEmpty() == false &&
classBinding.typeParameters.get(0) == scriptScope.getScriptClassInfo().getBaseClass()) {
classBinding = null;
}

if (classBinding == null) {
instanceBinding = scriptScope.getPainlessLookup().lookupPainlessInstanceBinding(methodName, userArgumentsSize);
// This extra check looks for a possible match where the class binding requires an implicit this
// reference. This is a temporary solution to allow the class binding access to data from the
// base script class without need for a user to add additional arguments. A long term solution
// will likely involve adding a class instance binding where any instance can have a class binding
// as part of its API. However, the situation at run-time is difficult and will modifications that
// are a substantial change if even possible to do.
classBinding = scriptScope.getPainlessLookup().lookupPainlessClassBinding(methodName, userArgumentsSize + 1);

if (classBinding != null) {
if (classBinding.typeParameters.isEmpty() == false &&
classBinding.typeParameters.get(0) == scriptScope.getScriptClassInfo().getBaseClass()) {
classBindingOffset = 1;
} else {
classBinding = null;
}
}

if (classBinding == null) {
instanceBinding = scriptScope.getPainlessLookup().lookupPainlessInstanceBinding(methodName, userArgumentsSize);

if (instanceBinding == null) {
throw userCallLocalNode.createError(new IllegalArgumentException(
"Unknown call [" + methodName + "] with [" + userArgumentNodes + "] arguments."));
if (instanceBinding == null) {
throw userCallLocalNode.createError(new IllegalArgumentException(
"Unknown call [" + methodName + "] with [" + userArgumentNodes + "] arguments."));
}
}
}
}
Expand All @@ -1775,10 +1780,18 @@ public void visitCallLocal(ECallLocal userCallLocalNode, SemanticScope semanticS
List<Class<?>> typeParameters;

if (localFunction != null) {
semanticScope.setUsesInstanceMethod();
semanticScope.putDecoration(userCallLocalNode, new StandardLocalFunction(localFunction));

typeParameters = new ArrayList<>(localFunction.getTypeParameters());
valueType = localFunction.getReturnType();
} else if (thisMethod != null) {
semanticScope.setUsesInstanceMethod();
semanticScope.putDecoration(userCallLocalNode, new ThisPainlessMethod(thisMethod));

scriptScope.markNonDeterministic(thisMethod.annotations.containsKey(NonDeterministicAnnotation.class));
typeParameters = new ArrayList<>(thisMethod.typeParameters);
valueType = thisMethod.returnType;
} else if (importedMethod != null) {
semanticScope.putDecoration(userCallLocalNode, new StandardPainlessMethod(importedMethod));

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,7 @@
import org.elasticsearch.painless.symbol.Decorations.StandardPainlessMethod;
import org.elasticsearch.painless.symbol.Decorations.StaticType;
import org.elasticsearch.painless.symbol.Decorations.TargetType;
import org.elasticsearch.painless.symbol.Decorations.ThisPainlessMethod;
import org.elasticsearch.painless.symbol.Decorations.TypeParameters;
import org.elasticsearch.painless.symbol.Decorations.UnaryType;
import org.elasticsearch.painless.symbol.Decorations.UpcastPainlessCast;
Expand Down Expand Up @@ -239,6 +240,7 @@
import org.elasticsearch.painless.symbol.IRDecorations.IRDSize;
import org.elasticsearch.painless.symbol.IRDecorations.IRDStoreType;
import org.elasticsearch.painless.symbol.IRDecorations.IRDSymbol;
import org.elasticsearch.painless.symbol.IRDecorations.IRDThisMethod;
import org.elasticsearch.painless.symbol.IRDecorations.IRDTypeParameters;
import org.elasticsearch.painless.symbol.IRDecorations.IRDUnaryType;
import org.elasticsearch.painless.symbol.IRDecorations.IRDValue;
Expand Down Expand Up @@ -1221,6 +1223,10 @@ public void visitCallLocal(ECallLocal callLocalNode, ScriptScope scriptScope) {
if (scriptScope.hasDecoration(callLocalNode, StandardLocalFunction.class)) {
LocalFunction localFunction = scriptScope.getDecoration(callLocalNode, StandardLocalFunction.class).getLocalFunction();
irInvokeCallMemberNode.attachDecoration(new IRDFunction(localFunction));
} else if (scriptScope.hasDecoration(callLocalNode, ThisPainlessMethod.class)) {
PainlessMethod thisMethod =
scriptScope.getDecoration(callLocalNode, ThisPainlessMethod.class).getThisPainlessMethod();
irInvokeCallMemberNode.attachDecoration(new IRDThisMethod(thisMethod));
} else if (scriptScope.hasDecoration(callLocalNode, StandardPainlessMethod.class)) {
PainlessMethod importedMethod =
scriptScope.getDecoration(callLocalNode, StandardPainlessMethod.class).getStandardPainlessMethod();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -417,6 +417,19 @@ public LocalFunction getLocalFunction() {
}
}

public static class ThisPainlessMethod implements Decoration {

private final PainlessMethod thisPainlessMethod;

public ThisPainlessMethod(PainlessMethod thisPainlessMethod) {
this.thisPainlessMethod = Objects.requireNonNull(thisPainlessMethod);
}

public PainlessMethod getThisPainlessMethod() {
return thisPainlessMethod;
}
}

public static class StandardPainlessClassBinding implements Decoration {

private final PainlessClassBinding painlessClassBinding;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -370,6 +370,19 @@ public IRDFunction(LocalFunction value) {
}
}

/** describes a method for a node on the script class; which method depends on node type */
public static class IRDThisMethod extends IRDecoration<PainlessMethod> {

public IRDThisMethod(PainlessMethod value) {
super(value);
}

@Override
public String toString() {
return PainlessLookupUtility.buildPainlessMethodKey(getValue().javaMethod.getName(), getValue().typeParameters.size());
}
}

/** describes the call to a class binding */
public static class IRDClassBinding extends IRDecoration<PainlessClassBinding> {

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0 and the Server Side Public License, v 1; you may not use this file except
* in compliance with, at your election, the Elastic License 2.0 or the Server
* Side Public License, v 1.
*/

package org.elasticsearch.painless;

import org.elasticsearch.painless.spi.Whitelist;
import org.elasticsearch.painless.spi.WhitelistLoader;
import org.elasticsearch.script.ScriptContext;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class ThisTests extends ScriptTestCase {

public abstract static class ThisBaseScript {

protected String baseString;

public ThisBaseScript(String baseString) {
this.baseString = baseString;
}

public String getBaseString() {
return baseString;
}

public void setBaseString(String testString) {
this.baseString = testString;
}

public int getBaseLength() {
return baseString.length();
}
}

public abstract static class ThisScript extends ThisBaseScript {

protected String thisString;

public ThisScript(String baseString, String thisString) {
super(baseString);

this.thisString = thisString;
}

public String thisString() {
return thisString;
}

public void thisString(String testString) {
this.thisString = testString;
}

public int thisLength() {
return thisString.length();
}

public abstract Object execute();

public interface Factory {

ThisScript newInstance(String baseString, String testString);
}

public static final String[] PARAMETERS = {};
public static final ScriptContext<ThisScript.Factory> CONTEXT =
new ScriptContext<>("this_test", ThisScript.Factory.class);
}

@Override
protected Map<ScriptContext<?>, List<Whitelist>> scriptContexts() {
Map<ScriptContext<?>, List<Whitelist>> contexts = new HashMap<>();
List<Whitelist> whitelists = new ArrayList<>(Whitelist.BASE_WHITELISTS);
whitelists.add(WhitelistLoader.loadFromResourceFiles(Whitelist.class, "org.elasticsearch.painless.this"));
contexts.put(ThisScript.CONTEXT, whitelists);
return contexts;
}

public Object exec(String script, String baseString, String testString) {
ThisScript.Factory factory = scriptEngine.compile(null, script, ThisScript.CONTEXT, new HashMap<>());
ThisScript testThisScript = factory.newInstance(baseString, testString);
return testThisScript.execute();
}

public void testThisMethods() {
assertEquals("basethis", exec("getBaseString() + thisString()", "base", "this"));
assertEquals(8, exec("getBaseLength() + thisLength()", "yyy", "xxxxx"));

List<String> result = new ArrayList<>();
result.add("this");
result.add("base");
assertEquals(result, exec("List result = []; " +
"thisString('this');" +
"setBaseString('base');" +
"result.add(thisString()); " +
"result.add(getBaseString());" +
"result;", "", ""));
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
class org.elasticsearch.painless.ThisTests$ThisBaseScript @no_import {
String getBaseString();
void setBaseString(String);
int getBaseLength();
}


class org.elasticsearch.painless.ThisTests$ThisScript @no_import {
String thisString();
void thisString(String);
int thisLength();
}

0 comments on commit 34829e8

Please sign in to comment.