1 files changed, 58 insertions, 1 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilder.java b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilder.java
index da43b6a2a4..9f2217cd62 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilder.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCompileActionBuilder.java
@@ -351,8 +351,10 @@ public class ProtoCompileActionBuilder {
             protosToCompile,
             transitiveSources,
             protosInDirectDeps,
+            /* protosInExports= */ null,
             protoSourceRoots,
             directProtoSourceRoots,
+            /* exportedProtoSourceRoots= */ null,
             ruleContext.getLabel(),
             ImmutableList.of(output),
             "Descriptor Set",
@@ -381,6 +383,34 @@ public class ProtoCompileActionBuilder {
         outReplacement);
   }
 
+  public static void registerActions(
+      RuleContext ruleContext,
+      List<ToolchainInvocation> toolchainInvocations,
+      Iterable<Artifact> protosToCompile,
+      NestedSet<Artifact> transitiveSources,
+      NestedSet<Artifact> protosInDirectDeps,
+      NestedSet<String> protoSourceRoots,
+      NestedSet<String> directProtoSourceRoots,
+      Label ruleLabel,
+      Iterable<Artifact> outputs,
+      String flavorName,
+      boolean allowServices) {
+    registerActions(
+        ruleContext,
+        toolchainInvocations,
+        protosToCompile,
+        transitiveSources,
+        protosInDirectDeps,
+        protoSourceRoots,
+        directProtoSourceRoots,
+        ruleLabel,
+        outputs,
+        flavorName,
+        allowServices,
+        /* protosInExports= */ null,
+        /* exportedProtoSourceRoots= */ null);
+  }
+
   /**
    * Registers actions to generate code from .proto files.
    *
@@ -404,7 +434,9 @@ public class ProtoCompileActionBuilder {
       Label ruleLabel,
       Iterable<Artifact> outputs,
       String flavorName,
-      boolean allowServices) {
+      boolean allowServices,
+      NestedSet<Artifact> protosInExports,
+      NestedSet<String> exportedProtoSourceRoots) {
     SpawnAction.Builder actions =
         createActions(
             ruleContext,
@@ -412,8 +444,10 @@ public class ProtoCompileActionBuilder {
             protosToCompile,
             transitiveSources,
             protosInDirectDeps,
+            protosInExports,
             protoSourceRoots,
             directProtoSourceRoots,
+            exportedProtoSourceRoots,
             ruleLabel,
             outputs,
             flavorName,
@@ -430,8 +464,10 @@ public class ProtoCompileActionBuilder {
       Iterable<Artifact> protosToCompile,
       NestedSet<Artifact> transitiveSources,
       @Nullable NestedSet<Artifact> protosInDirectDeps,
+      @Nullable NestedSet<Artifact> protosInExports,
       NestedSet<String> protoSourceRoots,
       NestedSet<String> directProtoSourceRoots,
+      @Nullable NestedSet<String> exportedProtoSourceRoots,
       Label ruleLabel,
       Iterable<Artifact> outputs,
       String flavorName,
@@ -468,7 +504,9 @@ public class ProtoCompileActionBuilder {
                 transitiveSources,
                 protoSourceRoots,
                 directProtoSourceRoots,
+                exportedProtoSourceRoots,
                 areDepsStrict(ruleContext) ? protosInDirectDeps : null,
+                arePublicImportsStrict(ruleContext) ? protosInExports : null,
                 ruleLabel,
                 allowServices,
                 ruleContext.getFragment(ProtoConfiguration.class).protocOpts()),
@@ -479,6 +517,10 @@ public class ProtoCompileActionBuilder {
     return result;
   }
 
+  public static boolean arePublicImportsStrict(RuleContext ruleContext) {
+    return ruleContext.getFragment(ProtoConfiguration.class).strictPublicImports();
+  }
+
   /**
    * Constructs command-line arguments to execute proto-compiler.
    *
@@ -506,7 +548,9 @@ public class ProtoCompileActionBuilder {
       NestedSet<Artifact> transitiveSources,
       NestedSet<String> transitiveProtoPathFlags,
       NestedSet<String> directProtoSourceRoots,
+      NestedSet<String> exportedProtoSourceRoots,
       @Nullable NestedSet<Artifact> protosInDirectDeps,
+      @Nullable NestedSet<Artifact> protosInExports,
       Label ruleLabel,
       boolean allowServices,
       ImmutableList<String> protocOpts) {
@@ -554,6 +598,19 @@ public class ProtoCompileActionBuilder {
       cmdLine.addFormatted(STRICT_DEPS_FLAG_TEMPLATE, ruleLabel);
     }
 
+    if (protosInExports != null) {
+      if (protosInExports.isEmpty()) {
+        // This line is necessary to trigger the check.
+        cmdLine.add("--allowed_public_imports=");
+      } else {
+        cmdLine.addAll(
+            "--allowed_public_imports",
+            VectorArg.join(":")
+                .each(protosInExports)
+                .mapped(new ExpandToPathFn(exportedProtoSourceRoots)));
+      }
+    }
+
     for (Artifact src : protosToCompile) {
       cmdLine.addPath(src.getExecPath());
     }