7 files changed, 117 insertions, 16 deletions
diff --git a/chrome/browser/browser_url_handler.cc b/chrome/browser/browser_url_handler.cc
index 0cd1e2f..b209e04 100644
--- a/chrome/browser/browser_url_handler.cc
+++ b/chrome/browser/browser_url_handler.cc
@@ -43,6 +43,21 @@ static bool HandleViewSource(GURL* url, Profile* profile) {
   return false;
 }
 
+// Turns a non view-source URL into the corresponding view-source URL.
+static bool ReverseViewSource(GURL* url, Profile* profile) {
+  // No action necessary if the URL is already view-source:
+  if (url->SchemeIs(chrome::kViewSourceScheme))
+    return false;
+
+  url_canon::Replacements<char> repl;
+  repl.SetScheme(chrome::kViewSourceScheme,
+      url_parse::Component(0, strlen(chrome::kViewSourceScheme)));
+  repl.SetPath(url->spec().c_str(),
+      url_parse::Component(0, url->spec().size()));
+  *url = url->ReplaceComponents(repl);
+  return true;
+}
+
 // Handles rewriting DOM UI URLs.
 static bool HandleDOMUI(GURL* url, Profile* profile) {
   if (!DOMUIFactory::UseDOMUIForURL(*url))
@@ -60,7 +75,7 @@ static bool HandleDOMUI(GURL* url, Profile* profile) {
   return true;
 }
 
-std::vector<BrowserURLHandler::URLHandler> BrowserURLHandler::url_handlers_;
+std::vector<BrowserURLHandler::HandlerPair> BrowserURLHandler::url_handlers_;
 
 // static
 void BrowserURLHandler::InitURLHandlers() {
@@ -68,18 +83,40 @@ void BrowserURLHandler::InitURLHandlers() {
     return;
 
   // Add the default URL handlers.
-  url_handlers_.push_back(&ExtensionDOMUI::HandleChromeURLOverride);
-  url_handlers_.push_back(&WillHandleBrowserAboutURL);  // about:
-  url_handlers_.push_back(&HandleDOMUI);                // chrome: & friends.
-  url_handlers_.push_back(&HandleViewSource);           // view-source:
+  url_handlers_.push_back(
+      HandlerPair(&ExtensionDOMUI::HandleChromeURLOverride, NULL));
+  // about:
+  url_handlers_.push_back(HandlerPair(&WillHandleBrowserAboutURL, NULL));
+  // chrome: & friends.
+  url_handlers_.push_back(HandlerPair(&HandleDOMUI, NULL));
+  // view-source:
+  url_handlers_.push_back(HandlerPair(&HandleViewSource, &ReverseViewSource));
 }
 
 // static
-void BrowserURLHandler::RewriteURLIfNecessary(GURL* url, Profile* profile) {
+void BrowserURLHandler::RewriteURLIfNecessary(GURL* url, Profile* profile,
+                                              bool* reverse_on_redirect) {
   if (url_handlers_.empty())
     InitURLHandlers();
   for (size_t i = 0; i < url_handlers_.size(); ++i) {
-    if ((*url_handlers_[i])(url, profile))
+    if ((*url_handlers_[i].first)(url, profile)) {
+      *reverse_on_redirect = (url_handlers_[i].second != NULL);
       return;
+    }
   }
 }
+
+// static
+bool BrowserURLHandler::ReverseURLRewrite(
+    GURL* url, const GURL& original, Profile* profile) {
+  for (size_t i = 0; i < url_handlers_.size(); ++i) {
+    GURL test_url(original);
+    if ((*url_handlers_[i].first)(&test_url, profile)) {
+      if (url_handlers_[i].second)
+        return (*url_handlers_[i].second)(url, profile);
+      else
+        return false;
+    }
+  }
+  return false;
+}
diff --git a/chrome/browser/browser_url_handler.h b/chrome/browser/browser_url_handler.h
index 593408e..e8b8a97 100644
--- a/chrome/browser/browser_url_handler.h
+++ b/chrome/browser/browser_url_handler.h
@@ -13,6 +13,7 @@
 #define CHROME_BROWSER_BROWSER_URL_HANDLER_H_
 
 #include <vector>
+#include <utility>
 
 class GURL;
 class Profile;
@@ -24,21 +25,27 @@ class BrowserURLHandler {
   // The type of functions that can process a URL.
   // If a handler handles |url|, it should :
   // - optionally modify |url| to the URL that should be sent to the renderer
-  // - optionally set |dispatcher| to the necessary DOMMessageDispatcher
-  // - return true.
   // If the URL is not handled by a handler, it should return false.
   typedef bool (*URLHandler)(GURL* url, Profile* profile);
 
-  // HandleBrowserURL gives all registered URLHandlers a shot at processing
+  // RewriteURLIfNecessary gives all registered URLHandlers a shot at processing
   // the given URL, and modifies it in place.
-  static void RewriteURLIfNecessary(GURL* url, Profile* profile);
+  // If the original URL needs to be adjusted if the modified URL is redirected,
+  // this function sets |reverse_on_redirect| to true.
+  static void RewriteURLIfNecessary(GURL* url, Profile* profile,
+                                    bool* reverse_on_redirect);
 
-  // We initialize the list of url_handlers_ lazily the first time MaybeHandle
-  // is called.
+  // Reverses the rewriting that was done for |original| using the new |url|.
+  static bool ReverseURLRewrite(GURL* url, const GURL& original,
+                                Profile* profile);
+
+  // We initialize the list of url_handlers_ lazily the first time
+  // RewriteURLIfNecessary is called.
   static void InitURLHandlers();
 
-  // The list of known URLHandlers.
-  static std::vector<URLHandler> url_handlers_;
+  // The list of known URLHandlers, optionally with reverse-rewriters.
+  typedef std::pair<URLHandler, URLHandler> HandlerPair;
+  static std::vector<HandlerPair> url_handlers_;
 };
 
 #endif  // CHROME_BROWSER_BROWSER_URL_HANDLER_H_
diff --git a/chrome/browser/tab_contents/navigation_controller.cc b/chrome/browser/tab_contents/navigation_controller.cc
index 9f1fedf..847f290 100644
--- a/chrome/browser/tab_contents/navigation_controller.cc
+++ b/chrome/browser/tab_contents/navigation_controller.cc
@@ -361,12 +361,15 @@ NavigationEntry* NavigationController::CreateNavigationEntry(
   // will actually be loaded. This real URL won't be shown to the user, just
   // used internally.
   GURL loaded_url(url);
-  BrowserURLHandler::RewriteURLIfNecessary(&loaded_url, profile_);
+  bool reverse_on_redirect = false;
+  BrowserURLHandler::RewriteURLIfNecessary(
+      &loaded_url, profile_, &reverse_on_redirect);
 
   NavigationEntry* entry = new NavigationEntry(NULL, -1, loaded_url, referrer,
                                                string16(), transition);
   entry->set_virtual_url(url);
   entry->set_user_typed_url(url);
+  entry->set_update_virtual_url_with_url(reverse_on_redirect);
   if (url.SchemeIsFile()) {
     std::wstring languages = profile()->GetPrefs()->GetString(
         prefs::kAcceptLanguages);
@@ -376,6 +379,15 @@ NavigationEntry* NavigationController::CreateNavigationEntry(
   return entry;
 }
 
+void NavigationController::UpdateVirtualURLToURL(
+    NavigationEntry* entry, const GURL& new_url) {
+  GURL new_virtual_url(new_url);
+  if (BrowserURLHandler::ReverseURLRewrite(
+          &new_virtual_url, entry->virtual_url(), profile_)) {
+    entry->set_virtual_url(new_virtual_url);
+  }
+}
+
 void NavigationController::AddTransientEntry(NavigationEntry* entry) {
   // Discard any current transient entry, we can only have one at a time.
   int index = 0;
@@ -619,6 +631,8 @@ void NavigationController::RendererDidNavigateToNewPage(
   }
 
   new_entry->set_url(params.url);
+  if (new_entry->update_virtual_url_with_url())
+    UpdateVirtualURLToURL(new_entry, params.url);
   new_entry->set_referrer(params.referrer);
   new_entry->set_page_id(params.page_id);
   new_entry->set_transition_type(params.transition);
@@ -654,6 +668,8 @@ void NavigationController::RendererDidNavigateToExistingPage(
   // be the same except during session restore, when no site instance will be
   // assigned.
   entry->set_url(params.url);
+  if (entry->update_virtual_url_with_url())
+    UpdateVirtualURLToURL(entry, params.url);
   DCHECK(entry->site_instance() == NULL ||
          entry->site_instance() == tab_contents_->GetSiteInstance());
   entry->set_site_instance(tab_contents_->GetSiteInstance());
@@ -688,6 +704,8 @@ void NavigationController::RendererDidNavigateToSamePage(
   existing_entry->set_unique_id(pending_entry_->unique_id());
 
   // The URL may have changed due to redirects.
+  if (existing_entry->update_virtual_url_with_url())
+    UpdateVirtualURLToURL(existing_entry, params.url);
   existing_entry->set_url(params.url);
 
   DiscardNonCommittedEntries();
@@ -707,6 +725,8 @@ void NavigationController::RendererDidNavigateInPage(
   // reference fragments, of course).
   NavigationEntry* new_entry = new NavigationEntry(*existing_entry);
   new_entry->set_page_id(params.page_id);
+  if (new_entry->update_virtual_url_with_url())
+    UpdateVirtualURLToURL(new_entry, params.url);
   new_entry->set_url(params.url);
   *did_replace_entry = IsRedirect(params) &&
                        IsLikelyAutoNavigation(base::TimeTicks::Now());
diff --git a/chrome/browser/tab_contents/navigation_controller.h b/chrome/browser/tab_contents/navigation_controller.h
index a227765..59af542 100644
--- a/chrome/browser/tab_contents/navigation_controller.h
+++ b/chrome/browser/tab_contents/navigation_controller.h
@@ -441,6 +441,10 @@ class NavigationController {
   NavigationEntry* CreateNavigationEntry(const GURL& url, const GURL& referrer,
                                          PageTransition::Type transition);
 
+  // Updates the virtual URL of an entry to match a new URL, for cases where
+  // the real renderer URL is derived from the virtual URL, like view-source:
+  void UpdateVirtualURLToURL(NavigationEntry* entry, const GURL& new_url);
+
   // Invoked after session/tab restore or cloning a tab. Resets the transition
   // type of the entries, updates the max page id and creates the active
   // contents. See RestoreFromState for a description of from_last_session.
diff --git a/chrome/browser/tab_contents/navigation_controller_unittest.cc b/chrome/browser/tab_contents/navigation_controller_unittest.cc
index 142ad8b..36344b0 100644
--- a/chrome/browser/tab_contents/navigation_controller_unittest.cc
+++ b/chrome/browser/tab_contents/navigation_controller_unittest.cc
@@ -1439,6 +1439,30 @@ TEST_F(NavigationControllerTest, SameSubframe) {
   EXPECT_EQ(controller().last_committed_entry_index(), 0);
 }
 
+// Test view source redirection is reflected in title bar.
+TEST_F(NavigationControllerTest, ViewSourceRedirect) {
+  const char kUrl[] = "view-source:http://redirect.to/google.com";
+  const char kResult[] = "http://google.com/";
+  const char kExpected[] = "view-source:http://google.com/";
+  const GURL url(kUrl);
+  const GURL result_url(kResult);
+
+  controller().LoadURL(url, GURL(), PageTransition::TYPED);
+
+  ViewHostMsg_FrameNavigate_Params params = {0};
+  params.page_id = 0;
+  params.url = result_url;
+  params.transition = PageTransition::SERVER_REDIRECT;
+  params.should_update_history = false;
+  params.gesture = NavigationGestureAuto;
+  params.is_post = false;
+  NavigationController::LoadCommittedDetails details;
+  controller().RendererDidNavigate(params, 0, &details);
+
+  EXPECT_EQ(ASCIIToUTF16(kExpected), contents()->GetTitle());
+  EXPECT_EQ(true, contents()->ShouldDisplayURL());
+}
+
 /* TODO(brettw) These test pass on my local machine but fail on the XP buildbot
    (but not Vista) cleaning up the directory after they run.
    This should be fixed.
diff --git a/chrome/browser/tab_contents/navigation_entry.cc b/chrome/browser/tab_contents/navigation_entry.cc
index 826e648..c8db037 100644
--- a/chrome/browser/tab_contents/navigation_entry.cc
+++ b/chrome/browser/tab_contents/navigation_entry.cc
@@ -57,6 +57,7 @@ NavigationEntry::NavigationEntry(SiteInstance* instance,
       page_type_(NORMAL_PAGE),
       url_(url),
       referrer_(referrer),
+      update_virtual_url_with_url_(false),
       title_(title),
       page_id_(page_id),
       transition_type_(transition_type),
diff --git a/chrome/browser/tab_contents/navigation_entry.h b/chrome/browser/tab_contents/navigation_entry.h
index a6ee1dc..df6ad06 100644
--- a/chrome/browser/tab_contents/navigation_entry.h
+++ b/chrome/browser/tab_contents/navigation_entry.h
@@ -253,6 +253,13 @@ class NavigationEntry {
     return virtual_url_.is_empty() ? url_ : virtual_url_;
   }
 
+  bool update_virtual_url_with_url() const {
+    return update_virtual_url_with_url_;
+  }
+  void set_update_virtual_url_with_url(bool update) {
+    update_virtual_url_with_url_ = update;
+  }
+
   // The title as set by the page. This will be empty if there is no title set.
   // The caller is responsible for detecting when there is no title and
   // displaying the appropriate "Untitled" label if this is being displayed to
@@ -405,6 +412,7 @@ class NavigationEntry {
   GURL url_;
   GURL referrer_;
   GURL virtual_url_;
+  bool update_virtual_url_with_url_;
   string16 title_;
   FaviconStatus favicon_;
   std::string content_state_;