tool
/
starrocks
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[Feature] support SecurityManager for UDF (#29906) 

Signed-off-by: stdpain <drfeng08@gmail.com>
This commit is contained in:
stdpain 2023-09-08 15:55:05 +08:00 committed by GitHub
parent a71ffe41c8
commit 4029e620a2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
13 changed files with 198 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
be/CMakeLists.txt
View File

@ -1044,6 +1044,7 @@ install(FILES
${BASE_DIR}/../conf/cn.conf
${BASE_DIR}/../conf/hadoop_env.sh
${BASE_DIR}/../conf/log4j2.properties
${BASE_DIR}/../conf/udf_security.policy
DESTINATION ${OUTPUT_DIR}/conf)
if ("${CMAKE_BUILD_TYPE}" STREQUAL "ASAN" OR "${CMAKE_BUILD_TYPE}" STREQUAL "LSAN")

2
be/src/exec/pipeline/exchange/local_exchange.h
View File

@ -100,7 +100,7 @@ class RandomPartitioner final : public Partitioner {
public:
RandomPartitioner(LocalExchangeSourceOperatorFactory* source) : Partitioner(source) {}
virtual ~RandomPartitioner() override = default;
~RandomPartitioner() override = default;
Status shuffle_channel_ids(const ChunkPtr& chunk, int32_t num_partitions) override;
};

21
be/src/udf/java/java_udf.cpp
View File

@ -91,15 +91,15 @@ void JVMFunctionHelper::_init() {
_object_class = JNI_FIND_CLASS("java/lang/Object");
_object_array_class = JNI_FIND_CLASS("[Ljava/lang/Object;");
_string_class = JNI_FIND_CLASS("java/lang/String");
_throwable_class = JNI_FIND_CLASS("java/lang/Throwable");
_jarrays_class = JNI_FIND_CLASS("java/util/Arrays");
_list_class = JNI_FIND_CLASS("java/util/List");
_exception_util_class = JNI_FIND_CLASS("org/apache/commons/lang3/exception/ExceptionUtils");
CHECK(_object_class);
CHECK(_string_class);
CHECK(_throwable_class);
CHECK(_jarrays_class);
CHECK(_list_class);
CHECK(_exception_util_class);
ADD_NUMBERIC_CLASS(boolean, Boolean, Z);
ADD_NUMBERIC_CLASS(byte, Byte, B);
@ -234,20 +234,13 @@ std::string JVMFunctionHelper::to_cxx_string(jstring str) {
}
std::string JVMFunctionHelper::dumpExceptionString(jthrowable throwable) {
std::stringstream ss;
// toString
jmethodID toString = getToStringMethod(_throwable_class);
CHECK(toString != nullptr) << "Not Found JNI method toString";
ss << to_string(throwable);
// e.getStackTrace()
jmethodID getStackTrace = _env->GetMethodID(_throwable_class, "getStackTrace", "()[Ljava/lang/StackTraceElement;");
CHECK(getStackTrace != nullptr) << "Not Found JNI method getStackTrace";
jobject stack_traces = _env->CallObjectMethod((jobject)throwable, getStackTrace);
auto get_stack_trace = _env->GetStaticMethodID(_exception_util_class, "getStackTrace",
"(Ljava/lang/Throwable;)Ljava/lang/String;");
CHECK(get_stack_trace != nullptr) << "Not Found JNI method getStackTrace";
jobject stack_traces = _env->CallStaticObjectMethod(_exception_util_class, get_stack_trace, (jobject)throwable);
LOCAL_REF_GUARD(stack_traces);
CHECK_FUNCTION_EXCEPTION(_env, "dump_string")
ss << array_to_string(stack_traces);
return ss.str();
return to_cxx_string((jstring)stack_traces);
}
jmethodID JVMFunctionHelper::getToStringMethod(jclass clazz) {

2
be/src/udf/java/java_udf.h
View File

@ -159,9 +159,9 @@ private:
jclass _object_class;
jclass _object_array_class;
jclass _string_class;
jclass _throwable_class;
jclass _jarrays_class;
jclass _list_class;
jclass _exception_util_class;
jmethodID _string_construct_with_bytes;

2
build.sh
View File

@ -406,6 +406,7 @@ if [ ${BUILD_FE} -eq 1 -o ${BUILD_SPARK_DPP} -eq 1 ]; then
cp -r -p ${STARROCKS_HOME}/bin/show_fe_version.sh ${STARROCKS_OUTPUT}/fe/bin/
cp -r -p ${STARROCKS_HOME}/bin/common.sh ${STARROCKS_OUTPUT}/fe/bin/
cp -r -p ${STARROCKS_HOME}/conf/fe.conf ${STARROCKS_OUTPUT}/fe/conf/
cp -r -p ${STARROCKS_HOME}/conf/udf_security.policy ${STARROCKS_OUTPUT}/fe/conf/
cp -r -p ${STARROCKS_HOME}/conf/hadoop_env.sh ${STARROCKS_OUTPUT}/fe/conf/
rm -rf ${STARROCKS_OUTPUT}/fe/lib/*
cp -r -p ${STARROCKS_HOME}/fe/fe-core/target/lib/* ${STARROCKS_OUTPUT}/fe/lib/
@ -435,6 +436,7 @@ if [ ${BUILD_BE} -eq 1 ]; then
cp -r -p ${STARROCKS_HOME}/be/output/bin/* ${STARROCKS_OUTPUT}/be/bin/
cp -r -p ${STARROCKS_HOME}/be/output/conf/be.conf ${STARROCKS_OUTPUT}/be/conf/
cp -r -p ${STARROCKS_HOME}/be/output/conf/udf_security.policy ${STARROCKS_OUTPUT}/be/conf/
cp -r -p ${STARROCKS_HOME}/be/output/conf/be_test.conf ${STARROCKS_OUTPUT}/be/conf/
cp -r -p ${STARROCKS_HOME}/be/output/conf/cn.conf ${STARROCKS_OUTPUT}/be/conf/
cp -r -p ${STARROCKS_HOME}/be/output/conf/hadoop_env.sh ${STARROCKS_OUTPUT}/be/conf/

4
conf/fe.conf
View File

@ -24,10 +24,10 @@
LOG_DIR = ${STARROCKS_HOME}/log
DATE = "$(date +%Y%m%d-%H%M%S)"
JAVA_OPTS="-Dlog4j2.formatMsgNoLookups=true -Xmx8192m -XX:+UseMembar -XX:SurvivorRatio=8 -XX:MaxTenuringThreshold=7 -XX:+PrintGCDateStamps -XX:+PrintGCDetails -XX:+UseConcMarkSweepGC -XX:+UseParNewGC -XX:+CMSClassUnloadingEnabled -XX:-CMSParallelRemarkEnabled -XX:CMSInitiatingOccupancyFraction=80 -XX:SoftRefLRUPolicyMSPerMB=0 -Xloggc:${LOG_DIR}/fe.gc.log.$DATE -XX:+PrintConcurrentLocks"
JAVA_OPTS="-Dlog4j2.formatMsgNoLookups=true -Xmx8192m -XX:+UseMembar -XX:SurvivorRatio=8 -XX:MaxTenuringThreshold=7 -XX:+PrintGCDateStamps -XX:+PrintGCDetails -XX:+UseConcMarkSweepGC -XX:+UseParNewGC -XX:+CMSClassUnloadingEnabled -XX:-CMSParallelRemarkEnabled -XX:CMSInitiatingOccupancyFraction=80 -XX:SoftRefLRUPolicyMSPerMB=0 -Xloggc:${LOG_DIR}/fe.gc.log.$DATE -XX:+PrintConcurrentLocks -Djava.security.policy=${STARROCKS_HOME}/conf/udf_security.policy"
# For jdk 11+, this JAVA_OPTS will be used as default JVM options
JAVA_OPTS_FOR_JDK_11="-Dlog4j2.formatMsgNoLookups=true -Xmx8192m -XX:+UseG1GC -Xlog:gc*:${LOG_DIR}/fe.gc.log.$DATE:time"
JAVA_OPTS_FOR_JDK_11="-Dlog4j2.formatMsgNoLookups=true -Xmx8192m -XX:+UseG1GC -Xlog:gc*:${LOG_DIR}/fe.gc.log.$DATE:time -Djava.security.policy=${STARROCKS_HOME}/conf/udf_security.policy"
##
## the lowercase properties are read by main program.

3
conf/udf_security.policy Normal file
View File

@ -0,0 +1,3 @@
grant {
permission java.security.AllPermission;
};

52
fe/fe-core/src/main/java/com/starrocks/sql/ast/CreateFunctionStmt.java
View File

@ -42,12 +42,12 @@ import java.io.InputStream;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.lang.reflect.Parameter;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.net.URLConnection;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.Permission;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
@ -210,6 +210,46 @@ public class CreateFunctionStmt extends DdlStmt {
}
}
public static class UDFInternalClassLoader extends URLClassLoader {
public UDFInternalClassLoader(String udfPath) throws IOException {
super(new URL[] {new URL("jar:" + udfPath + "!/")});
}
}
public class UDFSecurityManager extends SecurityManager {
private Class<?> clazz;
public UDFSecurityManager(Class<?> clazz) {
this.clazz = clazz;
}
@Override
public void checkPermission(Permission perm) {
if (isCreateFromUDFClassLoader()) {
super.checkPermission(perm);
}
}
public void checkPermission(Permission perm, Object context) {
if (isCreateFromUDFClassLoader()) {
super.checkPermission(perm, context);
}
}
private boolean isCreateFromUDFClassLoader() {
Class<?>[] classContext = getClassContext();
if (classContext.length >= 2) {
for (int i = 1; i < classContext.length; i++) {
if (classContext[i].getClassLoader() != null &&
clazz.equals(classContext[i].getClassLoader().getClass())) {
return true;
}
}
}
return false;
}
}
private UDFInternalClass mainClass;
private UDFInternalClass udafStateClass;
@ -311,8 +351,8 @@ public class CreateFunctionStmt extends DdlStmt {
}
try {
URL[] urls = {new URL("jar:" + objectFile + "!/")};
try (URLClassLoader classLoader = URLClassLoader.newInstance(urls)) {
System.setSecurityManager(new UDFSecurityManager(UDFInternalClass.class));
try (URLClassLoader classLoader = new UDFInternalClassLoader(objectFile)) {
mainClass.setClazz(classLoader.loadClass(className));
if (isAggregate) {
@ -329,9 +369,11 @@ public class CreateFunctionStmt extends DdlStmt {
throw new AnalysisException("Failed to load object_file: " + objectFile);
} catch (ClassNotFoundException e) {
throw new AnalysisException("Class '" + className + "' not found in object_file :" + objectFile);
} catch (Exception e) {
throw new AnalysisException("other exception when load class. exception:", e);
}
} catch (MalformedURLException e) {
throw new AnalysisException("Object file is invalid: " + objectFile);
} finally {
System.setSecurityManager(null);
}
}

6
java-extensions/udf-extensions/pom.xml
View File

@ -28,6 +28,12 @@
<version>4.13.1</version>
<scope>test</scope>
</dependency>
<!-- https://mvnrepository.com/artifact/org.apache.commons/commons-lang3 -->
<dependency>
<groupId>org.apache.commons</groupId>
<artifactId>commons-lang3</artifactId>
<version>3.9</version>
</dependency>
<dependency>
<groupId>com.starrocks</groupId>

10
java-extensions/udf-extensions/src/main/java/com/starrocks/udf/UDFClassLoader.java
View File

@ -30,10 +30,17 @@ public class UDFClassLoader extends URLClassLoader {
private Map<String, Class<?>> genClazzMap = new HashMap<>();
private static final int SINGLE_BATCH_UPDATE = 1;
private static final int BATCH_EVALUATE = 2;
private static final int BATCH_EVALUATE = 2;
public UDFClassLoader(String udfPath) throws IOException {
super(new URL[] {new URL("file://" + udfPath)});
if (System.getSecurityManager() == null && System.getProperties().get("java.security.policy") != null) {
synchronized (UDFClassLoader.class) {
if (System.getSecurityManager() == null) {
System.setSecurityManager(new UDFSecurityManager(UDFClassLoader.class));
}
}
}
}
@Override
@ -45,7 +52,6 @@ public class UDFClassLoader extends URLClassLoader {
return super.findClass(clazzName);
}
public Class<?> generateCallStubV(String name, Class<?> clazz, Method method, int genType) {
String clazzName = name.replace("/", ".");
if (!clazzName.startsWith(CallStubGenerator.GEN_KEYWORD)) {

51
java-extensions/udf-extensions/src/main/java/com/starrocks/udf/UDFSecurityManager.java Normal file
View File

@ -0,0 +1,51 @@
// Copyright 2021-present StarRocks, Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package com.starrocks.udf;
import java.security.Permission;
public class UDFSecurityManager extends SecurityManager {
private Class<?> clazz;
public UDFSecurityManager(Class<?> clazz) {
this.clazz = clazz;
}
@Override
public void checkPermission(Permission perm) {
if (isCreateFromUDFClassLoader()) {
super.checkPermission(perm);
}
}
public void checkPermission(Permission perm, Object context) {
if (isCreateFromUDFClassLoader()) {
super.checkPermission(perm, context);
}
}
private boolean isCreateFromUDFClassLoader() {
Class<?>[] classContext = getClassContext();
if (classContext.length >= 2) {
for (int i = 1; i < classContext.length; i++) {
if (classContext[i].getClassLoader() != null &&
clazz.equals(classContext[i].getClassLoader().getClass())) {
return true;
}
}
}
return false;
}
}

66
java-extensions/udf-extensions/src/test/java/com/starrocks/udf/SecurityTest.java Normal file
View File

@ -0,0 +1,66 @@
package com.starrocks.udf;
import org.junit.Test;
import java.io.File;
import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.security.AccessControlException;
public class SecurityTest {
public static class TestClassLoader extends ClassLoader {
public TestClassLoader(String clazzName, byte[] bytes) {
defineClass(clazzName, bytes, 0, bytes.length);
}
}
public static class ScalarAdd {
public String evaluate(String v1) throws IOException {
File.createTempFile(v1, ".txt");
return v1;
}
}
private static Method getFirstMethod(Class<?> clazz, String name) {
Method call = null;
for (Method declaredMethod : clazz.getDeclaredMethods()) {
if (declaredMethod.getName().equals(name)) {
call = declaredMethod;
}
}
return call;
}
@Test(expected = AccessControlException.class)
public void testUDFCreateFile()
throws NoSuchMethodException, ClassNotFoundException, IOException, InvocationTargetException,
IllegalAccessException {
System.setSecurityManager(new UDFSecurityManager(TestClassLoader.class));
Class<?> clazz = ScalarAdd.class;
final String genClassName = CallStubGenerator.CLAZZ_NAME.replace("/", ".");
Method m = clazz.getMethod("evaluate", String.class);
final byte[] updates =
CallStubGenerator.generateScalarCallStub(clazz, m);
ClassLoader classLoader = new TestClassLoader(genClassName, updates);
final Class<?> stubClazz = classLoader.loadClass(genClassName);
Method batchCall = getFirstMethod(stubClazz, "batchCallV");
String[] inputs1 = new String[1];
inputs1[0] = "prefix";
ScalarAdd concat = new ScalarAdd();
batchCall.invoke(null, 1, concat, inputs1);
System.setSecurityManager(null);
}
@Test
public void testNoUDFCreateFile() throws IOException {
System.setSecurityManager(new UDFSecurityManager(TestClassLoader.class));
ScalarAdd concat = new ScalarAdd();
concat.evaluate("test");
System.setSecurityManager(null);
}
}

3
java-extensions/udf-extensions/src/test/resources/PermissionCheck.policy Normal file
View File

@ -0,0 +1,3 @@
// detail see $JAVA_HOME/jre/lib/security/java.policy
grant {
};