diff --git a/presto-main/src/main/java/com/facebook/presto/server/AsyncPageTransportServlet.java b/presto-main/src/main/java/com/facebook/presto/server/AsyncPageTransportServlet.java index 7619c3bb6f5d..79a7911d7bd2 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/AsyncPageTransportServlet.java +++ b/presto-main/src/main/java/com/facebook/presto/server/AsyncPageTransportServlet.java @@ -57,6 +57,7 @@ import static com.facebook.presto.client.PrestoHeaders.PRESTO_TASK_INSTANCE_ID; import static com.facebook.presto.server.security.RoleType.INTERNAL; import static com.facebook.presto.spi.page.PagesSerdeUtil.PAGE_METADATA_SIZE; +import static com.facebook.presto.util.SanitizeMessage.getSanitizeMessage; import static com.facebook.presto.util.TaskUtils.DEFAULT_MAX_WAIT_TIME; import static com.facebook.presto.util.TaskUtils.randomizeWaitTime; import static com.google.common.net.HttpHeaders.CONTENT_LENGTH; @@ -116,9 +117,8 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response) protected void reportFailure(HttpServletResponse response, String message) throws IOException { - response.sendError(SC_BAD_REQUEST, message); + response.sendError(SC_BAD_REQUEST, getSanitizeMessage(message)); } - protected void parseURI(String requestURI, HttpServletRequest request, HttpServletResponse response) throws IOException { @@ -181,7 +181,7 @@ public void onError(AsyncEvent event) { String errorMessage = format("Server error to process task result request %s : %s", requestURI, event.getThrowable().getMessage()); log.error(event.getThrowable(), errorMessage); - response.sendError(SC_INTERNAL_SERVER_ERROR, errorMessage); + response.sendError(SC_INTERNAL_SERVER_ERROR, getSanitizeMessage(errorMessage)); } public void onStartAsync(AsyncEvent event) diff --git a/presto-main/src/main/java/com/facebook/presto/util/SanitizeMessage.java b/presto-main/src/main/java/com/facebook/presto/util/SanitizeMessage.java new file mode 100644 index 000000000000..d72f2fe12f2d --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/util/SanitizeMessage.java @@ -0,0 +1,56 @@ +/* + * 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 + * + * 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. + */ +package com.facebook.presto.util; + +public class SanitizeMessage +{ + private SanitizeMessage() {} + + public static String getSanitizeMessage(String message) + { + if (message == null) { + return ""; + } + StringBuilder stringBuilder = new StringBuilder(message.length()); + + // Traverse each character in the string and replace only when necessary + for (int i = 0; i < message.length(); i++) { + char ch = message.charAt(i); + switch (ch) { + case '&': + stringBuilder.append("&"); + break; + case '<': + stringBuilder.append("<"); + break; + case '>': + stringBuilder.append(">"); + break; + case '"': + stringBuilder.append("""); + break; + case '\'': + stringBuilder.append("'"); + break; + case '/': + stringBuilder.append("/"); + break; + default: + stringBuilder.append(ch); + break; + } + } + return stringBuilder.toString(); + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/util/TestSanitizeMessage.java b/presto-main/src/test/java/com/facebook/presto/util/TestSanitizeMessage.java new file mode 100644 index 000000000000..f37968f07090 --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/util/TestSanitizeMessage.java @@ -0,0 +1,35 @@ +/* + * 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 + * + * 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. + */ +package com.facebook.presto.util; + +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; + +public class TestSanitizeMessage +{ + @Test + public void testGetSanitizeMessage() + { + String message = "Server&> error/ to \" process < task result request ' "; + assertEquals(SanitizeMessage.getSanitizeMessage(message), "Server&> error/ to " process < task result request ' "); + } + + @Test + public void testGetSanitizeMessageNullCheck() + { + String message = null; + assertEquals(SanitizeMessage.getSanitizeMessage(message), ""); + } +}