[iOS] Support per-frame WebUI's in WebState (#31259)

This change replaces WebState's `*WebUI` related methods with ones that allow storing a set of WebUI's such that WebUI's can create child frames that also load WebUI which can be necessary for untrusted frames.

BUILD.gn

@@ -80,6 +80,7 @@ if (!is_ios) {
 } else {
   brave_all_unit_tests_deps += [
     "//brave/ios/testing:ios_brave_unit_tests",
+    "//brave/ios/web:ios_brave_web_inttests",
     "//brave/ios/web:ios_brave_web_unittests",
   ]
 }

chromium_src/ios/web/web_state/web_state_impl.h

@@ -0,0 +1,24 @@
+// Copyright (c) 2025 The Brave Authors. All rights reserved.
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at https://mozilla.org/MPL/2.0/.
+
+#ifndef BRAVE_CHROMIUM_SRC_IOS_WEB_WEB_STATE_WEB_STATE_IMPL_H_
+#define BRAVE_CHROMIUM_SRC_IOS_WEB_WEB_STATE_WEB_STATE_IMPL_H_
+
+namespace web {
+class WebUIIOS;
+}
+
+// Exposes an API to obtain the main frame WebUI which is required for some
+// Brave WebUI implementations
+#define ClearWebUI                    \
+  ClearWebUI();                       \
+  web::WebUIIOS* GetMainFrameWebUI(); \
+  size_t GetWebUICountForTesting
+
+#include <ios/web/web_state/web_state_impl.h>  // IWYU pragma: export
+
+#undef ClearWebUI
+
+#endif  // BRAVE_CHROMIUM_SRC_IOS_WEB_WEB_STATE_WEB_STATE_IMPL_H_

chromium_src/ios/web/web_state/web_state_impl.mm

@@ -0,0 +1,20 @@
+// Copyright (c) 2025 The Brave Authors. All rights reserved.
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at https://mozilla.org/MPL/2.0/.
+
+#include <ios/web/web_state/web_state_impl.mm>
+
+namespace web {
+
+WebUIIOS* WebStateImpl::GetMainFrameWebUI() {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  return RealizedState()->GetMainFrameWebUI();
+}
+
+size_t WebStateImpl::GetWebUICountForTesting() {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  return RealizedState()->GetWebUICountForTesting();  // IN-TEST
+}
+
+}  // namespace web

chromium_src/ios/web/web_state/web_state_impl_realized_web_state.h

@@ -0,0 +1,41 @@
+// Copyright (c) 2025 The Brave Authors. All rights reserved.
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at https://mozilla.org/MPL/2.0/.
+
+#ifndef BRAVE_CHROMIUM_SRC_IOS_WEB_WEB_STATE_WEB_STATE_IMPL_REALIZED_WEB_STATE_H_
+#define BRAVE_CHROMIUM_SRC_IOS_WEB_WEB_STATE_WEB_STATE_IMPL_REALIZED_WEB_STATE_H_
+
+#include <map>
+
+#include "ios/web/web_state/web_state_impl.h"
+
+// The goal of this override is to introduce a new map of WebUIIOS so that
+// WebStateImpl can handle WebUI in multiple frames in one page. It introduces
+// a set of replacements for the WebUI method implementations to use said map
+//
+// The reason we need to allow multiple WebUI's per frame is because some of our
+// WebUI applications such as AI Chat use subframes to host chrome-untrusted
+// frames
+#define web_ui_                                                                \
+  web_ui_;                                                                     \
+  std::map<std::string, std::unique_ptr<web::WebUIIOS>, std::less<>> web_uis_; \
+  void TearDown_ChromiumImpl();                                                \
+  void CreateWebUI_ChromiumImpl(const GURL& url);                              \
+  void ClearWebUI_ChromiumImpl();                                              \
+  bool HasWebUI_ChromiumImpl() const;                                          \
+  void HandleWebUIMessage_ChromiumImpl(const GURL&, std::string_view,          \
+                                       const base::Value::List&)
+// Exposes an API to obtain the main frame WebUI which is required for some
+// Brave WebUI implementations
+#define ClearWebUI                    \
+  ClearWebUI();                       \
+  web::WebUIIOS* GetMainFrameWebUI(); \
+  size_t GetWebUICountForTesting
+
+#include <ios/web/web_state/web_state_impl_realized_web_state.h>  // IWYU pragma: export
+
+#undef ClearWebUI
+#undef web_ui_
+
+#endif  // BRAVE_CHROMIUM_SRC_IOS_WEB_WEB_STATE_WEB_STATE_IMPL_REALIZED_WEB_STATE_H_

