This is an automated email from the ASF dual-hosted git repository. HTHou pushed a commit to branch codex/rest-request-limits-dev-1.3 in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit e6f19a30ef83ec93d1bd3deb734b8caaeeca3282 Author: HTHou <[email protected]> AuthorDate: Thu Aug 13 12:02:21 2026 +0800 Limit REST write request size on dev/1.3 --- .../org/apache/iotdb/db/conf/IoTDBDescriptor.java | 4 + .../iotdb/db/conf/rest/IoTDBRestServiceConfig.java | 62 +++++ .../db/conf/rest/IoTDBRestServiceDescriptor.java | 82 +++++- .../exception/RequestLimitExceededException.java | 25 ++ .../filter/RequestBodyMemoryReleaseFilter.java | 33 +++ .../rest/filter/RequestSizeLimitFilter.java | 192 ++++++++++++++ .../rest/filter/RestRequestBodyMemoryManager.java | 114 ++++++++ .../protocol/rest/handler/RequestLimitChecker.java | 56 ++++ .../protocol/rest/v1/handler/ExceptionHandler.java | 10 +- .../rest/v1/handler/RequestValidationHandler.java | 24 ++ .../protocol/rest/v1/impl/RestApiServiceImpl.java | 4 +- .../protocol/rest/v2/handler/ExceptionHandler.java | 10 +- .../rest/v2/handler/RequestValidationHandler.java | 39 +++ .../protocol/rest/v2/impl/RestApiServiceImpl.java | 8 +- .../conf/rest/IoTDBRestServiceDescriptorTest.java | 91 +++++++ .../rest/filter/RequestSizeLimitFilterTest.java | 286 +++++++++++++++++++++ .../rest/handler/RequestValidationLimitTest.java | 110 ++++++++ .../src/test/resources/iotdb-common.properties | 15 ++ .../src/test/resources/iotdb-system.properties | 15 ++ .../conf/iotdb-system.properties.template | 29 ++- 20 files changed, 1198 insertions(+), 11 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java index 5658e10cc60..a20d2f1e119 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java @@ -33,6 +33,7 @@ import org.apache.iotdb.commons.utils.NodeUrlUtils; import org.apache.iotdb.confignode.rpc.thrift.TCQConfig; import org.apache.iotdb.confignode.rpc.thrift.TGlobalConfig; import org.apache.iotdb.confignode.rpc.thrift.TRatisConfig; +import org.apache.iotdb.db.conf.rest.IoTDBRestServiceDescriptor; import org.apache.iotdb.db.consensus.DataRegionConsensusImpl; import org.apache.iotdb.db.exception.query.QueryProcessException; import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager; @@ -2062,6 +2063,9 @@ public class IoTDBDescriptor { // update trusted_uri_pattern loadTrustedUriPattern(properties); + // update REST request limits + IoTDBRestServiceDescriptor.getInstance().loadHotModifiedProps(properties); + // tvlist_sort_threshold conf.setTVListSortThreshold( Integer.parseInt( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/rest/IoTDBRestServiceConfig.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/rest/IoTDBRestServiceConfig.java index 64c0f65fe30..bca061a905f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/rest/IoTDBRestServiceConfig.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/rest/IoTDBRestServiceConfig.java @@ -59,6 +59,22 @@ public class IoTDBRestServiceConfig { private int restQueryDefaultRowSizeLimit = 10000; + /** Maximum accepted REST request body size in bytes. */ + private long restMaxRequestBodySizeInBytes = 16 * 1024 * 1024L; + + /** Maximum total in-flight REST request body size in bytes across concurrent requests. */ + private long restMaxTotalConcurrentRequestBodySizeInBytes = + getDefaultRequestBodyMemoryLimitInBytes(); + + /** Maximum row count accepted by a single REST write request. */ + private int restMaxInsertRows = 100000; + + /** Maximum column count accepted by a single REST write request. */ + private int restMaxInsertColumns = 1024; + + /** Maximum value cell count accepted by a single REST write request. */ + private long restMaxInsertValues = 1000000L; + /** Is client authentication required. */ private boolean clientAuth = false; @@ -173,4 +189,50 @@ public class IoTDBRestServiceConfig { public void setRestQueryDefaultRowSizeLimit(int restQueryDefaultRowSizeLimit) { this.restQueryDefaultRowSizeLimit = restQueryDefaultRowSizeLimit; } + + public long getRestMaxRequestBodySizeInBytes() { + return restMaxRequestBodySizeInBytes; + } + + public void setRestMaxRequestBodySizeInBytes(long restMaxRequestBodySizeInBytes) { + this.restMaxRequestBodySizeInBytes = restMaxRequestBodySizeInBytes; + } + + public long getRestMaxTotalConcurrentRequestBodySizeInBytes() { + return restMaxTotalConcurrentRequestBodySizeInBytes; + } + + public void setRestMaxTotalConcurrentRequestBodySizeInBytes( + long restMaxTotalConcurrentRequestBodySizeInBytes) { + this.restMaxTotalConcurrentRequestBodySizeInBytes = + restMaxTotalConcurrentRequestBodySizeInBytes; + } + + public int getRestMaxInsertRows() { + return restMaxInsertRows; + } + + public void setRestMaxInsertRows(int restMaxInsertRows) { + this.restMaxInsertRows = restMaxInsertRows; + } + + public int getRestMaxInsertColumns() { + return restMaxInsertColumns; + } + + public void setRestMaxInsertColumns(int restMaxInsertColumns) { + this.restMaxInsertColumns = restMaxInsertColumns; + } + + public long getRestMaxInsertValues() { + return restMaxInsertValues; + } + + public void setRestMaxInsertValues(long restMaxInsertValues) { + this.restMaxInsertValues = restMaxInsertValues; + } + + public static long getDefaultRequestBodyMemoryLimitInBytes() { + return Runtime.getRuntime().maxMemory() / 20; + } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/rest/IoTDBRestServiceDescriptor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/rest/IoTDBRestServiceDescriptor.java index ee8ca0705b5..1fc3d376bef 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/rest/IoTDBRestServiceDescriptor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/rest/IoTDBRestServiceDescriptor.java @@ -38,6 +38,16 @@ import java.nio.charset.StandardCharsets; public class IoTDBRestServiceDescriptor { private static final Logger logger = LoggerFactory.getLogger(IoTDBRestServiceDescriptor.class); + private static final String REST_QUERY_DEFAULT_ROW_SIZE_LIMIT = + "rest_query_default_row_size_limit"; + private static final String REST_MAX_REQUEST_BODY_SIZE_IN_BYTES = + "rest_max_request_body_size_in_bytes"; + private static final String REST_MAX_TOTAL_CONCURRENT_REQUEST_BODY_SIZE_IN_BYTES = + "rest_max_total_concurrent_request_body_size_in_bytes"; + private static final String REST_MAX_INSERT_ROWS = "rest_max_insert_rows"; + private static final String REST_MAX_INSERT_COLUMNS = "rest_max_insert_columns"; + private static final String REST_MAX_INSERT_VALUES = "rest_max_insert_values"; + private final IoTDBRestServiceConfig conf = new IoTDBRestServiceConfig(); protected IoTDBRestServiceDescriptor() { @@ -86,11 +96,7 @@ public class IoTDBRestServiceDescriptor { Integer.parseInt( properties.getProperty( "rest_service_port", Integer.toString(conf.getRestServicePort())))); - conf.setRestQueryDefaultRowSizeLimit( - Integer.parseInt( - properties.getProperty( - "rest_query_default_row_size_limit", - Integer.toString(conf.getRestQueryDefaultRowSizeLimit())))); + loadRuntimeLimitProps(properties); conf.setEnableSwagger( Boolean.parseBoolean( properties.getProperty("enable_swagger", Boolean.toString(conf.isEnableSwagger())))); @@ -111,6 +117,72 @@ public class IoTDBRestServiceDescriptor { "idle_timeout_in_seconds", Integer.toString(conf.getIdleTimeoutInSeconds())))); } + public synchronized void loadHotModifiedProps(TrimProperties properties) { + loadRuntimeLimitProps(properties); + } + + private void loadRuntimeLimitProps(TrimProperties properties) { + conf.setRestQueryDefaultRowSizeLimit( + Integer.parseInt( + properties.getProperty( + REST_QUERY_DEFAULT_ROW_SIZE_LIMIT, + Integer.toString(conf.getRestQueryDefaultRowSizeLimit())))); + conf.setRestMaxRequestBodySizeInBytes( + Long.parseLong( + properties.getProperty( + REST_MAX_REQUEST_BODY_SIZE_IN_BYTES, + Long.toString(conf.getRestMaxRequestBodySizeInBytes())))); + conf.setRestMaxTotalConcurrentRequestBodySizeInBytes( + parseMaxTotalConcurrentRequestBodySizeInBytes(properties)); + conf.setRestMaxInsertRows( + Integer.parseInt( + properties.getProperty( + REST_MAX_INSERT_ROWS, Integer.toString(conf.getRestMaxInsertRows())))); + conf.setRestMaxInsertColumns( + Integer.parseInt( + properties.getProperty( + REST_MAX_INSERT_COLUMNS, Integer.toString(conf.getRestMaxInsertColumns())))); + conf.setRestMaxInsertValues( + Long.parseLong( + properties.getProperty( + REST_MAX_INSERT_VALUES, Long.toString(conf.getRestMaxInsertValues())))); + } + + private long parseMaxTotalConcurrentRequestBodySizeInBytes(TrimProperties properties) { + long configuredLimit = + Long.parseLong( + properties.getProperty( + REST_MAX_TOTAL_CONCURRENT_REQUEST_BODY_SIZE_IN_BYTES, + Long.toString(conf.getRestMaxTotalConcurrentRequestBodySizeInBytes()))); + return configuredLimit == 0 + ? calculateRequestBodyMemoryLimitInBytes(properties) + : configuredLimit; + } + + static long calculateRequestBodyMemoryLimitInBytes(TrimProperties properties) { + String memoryAllocateProportion = properties.getProperty("datanode_memory_proportion", null); + if (memoryAllocateProportion == null) { + memoryAllocateProportion = + properties.getProperty("storage_query_schema_consensus_free_memory_proportion", null); + } + if (memoryAllocateProportion == null) { + return IoTDBRestServiceConfig.getDefaultRequestBodyMemoryLimitInBytes(); + } + + String[] proportions = memoryAllocateProportion.split(":"); + int proportionSum = 0; + for (String proportion : proportions) { + proportionSum += Integer.parseInt(proportion.trim()); + } + if (proportionSum == 0 || proportions.length < 6) { + return IoTDBRestServiceConfig.getDefaultRequestBodyMemoryLimitInBytes(); + } + return Runtime.getRuntime().maxMemory() + * Integer.parseInt(proportions[proportions.length - 1].trim()) + / proportionSum + / 2; + } + /** * get props url location * diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/exception/RequestLimitExceededException.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/exception/RequestLimitExceededException.java new file mode 100644 index 00000000000..95ba122d34f --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/exception/RequestLimitExceededException.java @@ -0,0 +1,25 @@ +/* + * 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. + */ + +package org.apache.iotdb.db.protocol.rest.exception; + +public class RequestLimitExceededException extends IllegalArgumentException { + + public RequestLimitExceededException(String message) { + super(message); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/filter/RequestBodyMemoryReleaseFilter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/filter/RequestBodyMemoryReleaseFilter.java new file mode 100644 index 00000000000..9cad6c431d5 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/filter/RequestBodyMemoryReleaseFilter.java @@ -0,0 +1,33 @@ +/* + * 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. + */ + +package org.apache.iotdb.db.protocol.rest.filter; + +import javax.ws.rs.container.ContainerRequestContext; +import javax.ws.rs.container.ContainerResponseContext; +import javax.ws.rs.container.ContainerResponseFilter; +import javax.ws.rs.ext.Provider; + +@Provider +public class RequestBodyMemoryReleaseFilter implements ContainerResponseFilter { + + @Override + public void filter( + ContainerRequestContext requestContext, ContainerResponseContext responseContext) { + RestRequestBodyMemoryManager.releaseReservation(requestContext); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/filter/RequestSizeLimitFilter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/filter/RequestSizeLimitFilter.java new file mode 100644 index 00000000000..4c65bc0f9ef --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/filter/RequestSizeLimitFilter.java @@ -0,0 +1,192 @@ +/* + * 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. + */ + +package org.apache.iotdb.db.protocol.rest.filter; + +import org.apache.iotdb.db.conf.rest.IoTDBRestServiceDescriptor; +import org.apache.iotdb.db.protocol.rest.model.ExecutionStatus; + +import javax.ws.rs.WebApplicationException; +import javax.ws.rs.container.ContainerRequestContext; +import javax.ws.rs.container.ContainerRequestFilter; +import javax.ws.rs.container.PreMatching; +import javax.ws.rs.core.MediaType; +import javax.ws.rs.core.Response; +import javax.ws.rs.ext.Provider; + +import java.io.FilterInputStream; +import java.io.IOException; +import java.io.InputStream; + +@Provider +@PreMatching +public class RequestSizeLimitFilter implements ContainerRequestFilter { + + private static final int PAYLOAD_TOO_LARGE_STATUS_CODE = 413; + private static final int SERVICE_UNAVAILABLE_STATUS_CODE = 503; + + @Override + public void filter(ContainerRequestContext requestContext) { + long maxBodySize = + IoTDBRestServiceDescriptor.getInstance().getConfig().getRestMaxRequestBodySizeInBytes(); + long memoryLimit = + IoTDBRestServiceDescriptor.getInstance() + .getConfig() + .getRestMaxTotalConcurrentRequestBodySizeInBytes(); + if (maxBodySize <= 0 && memoryLimit <= 0) { + return; + } + + int contentLength = requestContext.getLength(); + if (maxBodySize > 0 && contentLength > maxBodySize) { + requestContext.abortWith(buildPayloadTooLargeResponse(maxBodySize)); + return; + } + + RestRequestBodyMemoryManager.Reservation memoryReservation = + RestRequestBodyMemoryManager.newReservation(memoryLimit); + long memoryReservedByContentLength = 0; + if (contentLength > 0 && memoryReservation.isEnabled()) { + if (!memoryReservation.reserve(contentLength)) { + memoryReservation.close(); + requestContext.abortWith(buildMemoryQuotaExceededResponse(memoryLimit)); + return; + } + memoryReservedByContentLength = contentLength; + RestRequestBodyMemoryManager.registerReservation(requestContext, memoryReservation); + } + + requestContext.setEntityStream( + new LimitedInputStream( + requestContext.getEntityStream(), + maxBodySize, + memoryLimit, + memoryReservation, + memoryReservedByContentLength)); + if (memoryReservation.isEnabled() && memoryReservedByContentLength == 0) { + RestRequestBodyMemoryManager.registerReservation(requestContext, memoryReservation); + } + } + + private static Response buildPayloadTooLargeResponse(long maxBodySize) { + return Response.status(PAYLOAD_TOO_LARGE_STATUS_CODE) + .type(MediaType.APPLICATION_JSON_TYPE) + .entity( + new ExecutionStatus() + .code(PAYLOAD_TOO_LARGE_STATUS_CODE) + .message( + String.format( + "REST request body exceeds limit %d bytes. Use SET CONFIGURATION " + + "rest_max_request_body_size_in_bytes=<bytes> to increase it.", + maxBodySize))) + .build(); + } + + private static Response buildMemoryQuotaExceededResponse(long memoryLimit) { + return Response.status(SERVICE_UNAVAILABLE_STATUS_CODE) + .type(MediaType.APPLICATION_JSON_TYPE) + .entity( + new ExecutionStatus() + .code(SERVICE_UNAVAILABLE_STATUS_CODE) + .message( + String.format( + "REST request body memory quota exceeds limit %d bytes. Use SET " + + "CONFIGURATION rest_max_total_concurrent_request_body_size_in_bytes=" + + "<bytes> to increase it.", + memoryLimit))) + .build(); + } + + private static class LimitedInputStream extends FilterInputStream { + + private final long maxBodySize; + private final long memoryLimit; + private final RestRequestBodyMemoryManager.Reservation memoryReservation; + private long memoryCoveredBytes; + private long bytesRead; + + private LimitedInputStream( + InputStream in, + long maxBodySize, + long memoryLimit, + RestRequestBodyMemoryManager.Reservation memoryReservation, + long memoryCoveredBytes) { + super(in); + this.maxBodySize = maxBodySize; + this.memoryLimit = memoryLimit; + this.memoryReservation = memoryReservation; + this.memoryCoveredBytes = memoryCoveredBytes; + } + + @Override + public int read() throws IOException { + int result = super.read(); + if (result != -1) { + incrementBytesRead(1); + } + return result; + } + + @Override + public int read(byte[] bytes, int offset, int length) throws IOException { + int result = super.read(bytes, offset, length); + if (result > 0) { + incrementBytesRead(result); + } + return result; + } + + @Override + public long skip(long byteCount) throws IOException { + long skippedBytes = super.skip(byteCount); + if (skippedBytes > 0) { + incrementBytesRead(skippedBytes); + } + return skippedBytes; + } + + private void incrementBytesRead(long increment) { + bytesRead += increment; + if (maxBodySize > 0 && bytesRead > maxBodySize) { + memoryReservation.close(); + throw new WebApplicationException(buildPayloadTooLargeResponse(maxBodySize)); + } + reserveMemoryIfNecessary(); + } + + private void reserveMemoryIfNecessary() { + if (bytesRead <= memoryCoveredBytes) { + return; + } + long sizeToReserve = bytesRead - memoryCoveredBytes; + if (!memoryReservation.reserve(sizeToReserve)) { + memoryReservation.close(); + throw new WebApplicationException(buildMemoryQuotaExceededResponse(memoryLimit)); + } + memoryCoveredBytes = bytesRead; + } + + @Override + public void close() throws IOException { + try { + super.close(); + } finally { + memoryReservation.close(); + } + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/filter/RestRequestBodyMemoryManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/filter/RestRequestBodyMemoryManager.java new file mode 100644 index 00000000000..e341d6a9d44 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/filter/RestRequestBodyMemoryManager.java @@ -0,0 +1,114 @@ +/* + * 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. + */ + +package org.apache.iotdb.db.protocol.rest.filter; + +import javax.ws.rs.container.ContainerRequestContext; + +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; + +final class RestRequestBodyMemoryManager { + + static final String REQUEST_BODY_MEMORY_RESERVATION_PROPERTY = + RestRequestBodyMemoryManager.class.getName() + ".requestBodyMemoryReservation"; + private static final AtomicLong RESERVED_MEMORY_IN_BYTES = new AtomicLong(); + + private RestRequestBodyMemoryManager() {} + + static Reservation newReservation(long memoryLimitInBytes) { + return new Reservation(memoryLimitInBytes); + } + + static long getReservedMemoryInBytes() { + return RESERVED_MEMORY_IN_BYTES.get(); + } + + static void resetForTest() { + RESERVED_MEMORY_IN_BYTES.set(0); + } + + static void registerReservation( + ContainerRequestContext requestContext, Reservation memoryReservation) { + if (memoryReservation.isEnabled()) { + requestContext.setProperty(REQUEST_BODY_MEMORY_RESERVATION_PROPERTY, memoryReservation); + } + } + + static void releaseReservation(ContainerRequestContext requestContext) { + Object reservation = requestContext.getProperty(REQUEST_BODY_MEMORY_RESERVATION_PROPERTY); + if (reservation instanceof Reservation) { + ((Reservation) reservation).close(); + requestContext.removeProperty(REQUEST_BODY_MEMORY_RESERVATION_PROPERTY); + } + } + + private static boolean tryReserve(long sizeInBytes, long memoryLimitInBytes) { + if (sizeInBytes <= 0 || memoryLimitInBytes <= 0) { + return true; + } + + while (true) { + long currentReservedBytes = RESERVED_MEMORY_IN_BYTES.get(); + if (sizeInBytes > memoryLimitInBytes + || currentReservedBytes > memoryLimitInBytes - sizeInBytes) { + return false; + } + if (RESERVED_MEMORY_IN_BYTES.compareAndSet( + currentReservedBytes, currentReservedBytes + sizeInBytes)) { + return true; + } + } + } + + static final class Reservation implements AutoCloseable { + + private final long memoryLimitInBytes; + private final AtomicBoolean released = new AtomicBoolean(); + private long reservedMemoryInBytes; + + private Reservation(long memoryLimitInBytes) { + this.memoryLimitInBytes = memoryLimitInBytes; + } + + boolean isEnabled() { + return memoryLimitInBytes > 0; + } + + synchronized boolean reserve(long sizeInBytes) { + if (sizeInBytes <= 0 || !isEnabled()) { + return true; + } + if (released.get() || !tryReserve(sizeInBytes, memoryLimitInBytes)) { + return false; + } + reservedMemoryInBytes += sizeInBytes; + return true; + } + + @Override + public synchronized void close() { + if (released.compareAndSet(false, true)) { + long reservedBytes = reservedMemoryInBytes; + reservedMemoryInBytes = 0; + if (reservedBytes > 0) { + RESERVED_MEMORY_IN_BYTES.addAndGet(-reservedBytes); + } + } + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/handler/RequestLimitChecker.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/handler/RequestLimitChecker.java new file mode 100644 index 00000000000..7f7b8b50bf6 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/handler/RequestLimitChecker.java @@ -0,0 +1,56 @@ +/* + * 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. + */ + +package org.apache.iotdb.db.protocol.rest.handler; + +import org.apache.iotdb.db.conf.rest.IoTDBRestServiceConfig; +import org.apache.iotdb.db.conf.rest.IoTDBRestServiceDescriptor; +import org.apache.iotdb.db.protocol.rest.exception.RequestLimitExceededException; + +public class RequestLimitChecker { + + private RequestLimitChecker() {} + + public static void checkRowCount(String requestName, int rowCount) { + int maxRows = getConfig().getRestMaxInsertRows(); + if (maxRows > 0 && rowCount > maxRows) { + throw new RequestLimitExceededException( + String.format("%s row count %d exceeds limit %d", requestName, rowCount, maxRows)); + } + } + + public static void checkColumnCount(String requestName, int columnCount) { + int maxColumns = getConfig().getRestMaxInsertColumns(); + if (maxColumns > 0 && columnCount > maxColumns) { + throw new RequestLimitExceededException( + String.format( + "%s column count %d exceeds limit %d", requestName, columnCount, maxColumns)); + } + } + + public static void checkValueCount(String requestName, long valueCount) { + long maxValues = getConfig().getRestMaxInsertValues(); + if (maxValues > 0 && valueCount > maxValues) { + throw new RequestLimitExceededException( + String.format("%s value count %d exceeds limit %d", requestName, valueCount, maxValues)); + } + } + + private static IoTDBRestServiceConfig getConfig() { + return IoTDBRestServiceDescriptor.getInstance().getConfig(); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/handler/ExceptionHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/handler/ExceptionHandler.java index 862f799961d..af3cea191e2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/handler/ExceptionHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/handler/ExceptionHandler.java @@ -26,6 +26,7 @@ import org.apache.iotdb.db.exception.metadata.DatabaseNotSetException; import org.apache.iotdb.db.exception.query.QueryProcessException; import org.apache.iotdb.db.exception.sql.SemanticException; import org.apache.iotdb.db.exception.sql.StatementAnalyzeException; +import org.apache.iotdb.db.protocol.rest.exception.RequestLimitExceededException; import org.apache.iotdb.db.protocol.rest.v1.model.ExecutionStatus; import org.apache.iotdb.rpc.TSStatusCode; @@ -44,7 +45,10 @@ public class ExceptionHandler { public static ExecutionStatus tryCatchException(Exception e) { ExecutionStatus responseResult = new ExecutionStatus(); - if (e instanceof QueryProcessException) { + if (e instanceof RequestLimitExceededException) { + responseResult.setMessage(e.getMessage()); + responseResult.setCode(413); + } else if (e instanceof QueryProcessException) { responseResult.setMessage(e.getMessage()); responseResult.setCode(((QueryProcessException) e).getErrorCode()); } else if (e instanceof DatabaseNotSetException) { @@ -84,4 +88,8 @@ public class ExceptionHandler { LOGGER.warn(e.getMessage(), e); return responseResult; } + + public static int getHttpStatus(Exception e) { + return e instanceof RequestLimitExceededException ? 413 : Status.OK.getStatusCode(); + } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/handler/RequestValidationHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/handler/RequestValidationHandler.java index 32e97d36fec..4e72646aaad 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/handler/RequestValidationHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/handler/RequestValidationHandler.java @@ -17,12 +17,14 @@ package org.apache.iotdb.db.protocol.rest.v1.handler; +import org.apache.iotdb.db.protocol.rest.handler.RequestLimitChecker; import org.apache.iotdb.db.protocol.rest.v1.model.ExpressionRequest; import org.apache.iotdb.db.protocol.rest.v1.model.InsertTabletRequest; import org.apache.iotdb.db.protocol.rest.v1.model.SQL; import org.apache.commons.lang3.Validate; +import java.util.List; import java.util.Objects; public class RequestValidationHandler { @@ -40,8 +42,30 @@ public class RequestValidationHandler { Objects.requireNonNull(insertTabletRequest.getTimestamps(), "timestamps should not be null"); Objects.requireNonNull(insertTabletRequest.getIsAligned(), "isAligned should not be null"); Objects.requireNonNull(insertTabletRequest.getDeviceId(), "deviceId should not be null"); + Objects.requireNonNull( + insertTabletRequest.getMeasurements(), "measurements should not be null"); Objects.requireNonNull(insertTabletRequest.getDataTypes(), "dataTypes should not be null"); Objects.requireNonNull(insertTabletRequest.getValues(), "values should not be null"); + + if (insertTabletRequest.getMeasurements().size() != insertTabletRequest.getDataTypes().size()) { + throw new IllegalArgumentException("measurements and dataTypes should have the same size"); + } + if (insertTabletRequest.getValues().size() != insertTabletRequest.getDataTypes().size()) { + throw new IllegalArgumentException("values and dataTypes should have the same size"); + } + + int rowCount = insertTabletRequest.getTimestamps().size(); + int columnCount = insertTabletRequest.getMeasurements().size(); + RequestLimitChecker.checkRowCount("insertTablet request", rowCount); + RequestLimitChecker.checkColumnCount("insertTablet request", columnCount); + RequestLimitChecker.checkValueCount("insertTablet request", (long) rowCount * columnCount); + + for (List<Object> column : insertTabletRequest.getValues()) { + if (column.size() != rowCount) { + throw new IllegalArgumentException( + "each value column should have the same size as timestamps"); + } + } } public static void validateExpressionRequest(ExpressionRequest expressionRequest) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/impl/RestApiServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/impl/RestApiServiceImpl.java index c50baa9bb54..b5361cd8be6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/impl/RestApiServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/impl/RestApiServiceImpl.java @@ -285,7 +285,9 @@ public class RestApiServiceImpl extends RestApiService { .message(result.status.getMessage())) .build(); } catch (Exception e) { - return Response.ok().entity(ExceptionHandler.tryCatchException(e)).build(); + return Response.status(ExceptionHandler.getHttpStatus(e)) + .entity(ExceptionHandler.tryCatchException(e)) + .build(); } finally { long costTime = System.nanoTime() - startTime; Optional.ofNullable(insertTabletStatement) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v2/handler/ExceptionHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v2/handler/ExceptionHandler.java index 6e93b7b55f0..345706d726d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v2/handler/ExceptionHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v2/handler/ExceptionHandler.java @@ -26,6 +26,7 @@ import org.apache.iotdb.db.exception.metadata.DatabaseNotSetException; import org.apache.iotdb.db.exception.query.QueryProcessException; import org.apache.iotdb.db.exception.sql.SemanticException; import org.apache.iotdb.db.exception.sql.StatementAnalyzeException; +import org.apache.iotdb.db.protocol.rest.exception.RequestLimitExceededException; import org.apache.iotdb.db.protocol.rest.model.ExecutionStatus; import org.apache.iotdb.rpc.TSStatusCode; @@ -44,7 +45,10 @@ public class ExceptionHandler { public static ExecutionStatus tryCatchException(Exception e) { ExecutionStatus responseResult = new ExecutionStatus(); - if (e instanceof QueryProcessException) { + if (e instanceof RequestLimitExceededException) { + responseResult.setMessage(e.getMessage()); + responseResult.setCode(413); + } else if (e instanceof QueryProcessException) { responseResult.setMessage(e.getMessage()); responseResult.setCode(((QueryProcessException) e).getErrorCode()); } else if (e instanceof DatabaseNotSetException) { @@ -84,4 +88,8 @@ public class ExceptionHandler { LOGGER.warn(e.getMessage(), e); return responseResult; } + + public static int getHttpStatus(Exception e) { + return e instanceof RequestLimitExceededException ? 413 : Status.OK.getStatusCode(); + } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v2/handler/RequestValidationHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v2/handler/RequestValidationHandler.java index c1fc9c6327f..87a3ef27485 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v2/handler/RequestValidationHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v2/handler/RequestValidationHandler.java @@ -17,6 +17,7 @@ package org.apache.iotdb.db.protocol.rest.v2.handler; +import org.apache.iotdb.db.protocol.rest.handler.RequestLimitChecker; import org.apache.iotdb.db.protocol.rest.v2.model.ExpressionRequest; import org.apache.iotdb.db.protocol.rest.v2.model.InsertRecordsRequest; import org.apache.iotdb.db.protocol.rest.v2.model.InsertTabletRequest; @@ -56,6 +57,23 @@ public class RequestValidationHandler { Objects.requireNonNull( insertTabletRequest.getMeasurements(), "measurements should not be null"); Objects.requireNonNull(insertTabletRequest.getValues(), "values should not be null"); + if (insertTabletRequest.getMeasurements().size() != insertTabletRequest.getDataTypes().size()) { + throw new IllegalArgumentException("measurements and data_types should have the same size"); + } + if (insertTabletRequest.getValues().size() != insertTabletRequest.getDataTypes().size()) { + throw new IllegalArgumentException("values and data_types should have the same size"); + } + int rowCount = insertTabletRequest.getTimestamps().size(); + int columnCount = insertTabletRequest.getMeasurements().size(); + RequestLimitChecker.checkRowCount("insertTablet request", rowCount); + RequestLimitChecker.checkColumnCount("insertTablet request", columnCount); + RequestLimitChecker.checkValueCount("insertTablet request", (long) rowCount * columnCount); + for (List<Object> column : insertTabletRequest.getValues()) { + if (column.size() != rowCount) { + throw new IllegalArgumentException( + "each value column should have the same size as timestamps"); + } + } List<String> errorMessages = new ArrayList<>(); String device = insertTabletRequest.getDevice(); for (int i = 0; i < insertTabletRequest.getMeasurements().size(); i++) { @@ -80,10 +98,30 @@ public class RequestValidationHandler { Objects.requireNonNull(insertRecordsRequest.getValuesList(), "values_list should not be null"); Objects.requireNonNull( insertRecordsRequest.getMeasurementsList(), "measurements_list should not be null"); + int rowCount = insertRecordsRequest.getDevices().size(); + if (insertRecordsRequest.getTimestamps().size() != rowCount + || insertRecordsRequest.getMeasurementsList().size() != rowCount + || insertRecordsRequest.getDataTypesList().size() != rowCount + || insertRecordsRequest.getValuesList().size() != rowCount) { + throw new IllegalArgumentException( + "devices, timestamps, measurements_list, data_types_list and values_list should have " + + "the same size"); + } + RequestLimitChecker.checkRowCount("insertRecords request", rowCount); List<String> errorMessages = new ArrayList<>(); + long valueCount = 0; for (int i = 0; i < insertRecordsRequest.getDataTypesList().size(); i++) { String device = insertRecordsRequest.getDevices().get(i); List<String> measurements = insertRecordsRequest.getMeasurementsList().get(i); + List<String> dataTypes = insertRecordsRequest.getDataTypesList().get(i); + List<Object> values = insertRecordsRequest.getValuesList().get(i); + if (measurements.size() != dataTypes.size() || values.size() != dataTypes.size()) { + throw new IllegalArgumentException( + "each insertRecords row should have the same number of measurements, data_types and " + + "values"); + } + RequestLimitChecker.checkColumnCount("insertRecords request", measurements.size()); + valueCount += values.size(); for (int c = 0; c < insertRecordsRequest.getDataTypesList().get(i).size(); c++) { String dataType = insertRecordsRequest.getDataTypesList().get(i).get(c); String measurement = measurements.get(c); @@ -93,6 +131,7 @@ public class RequestValidationHandler { } } } + RequestLimitChecker.checkValueCount("insertRecords request", valueCount); if (!errorMessages.isEmpty()) { throw new RuntimeException(String.join(",", errorMessages)); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v2/impl/RestApiServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v2/impl/RestApiServiceImpl.java index 80e86c1a998..cc81600c8e7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v2/impl/RestApiServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v2/impl/RestApiServiceImpl.java @@ -367,7 +367,9 @@ public class RestApiServiceImpl extends RestApiService { return responseGenerateHelper(result); } catch (Exception e) { - return Response.ok().entity(ExceptionHandler.tryCatchException(e)).build(); + return Response.status(ExceptionHandler.getHttpStatus(e)) + .entity(ExceptionHandler.tryCatchException(e)) + .build(); } finally { long costTime = System.nanoTime() - startTime; Optional.ofNullable(insertRowsStatement) @@ -421,7 +423,9 @@ public class RestApiServiceImpl extends RestApiService { false); return responseGenerateHelper(result); } catch (Exception e) { - return Response.ok().entity(ExceptionHandler.tryCatchException(e)).build(); + return Response.status(ExceptionHandler.getHttpStatus(e)) + .entity(ExceptionHandler.tryCatchException(e)) + .build(); } finally { long costTime = System.nanoTime() - startTime; Optional.ofNullable(insertTabletStatement) diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/rest/IoTDBRestServiceDescriptorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/rest/IoTDBRestServiceDescriptorTest.java new file mode 100644 index 00000000000..bc424deba8f --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/rest/IoTDBRestServiceDescriptorTest.java @@ -0,0 +1,91 @@ +/* + * 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. + */ + +package org.apache.iotdb.db.conf.rest; + +import org.apache.iotdb.commons.conf.TrimProperties; + +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +public class IoTDBRestServiceDescriptorTest { + + private IoTDBRestServiceConfig config; + private long originalMaxRequestBodySizeInBytes; + private long originalMaxTotalConcurrentRequestBodySizeInBytes; + private int originalMaxInsertRows; + private int originalMaxInsertColumns; + private long originalMaxInsertValues; + + @Before + public void setUp() { + config = IoTDBRestServiceDescriptor.getInstance().getConfig(); + originalMaxRequestBodySizeInBytes = config.getRestMaxRequestBodySizeInBytes(); + originalMaxTotalConcurrentRequestBodySizeInBytes = + config.getRestMaxTotalConcurrentRequestBodySizeInBytes(); + originalMaxInsertRows = config.getRestMaxInsertRows(); + originalMaxInsertColumns = config.getRestMaxInsertColumns(); + originalMaxInsertValues = config.getRestMaxInsertValues(); + } + + @After + public void tearDown() { + config.setRestMaxRequestBodySizeInBytes(originalMaxRequestBodySizeInBytes); + config.setRestMaxTotalConcurrentRequestBodySizeInBytes( + originalMaxTotalConcurrentRequestBodySizeInBytes); + config.setRestMaxInsertRows(originalMaxInsertRows); + config.setRestMaxInsertColumns(originalMaxInsertColumns); + config.setRestMaxInsertValues(originalMaxInsertValues); + } + + @Test + public void testDefaultRequestBodyMemoryLimit() { + Assert.assertEquals( + Runtime.getRuntime().maxMemory() / 20, + IoTDBRestServiceDescriptor.calculateRequestBodyMemoryLimitInBytes(new TrimProperties())); + } + + @Test + public void testRequestBodyMemoryLimitUsesHalfOfFreeMemory() { + TrimProperties properties = new TrimProperties(); + properties.setProperty("datanode_memory_proportion", "1:1:1:1:1:5"); + + Assert.assertEquals( + Runtime.getRuntime().maxMemory() / 4, + IoTDBRestServiceDescriptor.calculateRequestBodyMemoryLimitInBytes(properties)); + } + + @Test + public void testHotReloadRequestLimits() { + TrimProperties properties = new TrimProperties(); + properties.setProperty("rest_max_request_body_size_in_bytes", "101"); + properties.setProperty("rest_max_total_concurrent_request_body_size_in_bytes", "102"); + properties.setProperty("rest_max_insert_rows", "103"); + properties.setProperty("rest_max_insert_columns", "104"); + properties.setProperty("rest_max_insert_values", "105"); + + IoTDBRestServiceDescriptor.getInstance().loadHotModifiedProps(properties); + + Assert.assertEquals(101, config.getRestMaxRequestBodySizeInBytes()); + Assert.assertEquals(102, config.getRestMaxTotalConcurrentRequestBodySizeInBytes()); + Assert.assertEquals(103, config.getRestMaxInsertRows()); + Assert.assertEquals(104, config.getRestMaxInsertColumns()); + Assert.assertEquals(105, config.getRestMaxInsertValues()); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/rest/filter/RequestSizeLimitFilterTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/rest/filter/RequestSizeLimitFilterTest.java new file mode 100644 index 00000000000..74a1eb702f1 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/rest/filter/RequestSizeLimitFilterTest.java @@ -0,0 +1,286 @@ +/* + * 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. + */ + +package org.apache.iotdb.db.protocol.rest.filter; + +import org.apache.iotdb.db.conf.rest.IoTDBRestServiceConfig; +import org.apache.iotdb.db.conf.rest.IoTDBRestServiceDescriptor; +import org.apache.iotdb.db.protocol.rest.model.ExecutionStatus; + +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +import javax.ws.rs.WebApplicationException; +import javax.ws.rs.container.ContainerRequestContext; +import javax.ws.rs.container.ContainerResponseContext; +import javax.ws.rs.core.MediaType; +import javax.ws.rs.core.Response; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.lang.reflect.Proxy; +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; + +public class RequestSizeLimitFilterTest { + + private IoTDBRestServiceConfig config; + private long originalMaxBodySize; + private long originalMaxTotalConcurrentRequestBodySize; + + @Before + public void setUp() { + config = IoTDBRestServiceDescriptor.getInstance().getConfig(); + originalMaxBodySize = config.getRestMaxRequestBodySizeInBytes(); + originalMaxTotalConcurrentRequestBodySize = + config.getRestMaxTotalConcurrentRequestBodySizeInBytes(); + RestRequestBodyMemoryManager.resetForTest(); + } + + @After + public void tearDown() { + config.setRestMaxRequestBodySizeInBytes(originalMaxBodySize); + config.setRestMaxTotalConcurrentRequestBodySizeInBytes( + originalMaxTotalConcurrentRequestBodySize); + RestRequestBodyMemoryManager.resetForTest(); + } + + @Test + public void testAbortContentLengthOverLimit() { + config.setRestMaxRequestBodySizeInBytes(4); + config.setRestMaxTotalConcurrentRequestBodySizeInBytes(10); + TestRequestContext context = TestRequestContext.withLength(5); + + new RequestSizeLimitFilter().filter(context.proxy()); + + assertPayloadTooLarge(context.abortedResponse(), 4); + } + + @Test + public void testRejectStreamOverLimit() throws IOException { + config.setRestMaxRequestBodySizeInBytes(4); + config.setRestMaxTotalConcurrentRequestBodySizeInBytes(10); + TestRequestContext context = TestRequestContext.withStream("12345"); + + new RequestSizeLimitFilter().filter(context.proxy()); + + Assert.assertNull(context.abortedResponse()); + try { + consume(context.entityStream()); + Assert.fail("Expected WebApplicationException"); + } catch (WebApplicationException e) { + assertPayloadTooLarge(e.getResponse(), 4); + } + Assert.assertEquals(0, RestRequestBodyMemoryManager.getReservedMemoryInBytes()); + } + + @Test + public void testAbortContentLengthOverMemoryLimit() { + config.setRestMaxRequestBodySizeInBytes(10); + config.setRestMaxTotalConcurrentRequestBodySizeInBytes(4); + TestRequestContext context = TestRequestContext.withLength(5); + + new RequestSizeLimitFilter().filter(context.proxy()); + + assertMemoryQuotaExceeded(context.abortedResponse(), 4); + Assert.assertEquals(0, RestRequestBodyMemoryManager.getReservedMemoryInBytes()); + } + + @Test + public void testRejectConcurrentRequestsOverMemoryLimit() { + config.setRestMaxRequestBodySizeInBytes(10); + config.setRestMaxTotalConcurrentRequestBodySizeInBytes(4); + TestRequestContext firstContext = TestRequestContext.withLength(3); + TestRequestContext secondContext = TestRequestContext.withLength(2); + + new RequestSizeLimitFilter().filter(firstContext.proxy()); + new RequestSizeLimitFilter().filter(secondContext.proxy()); + + Assert.assertNull(firstContext.abortedResponse()); + assertMemoryQuotaExceeded(secondContext.abortedResponse(), 4); + Assert.assertEquals(3, RestRequestBodyMemoryManager.getReservedMemoryInBytes()); + + new RequestBodyMemoryReleaseFilter().filter(firstContext.proxy(), responseContext()); + + Assert.assertEquals(0, RestRequestBodyMemoryManager.getReservedMemoryInBytes()); + } + + @Test + public void testRejectStreamOverMemoryLimit() throws IOException { + config.setRestMaxRequestBodySizeInBytes(10); + config.setRestMaxTotalConcurrentRequestBodySizeInBytes(4); + TestRequestContext context = TestRequestContext.withStream("12345"); + + new RequestSizeLimitFilter().filter(context.proxy()); + + Assert.assertNull(context.abortedResponse()); + try { + consume(context.entityStream()); + Assert.fail("Expected WebApplicationException"); + } catch (WebApplicationException e) { + assertMemoryQuotaExceeded(e.getResponse(), 4); + } + Assert.assertEquals(0, RestRequestBodyMemoryManager.getReservedMemoryInBytes()); + } + + @Test + public void testRejectSkippedStreamOverLimit() throws IOException { + config.setRestMaxRequestBodySizeInBytes(4); + config.setRestMaxTotalConcurrentRequestBodySizeInBytes(10); + TestRequestContext context = TestRequestContext.withStream("12345"); + + new RequestSizeLimitFilter().filter(context.proxy()); + + try { + context.entityStream().skip(5); + Assert.fail("Expected WebApplicationException"); + } catch (WebApplicationException e) { + assertPayloadTooLarge(e.getResponse(), 4); + } + Assert.assertEquals(0, RestRequestBodyMemoryManager.getReservedMemoryInBytes()); + } + + @Test + public void testReleaseMemoryOnResponse() { + config.setRestMaxRequestBodySizeInBytes(10); + config.setRestMaxTotalConcurrentRequestBodySizeInBytes(5); + TestRequestContext context = TestRequestContext.withLength(4); + + new RequestSizeLimitFilter().filter(context.proxy()); + + Assert.assertNull(context.abortedResponse()); + Assert.assertEquals(4, RestRequestBodyMemoryManager.getReservedMemoryInBytes()); + + new RequestBodyMemoryReleaseFilter().filter(context.proxy(), responseContext()); + + Assert.assertEquals(0, RestRequestBodyMemoryManager.getReservedMemoryInBytes()); + } + + @Test + public void testDisabledLimitsDoNotWrapStream() { + config.setRestMaxRequestBodySizeInBytes(-1); + config.setRestMaxTotalConcurrentRequestBodySizeInBytes(-1); + TestRequestContext context = TestRequestContext.withStream("12345"); + InputStream originalStream = context.entityStream(); + + new RequestSizeLimitFilter().filter(context.proxy()); + + Assert.assertSame(originalStream, context.entityStream()); + Assert.assertEquals(0, RestRequestBodyMemoryManager.getReservedMemoryInBytes()); + } + + private static void consume(InputStream inputStream) throws IOException { + byte[] buffer = new byte[8]; + while (inputStream.read(buffer) != -1) { + // consume the request body + } + } + + private static void assertPayloadTooLarge(Response response, long maxBodySize) { + Assert.assertEquals(413, response.getStatus()); + Assert.assertEquals(MediaType.APPLICATION_JSON_TYPE, response.getMediaType()); + Assert.assertTrue(response.getEntity() instanceof ExecutionStatus); + ExecutionStatus status = (ExecutionStatus) response.getEntity(); + Assert.assertEquals(Integer.valueOf(413), status.getCode()); + Assert.assertTrue(status.getMessage().contains(Long.toString(maxBodySize))); + } + + private static void assertMemoryQuotaExceeded(Response response, long memoryLimit) { + Assert.assertEquals(503, response.getStatus()); + Assert.assertEquals(MediaType.APPLICATION_JSON_TYPE, response.getMediaType()); + Assert.assertTrue(response.getEntity() instanceof ExecutionStatus); + ExecutionStatus status = (ExecutionStatus) response.getEntity(); + Assert.assertEquals(Integer.valueOf(503), status.getCode()); + Assert.assertTrue(status.getMessage().contains(Long.toString(memoryLimit))); + } + + private static ContainerResponseContext responseContext() { + return (ContainerResponseContext) + Proxy.newProxyInstance( + ContainerResponseContext.class.getClassLoader(), + new Class<?>[] {ContainerResponseContext.class}, + (proxy, method, args) -> { + throw new UnsupportedOperationException(method.getName()); + }); + } + + private static class TestRequestContext { + + private final int contentLength; + private final AtomicReference<InputStream> entityStream; + private final AtomicReference<Response> abortedResponse = new AtomicReference<>(); + private final Map<String, Object> properties = new HashMap<>(); + + private TestRequestContext(int contentLength, InputStream entityStream) { + this.contentLength = contentLength; + this.entityStream = new AtomicReference<>(entityStream); + } + + private static TestRequestContext withLength(int contentLength) { + return new TestRequestContext(contentLength, new ByteArrayInputStream(new byte[0])); + } + + private static TestRequestContext withStream(String body) { + return new TestRequestContext( + -1, new ByteArrayInputStream(body.getBytes(StandardCharsets.UTF_8))); + } + + private ContainerRequestContext proxy() { + return (ContainerRequestContext) + Proxy.newProxyInstance( + ContainerRequestContext.class.getClassLoader(), + new Class<?>[] {ContainerRequestContext.class}, + (proxy, method, args) -> { + switch (method.getName()) { + case "getLength": + return contentLength; + case "getEntityStream": + return entityStream.get(); + case "setEntityStream": + entityStream.set((InputStream) args[0]); + return null; + case "abortWith": + abortedResponse.set((Response) args[0]); + return null; + case "getProperty": + return properties.get((String) args[0]); + case "setProperty": + properties.put((String) args[0], args[1]); + return null; + case "removeProperty": + properties.remove((String) args[0]); + return null; + default: + throw new UnsupportedOperationException(method.getName()); + } + }); + } + + private InputStream entityStream() { + return entityStream.get(); + } + + private Response abortedResponse() { + return abortedResponse.get(); + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/rest/handler/RequestValidationLimitTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/rest/handler/RequestValidationLimitTest.java new file mode 100644 index 00000000000..86a1622d2a0 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/rest/handler/RequestValidationLimitTest.java @@ -0,0 +1,110 @@ +/* + * 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. + */ + +package org.apache.iotdb.db.protocol.rest.handler; + +import org.apache.iotdb.db.conf.rest.IoTDBRestServiceConfig; +import org.apache.iotdb.db.conf.rest.IoTDBRestServiceDescriptor; +import org.apache.iotdb.db.protocol.rest.exception.RequestLimitExceededException; +import org.apache.iotdb.db.protocol.rest.v1.model.InsertTabletRequest; +import org.apache.iotdb.db.protocol.rest.v2.model.InsertRecordsRequest; + +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.util.Arrays; +import java.util.Collections; + +public class RequestValidationLimitTest { + + private IoTDBRestServiceConfig config; + private int originalMaxInsertRows; + private int originalMaxInsertColumns; + private long originalMaxInsertValues; + + @Before + public void setUp() { + config = IoTDBRestServiceDescriptor.getInstance().getConfig(); + originalMaxInsertRows = config.getRestMaxInsertRows(); + originalMaxInsertColumns = config.getRestMaxInsertColumns(); + originalMaxInsertValues = config.getRestMaxInsertValues(); + } + + @After + public void tearDown() { + config.setRestMaxInsertRows(originalMaxInsertRows); + config.setRestMaxInsertColumns(originalMaxInsertColumns); + config.setRestMaxInsertValues(originalMaxInsertValues); + } + + @Test(expected = RequestLimitExceededException.class) + public void testV1InsertTabletRejectsTooManyRows() { + config.setRestMaxInsertRows(2); + + InsertTabletRequest request = new InsertTabletRequest(); + request.setDeviceId("root.sg.d1"); + request.setIsAligned(false); + request.setMeasurements(Collections.singletonList("s1")); + request.setDataTypes(Collections.singletonList("INT64")); + request.setTimestamps(Arrays.asList(1L, 2L, 3L)); + request.setValues(Collections.singletonList(Arrays.<Object>asList(1L, 2L, 3L))); + + org.apache.iotdb.db.protocol.rest.v1.handler.RequestValidationHandler + .validateInsertTabletRequest(request); + } + + @Test(expected = RequestLimitExceededException.class) + public void testV2InsertTabletRejectsTooManyColumns() { + config.setRestMaxInsertColumns(1); + + org.apache.iotdb.db.protocol.rest.v2.model.InsertTabletRequest request = + new org.apache.iotdb.db.protocol.rest.v2.model.InsertTabletRequest(); + request.setDevice("root.sg.d1"); + request.setIsAligned(false); + request.setMeasurements(Arrays.asList("s1", "s2")); + request.setDataTypes(Arrays.asList("INT64", "INT64")); + request.setTimestamps(Collections.singletonList(1L)); + request.setValues( + Arrays.asList( + Collections.<Object>singletonList(1L), Collections.<Object>singletonList(2L))); + + org.apache.iotdb.db.protocol.rest.v2.handler.RequestValidationHandler + .validateInsertTabletRequest(request); + } + + @Test(expected = RequestLimitExceededException.class) + public void testV2InsertRecordsRejectsTooManyValues() { + config.setRestMaxInsertRows(10); + config.setRestMaxInsertColumns(10); + config.setRestMaxInsertValues(2); + + InsertRecordsRequest request = new InsertRecordsRequest(); + request.setIsAligned(false); + request.setDevices(Arrays.asList("root.sg.d1", "root.sg.d2")); + request.setTimestamps(Arrays.asList(1L, 2L)); + request.setMeasurementsList( + Arrays.asList(Arrays.asList("s1", "s2"), Collections.singletonList("s1"))); + request.setDataTypesList( + Arrays.asList(Arrays.asList("INT64", "INT64"), Collections.singletonList("INT64"))); + request.setValuesList( + Arrays.asList(Arrays.<Object>asList(1L, 2L), Collections.<Object>singletonList(3L))); + + org.apache.iotdb.db.protocol.rest.v2.handler.RequestValidationHandler + .validateInsertRecordsRequest(request); + } +} diff --git a/iotdb-core/datanode/src/test/resources/iotdb-common.properties b/iotdb-core/datanode/src/test/resources/iotdb-common.properties index 95ae09870f1..c713332aa44 100644 --- a/iotdb-core/datanode/src/test/resources/iotdb-common.properties +++ b/iotdb-core/datanode/src/test/resources/iotdb-common.properties @@ -34,6 +34,21 @@ enable_rest_service=true # The request rowLimit/row_limit value cannot exceed this limit. # rest_query_default_row_size_limit=10000 +# Maximum REST request body size in bytes +# rest_max_request_body_size_in_bytes=16777216 + +# Maximum total REST request body size in bytes across concurrent requests +# rest_max_total_concurrent_request_body_size_in_bytes=0 + +# Maximum rows accepted by a single REST write request +# rest_max_insert_rows=100000 + +# Maximum columns accepted by a single REST write request +# rest_max_insert_columns=1024 + +# Maximum values accepted by a single REST write request +# rest_max_insert_values=1000000 + # is SSL enabled # enable_https=false diff --git a/iotdb-core/datanode/src/test/resources/iotdb-system.properties b/iotdb-core/datanode/src/test/resources/iotdb-system.properties index 6a19f218257..00082cba0a1 100644 --- a/iotdb-core/datanode/src/test/resources/iotdb-system.properties +++ b/iotdb-core/datanode/src/test/resources/iotdb-system.properties @@ -52,6 +52,21 @@ enable_rest_service=true # The request rowLimit/row_limit value cannot exceed this limit. # rest_query_default_row_size_limit=10000 +# Maximum REST request body size in bytes +# rest_max_request_body_size_in_bytes=16777216 + +# Maximum total REST request body size in bytes across concurrent requests +# rest_max_total_concurrent_request_body_size_in_bytes=0 + +# Maximum rows accepted by a single REST write request +# rest_max_insert_rows=100000 + +# Maximum columns accepted by a single REST write request +# rest_max_insert_columns=1024 + +# Maximum values accepted by a single REST write request +# rest_max_insert_values=1000000 + # is SSL enabled # enable_https=false diff --git a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template index 2bc03179bd1..cff101654cf 100644 --- a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template +++ b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template @@ -566,10 +566,37 @@ enable_swagger=false # The maximum row limit for REST and Grafana query responses. # The request rowLimit/row_limit value cannot exceed this limit. # A non-positive value is invalid and falls back to the default (10000); it no longer means unlimited. -# effectiveMode: restart +# effectiveMode: hot_reload # Datatype: int rest_query_default_row_size_limit=10000 +# Maximum REST request body size in bytes. Set to 0 or a negative value to disable the limit. +# effectiveMode: hot_reload +# Datatype: long +rest_max_request_body_size_in_bytes=16777216 + +# Maximum total in-flight REST request body size in bytes across concurrent requests. +# When set to 0, use half of the free memory from datanode_memory_proportion. +# Set to a negative value to disable the limit. +# effectiveMode: hot_reload +# Datatype: long +rest_max_total_concurrent_request_body_size_in_bytes=0 + +# Maximum rows accepted by a single REST write request. Set to 0 or a negative value to disable the limit. +# effectiveMode: hot_reload +# Datatype: int +rest_max_insert_rows=100000 + +# Maximum columns accepted by a single REST write request. Set to 0 or a negative value to disable the limit. +# effectiveMode: hot_reload +# Datatype: int +rest_max_insert_columns=1024 + +# Maximum values accepted by a single REST write request. Set to 0 or a negative value to disable the limit. +# effectiveMode: hot_reload +# Datatype: long +rest_max_insert_values=1000000 + # Is client authentication required # effectiveMode: restart # Datatype: boolean
