From 2e8d20c40faba1a2bc27c9218a23dd5ec7a64c61 Mon Sep 17 00:00:00 2001
From: Zengfeng Ye <zenye@microsoft.com>
Date: Wed, 15 Jan 2025 18:36:41 -0800
Subject: [PATCH] [iOS] Restore virtual url when restoring sessions

In some cases the navigation url and virtual url are different and the
url is a non HTTP(S) url (e.g. open downloaded PDF in tab, the url is
file:// but the virtual url is chrome://downloads), there is a problem
that the virtual url is dropped in the cases:
1. The url is forced to be set to virtual url
2. The virtual url is dropped because the url is not the same as the
real url read from the native session and the whole navigation item is
dropped.

This CL removes the restriction of not restoring the url and virtual url
for a non HTTP(S) url to fix the issue above. Based on current
implementation of session restoration:
1. Construct a new NavigationItem from the native
session(WKBackForwardListItem).
2. Restore newly constructed items from restored items if the URL is the
same. Otherwise use the newly constructed items as navigation items.

So the restored navigation item url has nothing to do with whether the
url is HTTP or not, the real navigation item will always be constructed
from native session.

Bug: 389729375
Change-Id: I40b28679bdd6077845cbc7b4121d6c66fe5d328a
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6171952
Reviewed-by: Sylvain Defresne <sdefresne@chromium.org>
Commit-Queue: Zengfeng Ye <zenye@microsoft.com>
Cr-Commit-Position: refs/heads/main@{#1407083}
---
 ios/web/navigation/navigation_item_impl.mm    | 19 +++----
 .../navigation_item_impl_unittest.mm          |  7 +--
 ios/web/navigation/navigation_manager_impl.mm | 11 ++--
 .../navigation_manager_impl_unittest.mm       | 55 ++++++++++++++++++-
 ios/web/public/test/web_state_test_util.mm    |  2 +-
 5 files changed, 71 insertions(+), 23 deletions(-)

diff --git a/ios/web/navigation/navigation_item_impl.mm b/ios/web/navigation/navigation_item_impl.mm
index 743c4bac80723..75c2d73bc5a0e 100644
--- a/ios/web/navigation/navigation_item_impl.mm
+++ b/ios/web/navigation/navigation_item_impl.mm
@@ -77,18 +77,17 @@ NavigationItemImpl::NavigationItemImpl(
   // also update the virtual URL reported by this object.
   url_ = original_request_url_ = GURL(storage.url());
 
-  // In the cases where the URL to be restored is not an HTTP URL, it is
-  // very likely that we can't restore the page (e.g. for files, it could
-  // be an external PDF that has been deleted), don't restore it to avoid
-  // issues.
+  // Restore the `virtual_url`. In case the `url` is invalid, it should be set
+  // to the the value saved for `virtual_url` (we never store `virtual_url` if
+  // equal to `url`, so restore to `url` only in that case).
   const GURL virtual_url(storage.virtual_url());
-  if (url_.SchemeIsHTTPOrHTTPS()) {
-    if (virtual_url.is_valid() && virtual_url != url_) {
-      virtual_url_ = virtual_url;
-    }
-  } else {
-    if (virtual_url.is_valid()) {
+  if (virtual_url.is_valid()) {
+    if (!url_.is_valid()) {
       url_ = virtual_url;
+    } else {
+      if (virtual_url != url_) {
+        virtual_url_ = virtual_url;
+      }
     }
   }
 }
diff --git a/ios/web/navigation/navigation_item_impl_unittest.mm b/ios/web/navigation/navigation_item_impl_unittest.mm
index 39c70ad869e3b..9bb323fe228f0 100644
--- a/ios/web/navigation/navigation_item_impl_unittest.mm
+++ b/ios/web/navigation/navigation_item_impl_unittest.mm
@@ -260,8 +260,7 @@ TEST_F(NavigationItemTest, NavigationItemImplRoundTripNonHTTPURL) {
 
   NavigationItemImpl decoded(storage);
 
-  EXPECT_NE(original.GetURL(), decoded.GetURL());
-  EXPECT_EQ(original.GetVirtualURL(), decoded.GetURL());
+  EXPECT_EQ(original.GetURL(), decoded.GetURL());
   EXPECT_EQ(original.GetVirtualURL(), decoded.GetVirtualURL());
 }
 
@@ -338,7 +337,7 @@ TEST_F(NavigationItemTest, DecodeFileScheme) {
   ASSERT_NE(storage.url(), storage.virtual_url());
 
   NavigationItemImpl navigation_item(storage);
-  EXPECT_EQ(GURL(storage.virtual_url()), navigation_item.GetURL());
+  EXPECT_EQ(GURL(storage.url()), navigation_item.GetURL());
   EXPECT_EQ(GURL(storage.virtual_url()), navigation_item.GetVirtualURL());
 }
 
@@ -350,7 +349,7 @@ TEST_F(NavigationItemTest, DecodeBlobScheme) {
   ASSERT_NE(storage.url(), storage.virtual_url());
 
   NavigationItemImpl navigation_item(storage);
-  EXPECT_EQ(GURL(storage.virtual_url()), navigation_item.GetURL());
+  EXPECT_EQ(GURL(storage.url()), navigation_item.GetURL());
   EXPECT_EQ(GURL(storage.virtual_url()), navigation_item.GetVirtualURL());
 }
 
diff --git a/ios/web/navigation/navigation_manager_impl.mm b/ios/web/navigation/navigation_manager_impl.mm
index fb0b52f407688..659c370b9763b 100644
--- a/ios/web/navigation/navigation_manager_impl.mm
+++ b/ios/web/navigation/navigation_manager_impl.mm
@@ -498,12 +498,13 @@ void NavigationManagerImpl::RestoreNativeSession() {
   // Native restore worked, abort unsafe restore.
   DiscardNonCommittedItems();
   last_committed_item_index_ = web_view_cache_.GetCurrentItemIndex();
-  if (restored_visible_item_ &&
-      restored_visible_item_->GetUserAgentType() != UserAgentType::NONE) {
-    NavigationItem* last_committed_item = GetLastCommittedItem();
+  if (restored_visible_item_) {
+    // Restore the state of the `restored_visible_item_` that were
+    // not overwritten by the session data restoration, including the
+    // user agent and virtual URL.
+    NavigationItemImpl* last_committed_item = GetLastCommittedItemImpl();
     if (last_committed_item) {
-      last_committed_item->SetUserAgentType(
-          restored_visible_item_->GetUserAgentType());
+      last_committed_item->RestoreStateFromItem(restored_visible_item_.get());
     }
   }
   restored_visible_item_.reset();
diff --git a/ios/web/navigation/navigation_manager_impl_unittest.mm b/ios/web/navigation/navigation_manager_impl_unittest.mm
index e9d2e79e502cb..791c4664c88ee 100644
--- a/ios/web/navigation/navigation_manager_impl_unittest.mm
+++ b/ios/web/navigation/navigation_manager_impl_unittest.mm
@@ -2852,7 +2852,7 @@ TEST_F(NavigationManagerSerialisationTest, RestoreFromProto) {
   proto::NavigationStorage storage;
   storage.set_last_committed_item_index(0);
   for (const char* url : kTestURLs) {
-    storage.add_items()->set_virtual_url(url);
+    storage.add_items()->set_url(url);
   }
   storage.set_last_committed_item_index(storage.items_size() - 1);
 
@@ -2912,7 +2912,7 @@ TEST_F(NavigationManagerSerialisationTest, RestoreFromProto_LastItemIndex) {
   proto::NavigationStorage storage;
   storage.set_last_committed_item_index(0);
   for (const char* url : kTestURLs) {
-    storage.add_items()->set_virtual_url(url);
+    storage.add_items()->set_url(url);
   }
   storage.set_last_committed_item_index(0);
 
@@ -2952,7 +2952,7 @@ TEST_F(NavigationManagerSerialisationTest, RestoreFromProto_IndexOutOfBound) {
     proto::WebStateStorage storage;
     proto::NavigationStorage* navigation_storage = storage.mutable_navigation();
     for (const char* url : kTestURLs) {
-      navigation_storage->add_items()->set_virtual_url(url);
+      navigation_storage->add_items()->set_url(url);
     }
     // Set an out-of-bound value for last committed item index.
     navigation_storage->set_last_committed_item_index(std::size(kTestURLs));
@@ -3005,4 +3005,53 @@ TEST_F(NavigationManagerSerialisationTest, RestoreFromProto_IndexOutOfBound) {
   }
 }
 
+// Tests that restoring a session with a specific virtual url which is different
+// with url works correctly.
+TEST_F(NavigationManagerSerialisationTest, RestoreVirtualURLFromProto) {
+  // Tests both HTTP url and non HTTP url.
+  const std::vector<std::pair<std::string, std::string>> urls_and_virtual_urls =
+      {
+          {"http://url.test", "http://virtual.test"},
+          {"file:///path/to/file.pdf", "http://virtual.test"},
+      };
+
+  proto::NavigationStorage storage;
+  storage.set_last_committed_item_index(0);
+
+  for (const auto& [url, virtual_url] : urls_and_virtual_urls) {
+    proto::NavigationItemStorage* item = storage.add_items();
+    item->set_url(url);
+    item->set_virtual_url(virtual_url);
+  }
+
+  // Create a WebState with a real navigation proxy as this is required to
+  // perform a session restore and access the view to force instantiation
+  // of the WKWebView.
+  std::unique_ptr<web::WebStateImpl> web_state =
+      std::make_unique<web::WebStateImpl>(
+          web::WebState::CreateParams(browser_state()));
+  std::ignore = web_state->GetView();
+
+  NavigationManagerImpl& navigation_manager =
+      web_state->GetNavigationManagerImpl();
+
+  navigation_manager.RestoreFromProto(storage);
+
+  const int urls_count = static_cast<int>(urls_and_virtual_urls.size());
+  ASSERT_TRUE(WaitUntilConditionOrTimeout(kWaitForActionTimeout, ^{
+    return navigation_manager.GetItemCount() == urls_count;
+  }));
+
+  EXPECT_EQ(0, navigation_manager.GetLastCommittedItemIndex());
+
+  for (int index = 0; index < urls_count; ++index) {
+    NavigationItem* item = navigation_manager.GetItemAtIndex(index);
+    GURL url = GURL(urls_and_virtual_urls[index].first);
+    GURL virtual_url = GURL(urls_and_virtual_urls[index].second);
+
+    EXPECT_EQ(item->GetURL(), url);
+    EXPECT_EQ(item->GetVirtualURL(), virtual_url);
+  }
+}
+
 }  // namespace web
diff --git a/ios/web/public/test/web_state_test_util.mm b/ios/web/public/test/web_state_test_util.mm
index da88b46306f3f..153da6a5720b5 100644
--- a/ios/web/public/test/web_state_test_util.mm
+++ b/ios/web/public/test/web_state_test_util.mm
@@ -171,7 +171,7 @@ std::unique_ptr<WebState> CreateUnrealizedWebStateWithItems(
     for (const PageInfo& info : items) {
       proto::NavigationItemStorage* item_storage =
           navigation_storage->add_items();
-      item_storage->set_virtual_url(info.url.spec());
+      item_storage->set_url(info.url.spec());
       item_storage->set_user_agent(storage.user_agent());
       item_storage->set_title(info.title);
     }