Add REST write request size and batch limits

external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/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.rest.protocol.exception;
+
+public class RequestLimitExceededException extends IllegalArgumentException {
+
+  public RequestLimitExceededException(String message) {
+    super(message);
+  }
+}

external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/filter/RequestSizeLimitFilter.java

@@ -0,0 +1,98 @@
+/*
+ * 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.rest.protocol.filter;
+
+import org.apache.iotdb.db.conf.rest.IoTDBRestServiceDescriptor;
+
+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 {
+
+  @Override
+  public void filter(ContainerRequestContext requestContext) {
+    long maxBodySize =
+        IoTDBRestServiceDescriptor.getInstance().getConfig().getRestMaxRequestBodySizeInBytes();
+    if (maxBodySize <= 0) {
+      return;
+    }
+
+    int contentLength = requestContext.getLength();
+    if (contentLength > maxBodySize) {
+      requestContext.abortWith(buildPayloadTooLargeResponse(maxBodySize));
+      return;
+    }
+
+    requestContext.setEntityStream(
+        new LimitedInputStream(requestContext.getEntityStream(), maxBodySize));
+  }
+
+  private static Response buildPayloadTooLargeResponse(long maxBodySize) {
+    return Response.status(413)
+        .type(MediaType.TEXT_PLAIN_TYPE)
+        .entity("REST request body exceeds limit " + maxBodySize + " bytes")
+        .build();
+  }
+
+  private static class LimitedInputStream extends FilterInputStream {
+
+    private final long maxBodySize;
+    private long bytesRead;
+
+    private LimitedInputStream(InputStream in, long maxBodySize) {
+      super(in);
+      this.maxBodySize = maxBodySize;
+    }
+
+    @Override
+    public int read() throws IOException {
+      int result = super.read();
+      if (result != -1) {
+        incrementBytesRead(1);
+      }
+      return result;
+    }
+
+    @Override
+    public int read(byte[] b, int off, int len) throws IOException {
+      int result = super.read(b, off, len);
+      if (result > 0) {
+        incrementBytesRead(result);
+      }
+      return result;
+    }
+
+    private void incrementBytesRead(int increment) {
+      bytesRead += increment;
+      if (bytesRead > maxBodySize) {
+        throw new WebApplicationException(buildPayloadTooLargeResponse(maxBodySize));
+      }
+    }
+  }
+}

external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/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.rest.protocol.handler;
+
+import org.apache.iotdb.db.conf.rest.IoTDBRestServiceConfig;
+import org.apache.iotdb.db.conf.rest.IoTDBRestServiceDescriptor;
+import org.apache.iotdb.rest.protocol.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();
+  }
+}

external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/table/v1/handler/ExceptionHandler.java

@@ -28,6 +28,7 @@
 import org.apache.iotdb.db.exception.sql.SemanticException;
 import org.apache.iotdb.db.exception.sql.StatementAnalyzeException;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.parser.ParsingException;
+import org.apache.iotdb.rest.protocol.exception.RequestLimitExceededException;
 import org.apache.iotdb.rest.protocol.model.ExecutionStatus;
 import org.apache.iotdb.rpc.TSStatusCode;
 
@@ -46,7 +47,10 @@ private 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) {
@@ -92,4 +96,8 @@ public static ExecutionStatus tryCatchException(Exception e) {
     LOGGER.warn(e.getMessage(), e);
     return responseResult;
   }
+
+  public static int getHttpStatus(Exception e) {
+    return e instanceof RequestLimitExceededException ? 413 : Status.OK.getStatusCode();
+  }
 }

external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/table/v1/handler/RequestValidationHandler.java

@@ -17,6 +17,7 @@
 
 package org.apache.iotdb.rest.protocol.table.v1.handler;
 
+import org.apache.iotdb.rest.protocol.handler.RequestLimitChecker;
 import org.apache.iotdb.rest.protocol.table.v1.model.InsertTabletRequest;
 import org.apache.iotdb.rest.protocol.table.v1.model.SQL;
 
@@ -62,6 +63,13 @@
       errorMessages.add("values and timestamps should have the same size");
     }
 
+    int rowCount = insertTabletRequest.getTimestamps().size();
+    int columnCount = insertTabletRequest.getColumnNames().size();
+    RequestLimitChecker.checkRowCount("table insertTablet request", rowCount);
+    RequestLimitChecker.checkColumnCount("table insertTablet request", columnCount);
+    RequestLimitChecker.checkValueCount(
+        "table insertTablet request", (long) rowCount * columnCount);
+
     for (int i = 0; i < insertTabletRequest.getDataTypes().size(); i++) {
       String dataType = insertTabletRequest.getDataTypes().get(i);
       if (isDataType(dataType)) {

external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/table/v1/impl/RestApiServiceImpl.java

@@ -109,7 +109,9 @@ public Response executeQueryInternal(
         return res;
       }
     } catch (Exception e) {
-      return Response.ok().entity(ExceptionHandler.tryCatchException(e)).build();
+      return Response.status(ExceptionHandler.getHttpStatus(e))
+          .entity(ExceptionHandler.tryCatchException(e))
+          .build();
     } finally {
       if (queryId != null) {
         COORDINATOR.cleanupQueryExecution(queryId);
@@ -132,7 +134,9 @@ public Response executeQueryStatement(SQL sql, SecurityContext securityContext)
 
       return executeQueryInternal(sql, statement, clientSession, relationSqlParser);
     } 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(statement)
@@ -174,7 +178,9 @@ public Response insertTablet(
 
       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)
@@ -231,7 +237,9 @@ public Response executeNonQueryStatement(SQL sql, SecurityContext securityContex
       }
       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(statement)

external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/handler/ExceptionHandler.java

@@ -27,6 +27,7 @@
 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.rest.protocol.exception.RequestLimitExceededException;
 import org.apache.iotdb.rest.protocol.v1.model.ExecutionStatus;
 import org.apache.iotdb.rpc.TSStatusCode;
 
@@ -45,7 +46,10 @@ private 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) {
@@ -88,4 +92,8 @@ public static ExecutionStatus tryCatchException(Exception e) {
     LOGGER.warn(e.getMessage(), e);
     return responseResult;
   }
+
+  public static int getHttpStatus(Exception e) {
+    return e instanceof RequestLimitExceededException ? 413 : Status.OK.getStatusCode();
+  }
 }

external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/handler/RequestValidationHandler.java

@@ -17,12 +17,14 @@
 
 package org.apache.iotdb.rest.protocol.v1.handler;
 
+import org.apache.iotdb.rest.protocol.handler.RequestLimitChecker;
 import org.apache.iotdb.rest.protocol.v1.model.ExpressionRequest;
 import org.apache.iotdb.rest.protocol.v1.model.InsertTabletRequest;
 import org.apache.iotdb.rest.protocol.v1.model.SQL;
 
 import org.apache.tsfile.external.commons.lang3.Validate;
 
+import java.util.List;
 import java.util.Objects;
 
 public class RequestValidationHandler {
@@ -40,8 +42,30 @@
     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) {

external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/impl/GrafanaApiServiceImpl.java

@@ -133,7 +133,9 @@ public Response variables(SQL sql, SecurityContext securityContext) {
         return QueryDataSetHandler.fillGrafanaVariablesResult(queryExecution, statement);
       }
     } catch (Exception e) {
-      return Response.ok().entity(ExceptionHandler.tryCatchException(e)).build();
+      return Response.status(ExceptionHandler.getHttpStatus(e))
+          .entity(ExceptionHandler.tryCatchException(e))
+          .build();
     } finally {
       if (queryId != null) {
         COORDINATOR.cleanupQueryExecution(queryId);
@@ -206,7 +208,9 @@ public Response expression(ExpressionRequest expressionRequest, SecurityContext
         }
       }
     } catch (Exception e) {
-      return Response.ok().entity(ExceptionHandler.tryCatchException(e)).build();
+      return Response.status(ExceptionHandler.getHttpStatus(e))
+          .entity(ExceptionHandler.tryCatchException(e))
+          .build();
     } finally {
       if (queryId != null) {
         COORDINATOR.cleanupQueryExecution(queryId);
@@ -269,7 +273,9 @@ public Response node(List<String> requestBody, SecurityContext securityContext)
         return QueryDataSetHandler.fillGrafanaNodesResult(null);
       }
     } catch (Exception e) {
-      return Response.ok().entity(ExceptionHandler.tryCatchException(e)).build();
+      return Response.status(ExceptionHandler.getHttpStatus(e))
+          .entity(ExceptionHandler.tryCatchException(e))
+          .build();
     } finally {
       if (queryId != null) {
         COORDINATOR.cleanupQueryExecution(queryId);