Skip to content

Commit

Permalink
[HotFix] Fix TaskOutputParameterParser might OOM if meed a bad output…
Browse files Browse the repository at this point in the history
… param expression (#15264)
  • Loading branch information
ruanwenjun authored Dec 4, 2023
1 parent 7bfc6dc commit db3d84b
Show file tree
Hide file tree
Showing 5 changed files with 126 additions and 19 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -36,46 +36,78 @@
@NotThreadSafe
public class TaskOutputParameterParser {

private final Map<String, String> taskOutputParams = new HashMap<>();
// Used to avoid '${setValue(' which loss the end of ')}'
private final int maxOneParameterRows;

// Used to avoid '${setValue(' which length is too long, this may case OOM
private final int maxOneParameterLength;

private final Map<String, String> taskOutputParams;

private List<String> currentTaskOutputParam;

public void appendParseLog(String log) {
if (log == null) {
private long currentTaskOutputParamLength;

public TaskOutputParameterParser() {
// the default max rows of one parameter is 1024, this should be enough
this(1024, Integer.MAX_VALUE);
}

public TaskOutputParameterParser(int maxOneParameterRows, int maxOneParameterLength) {
this.maxOneParameterRows = maxOneParameterRows;
this.maxOneParameterLength = maxOneParameterLength;
this.taskOutputParams = new HashMap<>();
this.currentTaskOutputParam = null;
this.currentTaskOutputParamLength = 0;
}

public void appendParseLog(String logLine) {
if (logLine == null) {
return;
}

if (currentTaskOutputParam != null) {
if (currentTaskOutputParam.size() > maxOneParameterRows
|| currentTaskOutputParamLength > maxOneParameterLength) {
log.warn(
"The output param expression '{}' is too long, the max rows is {}, max length is {}, will skip this param",
String.join("\n", currentTaskOutputParam), maxOneParameterLength, maxOneParameterRows);
currentTaskOutputParam = null;
currentTaskOutputParamLength = 0;
return;
}
// continue to parse the rest of line
int i = log.indexOf(")}");
int i = logLine.indexOf(")}");
if (i == -1) {
// the end of var pool not found
currentTaskOutputParam.add(log);
currentTaskOutputParam.add(logLine);
currentTaskOutputParamLength += logLine.length();
} else {
// the end of var pool found
currentTaskOutputParam.add(log.substring(0, i + 2));
currentTaskOutputParam.add(logLine.substring(0, i + 2));
Pair<String, String> keyValue = parseOutputParam(String.join("\n", currentTaskOutputParam));
if (keyValue.getKey() != null && keyValue.getValue() != null) {
taskOutputParams.put(keyValue.getKey(), keyValue.getValue());
}
currentTaskOutputParam = null;
currentTaskOutputParamLength = 0;
// continue to parse the rest of line
if (i + 2 != log.length()) {
appendParseLog(log.substring(i + 2));
if (i + 2 != logLine.length()) {
appendParseLog(logLine.substring(i + 2));
}
}
return;
}

int indexOfVarPoolBegin = log.indexOf("${setValue(");
int indexOfVarPoolBegin = logLine.indexOf("${setValue(");
if (indexOfVarPoolBegin == -1) {
indexOfVarPoolBegin = log.indexOf("#{setValue(");
indexOfVarPoolBegin = logLine.indexOf("#{setValue(");
}
if (indexOfVarPoolBegin == -1) {
return;
}
currentTaskOutputParam = new ArrayList<>();
appendParseLog(log.substring(indexOfVarPoolBegin));
appendParseLog(logLine.substring(indexOfVarPoolBegin));
}

public Map<String, String> getTaskOutputParams() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,14 +19,15 @@

import static org.junit.jupiter.api.Assertions.assertEquals;

import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

import lombok.SneakyThrows;

import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;

Expand All @@ -35,23 +36,23 @@
class TaskOutputParameterParserTest {

@Test
void testEmptyLog() throws IOException, URISyntaxException {
void testEmptyLog() {
List<String> varPools = getLogs("/outputParam/emptyVarPoolLog.txt");
TaskOutputParameterParser taskOutputParameterParser = new TaskOutputParameterParser();
varPools.forEach(taskOutputParameterParser::appendParseLog);
Assertions.assertTrue(taskOutputParameterParser.getTaskOutputParams().isEmpty());
}

@Test
void testOneLineLog() throws IOException, URISyntaxException {
void testOneLineLog() {
List<String> varPools = getLogs("/outputParam/onelineVarPoolLog.txt");
TaskOutputParameterParser taskOutputParameterParser = new TaskOutputParameterParser();
varPools.forEach(taskOutputParameterParser::appendParseLog);
assertEquals(ImmutableMap.of("name", "name=tom"), taskOutputParameterParser.getTaskOutputParams());
}

@Test
void testOneVarPollInMultiLineLog() throws IOException, URISyntaxException {
void testOneVarPoolInMultiLineLog() {
List<String> varPools = getLogs("/outputParam/oneVarPollInMultiLineLog.txt");
TaskOutputParameterParser taskOutputParameterParser = new TaskOutputParameterParser();
varPools.forEach(taskOutputParameterParser::appendParseLog);
Expand All @@ -63,14 +64,46 @@ void testOneVarPollInMultiLineLog() throws IOException, URISyntaxException {
}

@Test
void testVarPollInMultiLineLog() throws IOException, URISyntaxException {
List<String> varPools = getLogs("/outputParam/multipleVarPoll.txt");
void testVarPoolInMultiLineLog() {
List<String> varPools = getLogs("/outputParam/multipleVarPool.txt");
TaskOutputParameterParser taskOutputParameterParser = new TaskOutputParameterParser();
varPools.forEach(taskOutputParameterParser::appendParseLog);
assertEquals(ImmutableMap.of("name", "tom", "age", "1"), taskOutputParameterParser.getTaskOutputParams());
}

private List<String> getLogs(String file) throws IOException, URISyntaxException {
@Test
void textVarPoolExceedMaxRows() {
List<String> varPools = getLogs("/outputParam/maxRowsVarPool.txt");
TaskOutputParameterParser taskOutputParameterParser = new TaskOutputParameterParser(2, Integer.MAX_VALUE);
varPools.forEach(taskOutputParameterParser::appendParseLog);
assertEquals(Collections.emptyMap(), taskOutputParameterParser.getTaskOutputParams());

taskOutputParameterParser = new TaskOutputParameterParser();
varPools.forEach(taskOutputParameterParser::appendParseLog);
assertEquals(ImmutableMap.of("name", "name=tom\n" +
"name=name=tom\n" +
"name=name=tom\n" +
"name=name=tom\n" +
"name=name=tom"), taskOutputParameterParser.getTaskOutputParams());

}

@Test
void textVarPoolExceedMaxLength() {
List<String> varPools = getLogs("/outputParam/maxLengthVarPool.txt");
TaskOutputParameterParser taskOutputParameterParser = new TaskOutputParameterParser(2, 10);
varPools.forEach(taskOutputParameterParser::appendParseLog);
assertEquals(Collections.emptyMap(), taskOutputParameterParser.getTaskOutputParams());

taskOutputParameterParser = new TaskOutputParameterParser();
varPools.forEach(taskOutputParameterParser::appendParseLog);
assertEquals(ImmutableMap.of("name", "123456789\n" +
"12345\n"), taskOutputParameterParser.getTaskOutputParams());

}

@SneakyThrows
private List<String> getLogs(String file) {
URI uri = TaskOutputParameterParserTest.class.getResource(file).toURI();
return Files.lines(Paths.get(uri)).collect(Collectors.toList());
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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
#
# http://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.
#

${setValue(name=123456789
12345
)}
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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
#
# http://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.
#

${setValue(name=name=tom
name=name=tom
name=name=tom
name=name=tom
name=name=tom)}

0 comments on commit db3d84b

Please sign in to comment.