chromium_src/ios/web/web_state/web_state_impl_realized_web_state.mm

@@ -0,0 +1,85 @@
+// Copyright (c) 2025 The Brave Authors. All rights reserved.
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at https://mozilla.org/MPL/2.0/.
+
+#include "ios/web/web_state/web_state_impl_realized_web_state.h"
+
+// Support for replacing the following methods below
+#define ClearWebUI ClearWebUI_ChromiumImpl
+#define CreateWebUI(url) CreateWebUI_ChromiumImpl(url)
+#define HandleWebUIMessage HandleWebUIMessage_ChromiumImpl
+#define HasWebUI HasWebUI_ChromiumImpl
+#define TearDown TearDown_ChromiumImpl
+
+#include <ios/web/web_state/web_state_impl_realized_web_state.mm>
+
+#undef TearDown
+#undef HasWebUI
+#undef HandleWebUIMessage
+#undef CreateWebUI
+#undef ClearWebUI
+
+namespace web {
+
+// These new implementations mimic their non-Brave counterparts in
+// web_state_impl_realized_web_state.mm but store and fetch from the newly added
+// map instead of the `web_ui_` field
+//
+// See web_state_impl_realized_web_state.h for reasons why we are replacing
+// these implementations
+
+void WebStateImpl::RealizedWebState::TearDown() {
+  // Call the original ClearWebUI which would have been replaced with
+  // ClearWebUI_ChromiumImpl
+  ClearWebUI();
+  TearDown_ChromiumImpl();
+}
+
+void WebStateImpl::RealizedWebState::CreateWebUI(const GURL& url) {
+  const std::string_view host = url.host_piece();
+  if (web_uis_.contains(host)) {
+    // Don't recreate WebUI for the same host. At the moment this is a required
+    // limitation as we don't have the neccessary info to associate a WebUI
+    // with a specific frame like desktop/android does.
+    return;
+  }
+  auto web_ui = CreateWebUIIOS(url);
+  if (web_ui) {
+    web_uis_.insert({std::string(host), std::move(web_ui)});
+  }
+}
+
+void WebStateImpl::RealizedWebState::ClearWebUI() {
+  web_uis_.clear();
+}
+
+void WebStateImpl::RealizedWebState::HandleWebUIMessage(
+    const GURL& source_url,
+    std::string_view message,
+    const base::Value::List& args) {
+  const std::string_view host = source_url.host_piece();
+  auto web_ui = web_uis_.find(host);
+  if (web_ui != web_uis_.end() && web_ui->second) {
+    web_ui->second->ProcessWebUIIOSMessage(source_url, message, args);
+  }
+}
+
+bool WebStateImpl::RealizedWebState::HasWebUI() const {
+  return !web_uis_.empty();
+}
+
+// Implements the new method we expose to get the main frame's WebUI
+WebUIIOS* WebStateImpl::RealizedWebState::GetMainFrameWebUI() {
+  if (web_uis_.empty()) {
+    return nullptr;
+  }
+  // The first WebUI created should always be the main frame
+  return web_uis_.begin()->second.get();
+}
+
+size_t WebStateImpl::RealizedWebState::GetWebUICountForTesting() {
+  return web_uis_.size();
+}
+
+}  // namespace web

ios/web/BUILD.gn

@@ -14,3 +14,13 @@ test("ios_brave_web_unittests") {
     "//brave/ios/web/webui:ios_web_webui_unittests",
   ]
 }
+
+test("ios_brave_web_inttests") {
+  testonly = true
+  deps = [
+    "//ios/web:run_all_unittests",
+
+    # Add individual test source_set targets here.
+    "//brave/ios/web/webui:ios_web_webui_inttests",
+  ]
+}

ios/web/test/BUILD.gn

@@ -0,0 +1,23 @@
+# Copyright (c) 2025 The Brave Authors. All rights reserved.
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this file,
+# You can obtain one at https://mozilla.org/MPL/2.0/.
+
+import("//brave/resources/brave_grit.gni")
+import("//tools/grit/repack.gni")
+
+repack("packed_resources") {
+  testonly = true
+  sources = [ "$root_gen_dir/brave/ios/web/test/test_resources.pak" ]
+  deps = [ ":resources" ]
+  output = "$target_gen_dir/resources.pak"
+  copy_data_to_bundle = true
+}
+
+brave_grit("resources") {
+  source = "test_resources.grd"
+  outputs = [
+    "grit/test_resources.h",
+    "test_resources.pak",
+  ]
+}

ios/web/test/data/webui_test.html

@@ -0,0 +1,7 @@
+<!doctype html>
+<html>
+<body>
+  WebUI page
+  <iframe src="testwebui://testwebui2"></iframe>
+</body>
+</html>

