From d83e662cb5a738f78eba168ae4d064dfa2b2c8d1 Mon Sep 17 00:00:00 2001
From: tanzhengwei <tanzhengwei@google.com>
Date: Wed, 9 Dec 2020 01:56:51 -0800
Subject: [PATCH] Add validation of scope for GraphlessBlazeQueryEnvironment.

PiperOrigin-RevId: 346504986
---
 .../common/AbstractBlazeQueryEnvironment.java    | 16 ++++++++++++++++
 .../lib/query2/query/BlazeQueryEnvironment.java  | 14 +-------------
 .../query/GraphlessBlazeQueryEnvironment.java    |  5 ++++-
 3 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/query2/common/AbstractBlazeQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/common/AbstractBlazeQueryEnvironment.java
index 0dea641b92b748..2e49f801cf59b0 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/common/AbstractBlazeQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/common/AbstractBlazeQueryEnvironment.java
@@ -16,6 +16,7 @@
 import static java.util.concurrent.TimeUnit.MINUTES;
 
 import com.google.common.base.Preconditions;
+import com.google.common.base.Predicates;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
@@ -45,6 +46,7 @@
 import com.google.devtools.build.lib.util.DetailedExitCode;
 import java.io.IOException;
 import java.util.Collection;
+import java.util.Iterator;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -277,6 +279,20 @@ public Map<Label, Target> getTargets(Iterable<Label> labels)
     return resultBuilder.build();
   }
 
+  protected void validateScopeOfTargets(Set<Target> targets) throws QueryException {
+    // Sets.filter would be more convenient here, but can't deal with exceptions.
+    if (labelFilter != Predicates.<Label>alwaysTrue()) {
+      // The labelFilter is always true for bazel query; it's only used for genquery rules.
+      Iterator<Target> targetIterator = targets.iterator();
+      while (targetIterator.hasNext()) {
+        Target target = targetIterator.next();
+        if (!validateScope(target.getLabel(), strictScope)) {
+          targetIterator.remove();
+        }
+      }
+    }
+  }
+
   protected boolean validateScope(Label label, boolean strict) throws QueryException {
     if (!labelFilter.test(label)) {
       String error = String.format("target '%s' is not within the scope of the query", label);
diff --git a/src/main/java/com/google/devtools/build/lib/query2/query/BlazeQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/query/BlazeQueryEnvironment.java
index 05fc3bd43fab9d..8312c340786458 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/query/BlazeQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/query/BlazeQueryEnvironment.java
@@ -17,7 +17,6 @@
 
 import com.google.common.base.Preconditions;
 import com.google.common.base.Predicate;
-import com.google.common.base.Predicates;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
@@ -63,7 +62,6 @@
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -172,17 +170,7 @@ private void getTargetsMatchingPatternImpl(String pattern, Callback<Target> call
     // becomes quadratic in runtime.
     Set<Target> targets = new LinkedHashSet<>(resolvedTargetPatterns.get(pattern));
 
-    // Sets.filter would be more convenient here, but can't deal with exceptions.
-    if (labelFilter != Predicates.<Label>alwaysTrue()) {
-      // The labelFilter is always true for bazel query; it's only used for genquery rules.
-      Iterator<Target> targetIterator = targets.iterator();
-      while (targetIterator.hasNext()) {
-        Target target = targetIterator.next();
-        if (!validateScope(target.getLabel(), strictScope)) {
-          targetIterator.remove();
-        }
-      }
-    }
+    validateScopeOfTargets(targets);
 
     Set<PackageIdentifier> packages = CompactHashSet.create();
     for (Target target : targets) {
diff --git a/src/main/java/com/google/devtools/build/lib/query2/query/GraphlessBlazeQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/query/GraphlessBlazeQueryEnvironment.java
index 1b4ca951920387..82046d29ff1dd2 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/query/GraphlessBlazeQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/query/GraphlessBlazeQueryEnvironment.java
@@ -51,6 +51,7 @@
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -153,7 +154,9 @@ public QueryTaskFuture<Void> getTargetsMatchingPattern(
 
   private void getTargetsMatchingPatternImpl(String pattern, Callback<Target> callback)
       throws QueryException, InterruptedException {
-    callback.process(resolvedTargetPatterns.get(pattern));
+    Set<Target> targets = new LinkedHashSet<>(resolvedTargetPatterns.get(pattern));
+    validateScopeOfTargets(targets);
+    callback.process(targets);
   }
 
   @Override