From 0aa46e4ee194dd63d9348a16a0b219e548839303 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Thu, 26 Oct 2023 16:19:07 -0700
Subject: [PATCH] Allow rec groups of public function types in closed world
 (#6053)

Closed-world mode allows function types to escape if they are on exported functions,
because that has been possible since wasm MVP and cannot be avoided. But we need to
also allow all types in those type's rec groups as well. Consider this case:

(module
 (rec
  (type $0 (func))
  (type $1 (func))
 )
 (func "0" (type $0)
  (nop)
 )
 (func "1" (type $1)
  (nop)
 )
)

The two exported functions make the two types public, so this module validates in
closed world mode. Now imagine that metadce removes one export:

(module
 (rec
  (type $0 (func))
  (type $1 (func))
 )
 (func "0" (type $0)
  (nop)
 )
 ;; The export "1" is gone.
)

Before this PR that no longer validates, because it only marks the type $0 as public.
But when a type is public that makes its entire rec group public, so $1 is errored on.

To fix that, this PR allows all types in a rec group of an exported function's type, which
makes that last module validate.
---
 src/wasm/wasm-validator.cpp                   | 22 ++++++++++++++-----
 .../validation/closed-world-interface.wast    | 11 +++++-----
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/wasm/wasm-validator.cpp b/src/wasm/wasm-validator.cpp
index 179235fbcc2..cd90e767eca 100644
--- a/src/wasm/wasm-validator.cpp
+++ b/src/wasm/wasm-validator.cpp
@@ -3734,13 +3734,23 @@ static void validateFeatures(Module& module, ValidationInfo& info) {
 
 static void validateClosedWorldInterface(Module& module, ValidationInfo& info) {
   // Error if there are any publicly exposed heap types beyond the types of
-  // publicly exposed functions.
-  std::unordered_set<HeapType> publicFuncTypes;
-  ModuleUtils::iterImportedFunctions(
-    module, [&](Function* func) { publicFuncTypes.insert(func->type); });
+  // publicly exposed functions. Note that we must include all types in the rec
+  // groups that are used, as if a type if public then all types in its rec
+  // group are as well.
+  std::unordered_set<RecGroup> publicRecGroups;
+  ModuleUtils::iterImportedFunctions(module, [&](Function* func) {
+    publicRecGroups.insert(func->type.getRecGroup());
+  });
   for (auto& ex : module.exports) {
     if (ex->kind == ExternalKind::Function) {
-      publicFuncTypes.insert(module.getFunction(ex->value)->type);
+      publicRecGroups.insert(module.getFunction(ex->value)->type.getRecGroup());
+    }
+  }
+
+  std::unordered_set<HeapType> publicTypes;
+  for (auto& group : publicRecGroups) {
+    for (auto type : group) {
+      publicTypes.insert(type);
     }
   }
 
@@ -3749,7 +3759,7 @@ static void validateClosedWorldInterface(Module& module, ValidationInfo& info) {
   auto ignorable = getIgnorablePublicTypes();
 
   for (auto type : ModuleUtils::getPublicHeapTypes(module)) {
-    if (!publicFuncTypes.count(type) && !ignorable.count(type)) {
+    if (!publicTypes.count(type) && !ignorable.count(type)) {
       auto name = type.toString();
       if (auto it = module.typeNames.find(type); it != module.typeNames.end()) {
         name = it->second.name.toString();
diff --git a/test/lit/validation/closed-world-interface.wast b/test/lit/validation/closed-world-interface.wast
index 06cf1cef715..daedaf990de 100644
--- a/test/lit/validation/closed-world-interface.wast
+++ b/test/lit/validation/closed-world-interface.wast
@@ -3,10 +3,6 @@
 ;; RUN: not wasm-opt -all --closed-world %s 2>&1 | filecheck %s
 
 
-;; This is pulled in because it is part of a rec group with $partial-pair-0.
-;; CHECK:      publicly exposed type disallowed with a closed world: $partial-pair-1, on
-;; CHECK-NEXT: (func)
-
 ;; This is pulled in by a global.
 ;; CHECK:      publicly exposed type disallowed with a closed world: $array, on
 ;; CHECK-NEXT: (array (mut i32))
@@ -32,8 +28,13 @@
     (type $exported-pair-1 (func (param (ref $exported-pair-0))))
   )
   (rec
+    ;; This is on an exported function.
     (type $partial-pair-0 (func))
+    ;; The latter type types are not public, but allowed to be because the
+    ;; entire rec group is allowed due to the first.
     (type $partial-pair-1 (func))
+    ;; Test a non-function type.
+    (type $partial-pair-2 (struct))
   )
 
   (type $private (func (param v128)))
@@ -61,7 +62,7 @@
   ;; Ok even though it is an import instead of an export.
   (func $5 (import "env" "test5") (type $exported-pair-1))
 
-  ;; Not ok because another type in the group is not on the boundary.
+  ;; Ok, and we also allow the other type in the group.
   (func $6 (export "test6") (type $partial-pair-0)
     (unreachable)
   )