ios/web/test/data/webui_test_2.html

@@ -0,0 +1,6 @@
+<!doctype html>
+<html>
+<body>
+  WebUI 2 page
+</body>
+</html>

ios/web/test/test_resources.grd

@@ -0,0 +1,15 @@
+<?xml version="1.0" encoding="utf-8"?>
+<grit latest_public_release="0" current_release="1">
+  <outputs>
+    <output filename="grit/test_resources.h" type="rc_header">
+      <emit emit_type='prepend'></emit>
+    </output>
+    <output filename="test_resources.pak" type="data_package" />
+  </outputs>
+  <release seq="1">
+    <includes>
+      <include name="IDR_WEBUI_TEST_HTML" file="data/webui_test.html" type="BINDATA" />
+      <include name="IDR_WEBUI_TEST_HTML_2" file="data/webui_test_2.html" type="BINDATA" />
+    </includes>
+  </release>
+</grit>

ios/web/webui/BUILD.gn

@@ -48,3 +48,23 @@ source_set("ios_web_webui_unittests") {
     "//testing/gtest",
   ]
 }
+
+source_set("ios_web_webui_inttests") {
+  testonly = true
+  sources = [ "web_ui_inttest.mm" ]
+  deps = [
+    "//base",
+    "//base/test:test_support",
+    "//brave/ios/web/test:packed_resources",
+    "//brave/ios/web/test:resources",
+    "//brave/ios/web/webui",
+    "//ios/web/public/test",
+    "//ios/web/public/test:test_fixture",
+    "//ios/web/public/test:util",
+    "//ios/web/public/test/fakes",
+    "//ios/web/public/webui",
+    "//ios/web/test:test_constants",
+    "//ios/web/web_state:web_state_impl_header",
+    "//url",
+  ]
+}

ios/web/webui/DEPS

@@ -6,5 +6,7 @@ include_rules = [
   "+ios/chrome/browser/shared/model/url",
   "+ios/components/webui",
   "+ios/web/webui",
+  "+ios/web/web_state",
+  "+ios/web/test",
   "+services/network/public/mojom",
 ]

ios/web/webui/web_ui_inttest.mm

