Merge pull request #3711 from atlanhq/plt-2751-idx-limit

PLT-2751 : Abuse protection in Indexsearch endpoint: Put limit in indexsearch size upto 100000
atlanhq · Nov 11, 2024 · 8584f0f · 8584f0f
2 parents dd3dc3d + 02a7b1b
commit 8584f0f
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 0 deletions.
diff --git a/intg/src/main/java/org/apache/atlas/AtlasConfiguration.java b/intg/src/main/java/org/apache/atlas/AtlasConfiguration.java
@@ -113,6 +113,8 @@ public enum AtlasConfiguration {
     HERACLES_API_SERVER_URL("atlas.heracles.api.service.url", "http://heracles-service.heracles.svc.cluster.local"),
 
     INDEXSEARCH_ASYNC_SEARCH_KEEP_ALIVE_TIME_IN_SECONDS("atlas.indexsearch.async.search.keep.alive.time.in.seconds", 300),
+    ATLAS_INDEXSEARCH_QUERY_SIZE_MAX_LIMIT("atlas.indexsearch.query.size.max.limit", 100000),
+    ATLAS_INDEXSEARCH_LIMIT_UTM_TAGS("atlas.indexsearch.limit.ignore.utm.tags", "project_sdk_python"),
 
     ATLAS_MAINTENANCE_MODE("atlas.maintenance.mode", false),
 

diff --git a/intg/src/main/java/org/apache/atlas/AtlasErrorCode.java b/intg/src/main/java/org/apache/atlas/AtlasErrorCode.java
@@ -121,6 +121,7 @@ public enum AtlasErrorCode {
     INVALID_DSL_DUPLICATE_ALIAS(400, "ATLAS-400-00-066", "DSL Semantic Error - Duplicate alias found: '{0}' for type '{1}' already present."),
     INVALID_DSL_INVALID_DATE(400, "ATLAS-400-00-067", "DSL Semantic Error - Date format: {0}."),
     INVALID_DSL_HAS_PROPERTY(400, "ATLAS-400-00-068", "DSL Semantic Error - Property needs to be a primitive type: {0}"),
+    INVALID_DSL_QUERY_SIZE(400, "ATLAS-400-00-103", "DSL Error - Please provide query size less than {0}"),
     RELATIONSHIP_UPDATE_END_CHANGE_NOT_ALLOWED(404, "ATLAS-400-00-069", "change of relationship end is not permitted. relationship-type={0}, relationship-guid={1}, end-guid={2}, updated-end-guid={3}"),
     RELATIONSHIP_UPDATE_TYPE_CHANGE_NOT_ALLOWED(404, "ATLAS-400-00-06A", "change of relationship type is not permitted. relationship-guid={0}, current-type={1}, new-type={2}"),
     CLASSIFICATION_UPDATE_FROM_PROPAGATED_ENTITY(400, "ATLAS-400-00-06B", "Update to classification {0} is not allowed from propagated entity"),

diff --git a/intg/src/main/java/org/apache/atlas/model/discovery/IndexSearchParams.java b/intg/src/main/java/org/apache/atlas/model/discovery/IndexSearchParams.java
@@ -44,6 +44,10 @@ public void setDsl(Map dsl) {
         queryString = AtlasType.toJson(dsl);
     }
 
+    public long getQuerySize() {
+        return dsl.get("size") != null ? ((Number)dsl.get("size")).longValue() : 10;
+    }
+
     public boolean isAllowDeletedRelations() {
         return allowDeletedRelations;
     }

diff --git a/webapp/src/main/java/org/apache/atlas/web/rest/DiscoveryREST.java b/webapp/src/main/java/org/apache/atlas/web/rest/DiscoveryREST.java
@@ -67,6 +67,7 @@
 import java.util.List;
 import java.util.Set;
 import java.util.Arrays;
+
 import static org.apache.atlas.repository.Constants.QUALIFIED_NAME;
 import static org.apache.atlas.repository.Constants.REQUEST_HEADER_HOST;
 import static org.apache.atlas.repository.Constants.REQUEST_HEADER_USER_AGENT;
@@ -400,6 +401,17 @@ public AtlasSearchResult indexSearch(@Context HttpServletRequest servletRequest,
                 perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "DiscoveryREST.indexSearch(" + parameters + ")");
             }
 
+            if (parameters.getQuerySize() > AtlasConfiguration.ATLAS_INDEXSEARCH_QUERY_SIZE_MAX_LIMIT.getLong()) {
+                if(CollectionUtils.isEmpty(parameters.getUtmTags())) {
+                    throw new AtlasBaseException(AtlasErrorCode.INVALID_DSL_QUERY_SIZE, String.valueOf(AtlasConfiguration.ATLAS_INDEXSEARCH_QUERY_SIZE_MAX_LIMIT.getLong()));
+                }
+                for (String utmTag : parameters.getUtmTags()) {
+                    if (Arrays.stream(AtlasConfiguration.ATLAS_INDEXSEARCH_LIMIT_UTM_TAGS.getStringArray()).anyMatch(utmTag::equalsIgnoreCase)) {
+                            throw new AtlasBaseException(AtlasErrorCode.INVALID_DSL_QUERY_SIZE, String.valueOf(AtlasConfiguration.ATLAS_INDEXSEARCH_QUERY_SIZE_MAX_LIMIT.getLong()));
+                    }
+                }
+            }
+
             if (StringUtils.isEmpty(parameters.getQuery())) {
                 AtlasBaseException abe = new AtlasBaseException(AtlasErrorCode.BAD_REQUEST, "Invalid search query");
                 if (enableSearchLogging && parameters.isSaveSearchLog()) {