@@ -0,0 +1,149 @@
+// Copyright (c) 2025 The Brave Authors. All rights reserved.
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at https://mozilla.org/MPL/2.0/.
+
+#include <memory>
+#include <string>
+
+#include "base/functional/bind.h"
+#include "base/run_loop.h"
+#include "base/test/ios/wait_util.h"
+#include "brave/ios/web/test/grit/test_resources.h"
+#include "brave/ios/web/webui/brave_web_ui_ios_data_source.h"
+#include "ios/web/public/navigation/navigation_manager.h"
+#include "ios/web/public/test/navigation_test_util.h"
+#include "ios/web/public/test/web_test_with_web_state.h"
+#include "ios/web/public/test/web_view_content_test_util.h"
+#include "ios/web/public/test/web_view_interaction_test_util.h"
+#include "ios/web/public/webui/web_ui_ios_controller.h"
+#include "ios/web/public/webui/web_ui_ios_controller_factory.h"
+#include "ios/web/public/webui/web_ui_ios_data_source.h"
+#include "ios/web/test/test_url_constants.h"
+#include "ios/web/web_state/web_state_impl.h"
+#include "third_party/abseil-cpp/absl/strings/str_format.h"
+#include "url/gurl.h"
+#include "url/scheme_host_port.h"
+
+using base::test::ios::kWaitForPageLoadTimeout;
+using base::test::ios::WaitUntilConditionOrTimeout;
+using web::test::WaitForWebViewContainingText;
+using web::test::WaitForWebViewContainingTextInFrame;
+
+namespace web {
+
+namespace {
+
+// Hostname for test WebUI page.
+const char kTestWebUIURLHost[] = "testwebui";
+const char kTestWebUIURLHost2[] = "testwebui2";
+
+// Text present on the sample WebUI page.
+const char kWebUIPageText[] = "WebUI page";
+const char kWebUIPage2Text[] = "WebUI 2 page";
+
+// Controller for test WebUI.
+class TestUI : public WebUIIOSController {
+ public:
+  // Constructs controller from `web_ui` and `ui_handler` which will communicate
+  // with test WebUI page.
+  TestUI(WebUIIOS* web_ui, const std::string& host, int resource_id)
+      : WebUIIOSController(web_ui, host) {
+    // Need to use a BraveWebUIIOSDataSource instead of standard
+    // WebUIIOSDataSource because we the test html has an iframe pointing to
+    // WebUI and so we need to override the CSP
+    BraveWebUIIOSDataSource* source = BraveWebUIIOSDataSource::Create(host);
+
+    source->SetDefaultResource(resource_id);
+    source->OverrideContentSecurityPolicy(
+        network::mojom::CSPDirectiveName::FrameSrc,
+        absl::StrFormat("frame-src %s://%s", kTestWebUIScheme,
+                        kTestWebUIURLHost2));
+
+    web::WebState* web_state = web_ui->GetWebState();
+    web::WebUIIOSDataSource::Add(web_state->GetBrowserState(), source);
+  }
+
+  ~TestUI() override = default;
+};
+
+// Factory that creates TestUI controller.
+class TestWebUIControllerFactory : public WebUIIOSControllerFactory {
+ public:
+  // Constructs a controller factory.
+  TestWebUIControllerFactory() {}
+
+  // WebUIIOSControllerFactory overrides.
+  std::unique_ptr<WebUIIOSController> CreateWebUIIOSControllerForURL(
+      WebUIIOS* web_ui,
+      const GURL& url) const override {
+    if (!url.SchemeIs(kTestWebUIScheme)) {
+      return nullptr;
+    }
+    if (url.host() == kTestWebUIURLHost) {
+      return std::make_unique<TestUI>(web_ui, url.host(), IDR_WEBUI_TEST_HTML);
+    }
+    DCHECK_EQ(url.host(), kTestWebUIURLHost2);
+    return std::make_unique<TestUI>(web_ui, url.host(), IDR_WEBUI_TEST_HTML_2);
+  }
+
+  NSInteger GetErrorCodeForWebUIURL(const GURL& url) const override {
+    if (url.SchemeIs(kTestWebUIScheme)) {
+      return 0;
+    }
+    return NSURLErrorUnsupportedURL;
+  }
+};
+}  // namespace
+
+// A test fixture for verifying WebUI. This test fixture is copied from
+// //ios/web/webui/web_ui_inttest.mm
+class WebUITest : public WebTestWithWebState {
+ protected:
+  WebUITest() : WebTestWithWebState() {}
+
+  void SetUp() override {
+    WebTestWithWebState::SetUp();
+    factory_ = std::make_unique<TestWebUIControllerFactory>();
+    WebUIIOSControllerFactory::RegisterFactory(factory_.get());
+
+    url::SchemeHostPort tuple(kTestWebUIScheme, kTestWebUIURLHost, 0);
+    GURL url(tuple.Serialize());
+    test::LoadUrl(web_state(), url);
+
+    // LoadIfNecessary is needed because the view is not created (but needed)
+    // when loading the page. TODO(crbug.com/41309809): Remove this call.
+    web_state()->GetNavigationManager()->LoadIfNecessary();
+
+    ASSERT_TRUE(WaitUntilConditionOrTimeout(kWaitForPageLoadTimeout, ^{
+      base::RunLoop().RunUntilIdle();
+      return !web_state()->IsLoading();
+    }));
+
+    ASSERT_EQ(url.spec(), BaseUrl());
+  }
+
+  void TearDown() override {
+    WebUIIOSControllerFactory::DeregisterFactory(factory_.get());
+    WebTestWithWebState::TearDown();
+  }
+
+ private:
+  std::unique_ptr<TestWebUIControllerFactory> factory_;
+};
+
+// Tests that a both the main web UI page and its child frame also containing
+// WebUI both load correctly and that WebState is holding onto both of the
+// WebUIIOS references
+TEST_F(WebUITest, LoadWebUIPageWithWebUIChildFrame) {
+  auto* web_state_impl = WebStateImpl::FromWebState(web_state());
+  ASSERT_TRUE(web_state_impl->HasWebUI());
+  EXPECT_TRUE(WaitForWebViewContainingText(web_state(), kWebUIPageText));
+  ASSERT_TRUE(web_state_impl->GetMainFrameWebUI() != nullptr);
+  EXPECT_TRUE(
+      WaitForWebViewContainingTextInFrame(web_state(), kWebUIPage2Text));
+  size_t expected_web_ui_count = 2;
+  ASSERT_EQ(web_state_impl->GetWebUICountForTesting(), expected_web_ui_count);
+}
+
+}  // namespace web

resources/resource_ids.spec

@@ -218,6 +218,10 @@
     "META": {"sizes": {"includes": [25]}},
     "includes": [30580],
   },
+  "brave/ios/web/test/test_resources.grd": {
+    "META": {"sizes": {"includes": [10]}},
+    "includes": [30600],
+  },
   # WARNING: The upstream ChromeOS/Ash strings currently run through 36930. We
   # must be careful not to exceed that maximum when adding new strings here.
 }