SessionRestore: Store session cookies and restore them if chrome crashes or auto-restarts.

chrome/browser/sessions/session_service.cc

Index: chrome/browser/sessions/session_service.cc
diff --git a/chrome/browser/sessions/session_service.cc b/chrome/browser/sessions/session_service.cc
index 1c570ac88b87ac2ecf5d90cc2b44e9e4dd8996a0..9aad390df8969892370ff7915afa8eb7adc12fed 100644
--- a/chrome/browser/sessions/session_service.cc
+++ b/chrome/browser/sessions/session_service.cc
@@ -9,13 +9,18 @@
 #include <set>
 #include <vector>
 
+#include "base/command_line.h"
 #include "base/file_util.h"
+#include "base/memory/ref_counted.h"
 #include "base/memory/scoped_vector.h"
 #include "base/message_loop.h"
 #include "base/metrics/histogram.h"
 #include "base/pickle.h"
 #include "base/threading/thread.h"
 #include "chrome/browser/extensions/extension_tab_helper.h"
+#include "chrome/browser/net/chrome_cookie_notification_details.h"
+#include "chrome/browser/net/sqlite_persistent_cookie_store.h"
+#include "chrome/browser/prefs/pref_service.h"
 #include "chrome/browser/prefs/session_startup_pref.h"
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/sessions/restore_tab_helper.h"
@@ -29,12 +34,17 @@
 #include "chrome/browser/ui/browser_window.h"
 #include "chrome/browser/ui/tab_contents/tab_contents_wrapper.h"
 #include "chrome/common/chrome_notification_types.h"
+#include "chrome/common/chrome_switches.h"
 #include "chrome/common/extensions/extension.h"
 #include "content/browser/tab_contents/navigation_details.h"
 #include "content/browser/tab_contents/navigation_entry.h"
 #include "content/browser/tab_contents/tab_contents.h"
-#include "content/public/browser/notification_service.h"
+#include "content/public/browser/browser_thread.h"
 #include "content/public/browser/notification_details.h"
+#include "content/public/browser/notification_service.h"
+#include "net/base/cookie_monster.h"
+#include "net/url_request/url_request_context.h"
+#include "net/url_request/url_request_context_getter.h"
 
 #if defined(OS_MACOSX)
 #include "chrome/browser/app_controller_cppsafe_mac.h"
@@ -68,6 +78,9 @@ static const int kWritesPerReset = 250;
 
 namespace {
 
+const FilePath::CharType kSessionCookiesDatabaseName[] =
+    FILE_PATH_LITERAL("SessionCookies");
+
 // The callback from GetLastSession is internally routed to SessionService
 // first and then the caller. This is done so that the SessionWindows can be
 // recreated from the SessionCommands and the SessionWindows passed to the
@@ -139,6 +152,94 @@ struct PinnedStatePayload {
 
 }  // namespace
 
+// |SessionService::SessionCookieStore| stores session cookies in a
+// |CookieMonster| and restores them. The constructor should be called in the
+// UI thread, and all other functions should be called in the IO thread.
+class SessionService::SessionCookieStore
+    : public base::RefCountedThreadSafe<
+          SessionCookieStore,
+          content::BrowserThread::DeleteOnIOThread> {
+ public:
+  typedef base::Callback<void()> CookieOperationCallback;
+
+  explicit SessionCookieStore(const FilePath& path)
+      : path_(path) {
+    content::BrowserThread::PostTask(
+        content::BrowserThread::IO, FROM_HERE,
+        base::Bind(&SessionCookieStore::InitializeOnIOThread, this));
+  }
+
+  ~SessionCookieStore() {
+    DCHECK(content::BrowserThread::CurrentlyOn(content::BrowserThread::IO));
+  }
+
+  void RestoreSessionCookies(
+      net::URLRequestContextGetter* url_request_context_getter) {
+    net::CookieMonster* real_cookie_monster =
+        url_request_context_getter->GetURLRequestContext()->cookie_store()->
+        GetCookieMonster();
+    real_cookie_monster->MergeSessionCookies(session_cookie_monster_);
+  }
+
+  void ForgetSessionCookies() {
+    session_cookie_monster_->DeleteAllAsync(
+        net::CookieMonster::DeleteCallback());
+  }
+
+  class CanonicalCookieWrapper :
+      public base::RefCountedThreadSafe<CanonicalCookieWrapper> {
+   public:
+    explicit CanonicalCookieWrapper(
+        const net::CookieMonster::CanonicalCookie* cookie)
+        : cookie_(new net::CookieMonster::CanonicalCookie(*cookie)) {}
+    ~CanonicalCookieWrapper() {}
+
+    scoped_ptr<net::CookieMonster::CanonicalCookie> cookie_;
+  };
+
+  void AddSessionCookie(CanonicalCookieWrapper* wrapper) {
+    // Set the cookie as permanent so that |session_cookie_monster_| syncs it
+    // to its backend store.
+    wrapper->cookie_->SetPersistent(true);
+    net::CookieList cookie_list;
+    cookie_list.push_back(*(wrapper->cookie_));
+    session_cookie_monster_->InitializeFrom(cookie_list);
+  }
+
+  void RemoveSessionCookie(CanonicalCookieWrapper* wrapper,
+                           const CookieOperationCallback& callback) {
+    session_cookie_monster_->DeleteCanonicalCookieAsync(
+        *(wrapper->cookie_),
+        base::Bind(&SessionCookieStore::OnCookieDeleted,
+                   this,
+                   callback));
+  }
+
+ private:
+  void InitializeOnIOThread() {
+    session_cookie_store_ = new SQLitePersistentCookieStore(path_);
+    session_cookie_monster_ =
+        new net::CookieMonster(session_cookie_store_.get(), NULL);
+  }
+
+  void OnCookieDeleted(const CookieOperationCallback& callback,
+                       bool success) {
+    if (!callback.is_null()) {
+      content::BrowserThread::PostTask(
+          content::BrowserThread::UI, FROM_HERE,
+          base::Bind(&SessionCookieStore::InvokeCallback, this, callback));
+    }
+  }
+
+  void InvokeCallback(const CookieOperationCallback& callback) {
+    callback.Run();
+  }
+
+  FilePath path_;
+  scoped_refptr<SQLitePersistentCookieStore> session_cookie_store_;
+  scoped_refptr<net::CookieMonster> session_cookie_monster_;
+};
+
 // SessionService -------------------------------------------------------------
 
 SessionService::SessionService(Profile* profile)
@@ -147,7 +248,10 @@ SessionService::SessionService(Profile* profile)
       move_on_new_browser_(false),
       save_delay_in_millis_(base::TimeDelta::FromMilliseconds(2500)),
       save_delay_in_mins_(base::TimeDelta::FromMinutes(10)),
-      save_delay_in_hrs_(base::TimeDelta::FromHours(8)) {
+      save_delay_in_hrs_(base::TimeDelta::FromHours(8)),
+      enable_restore_session_data_(
+          CommandLine::ForCurrentProcess()->HasSwitch(
+              switches::kEnableRestoreSessionState)) {
   Init();
 }
 
@@ -157,7 +261,10 @@ SessionService::SessionService(const FilePath& save_path)
       move_on_new_browser_(false),
       save_delay_in_millis_(base::TimeDelta::FromMilliseconds(2500)),
       save_delay_in_mins_(base::TimeDelta::FromMinutes(10)),
-      save_delay_in_hrs_(base::TimeDelta::FromHours(8)) {
+      save_delay_in_hrs_(base::TimeDelta::FromHours(8)),
+      enable_restore_session_data_(
+          CommandLine::ForCurrentProcess()->HasSwitch(
+              switches::kEnableRestoreSessionState)) {
   Init();
 }
 
@@ -180,6 +287,12 @@ void SessionService::MoveCurrentSessionToLastSession() {
 
   Save();
 
+  content::BrowserThread::PostTask(
+      content::BrowserThread::IO, FROM_HERE,
+      base::Bind(
+          &SessionCookieStore::ForgetSessionCookies,
+          session_cookie_store_));
+
   if (!backend_thread()) {
     backend()->MoveCurrentSessionToLastSession();
   } else {
@@ -465,6 +578,33 @@ void SessionService::Save() {
   }
 }
 
+void SessionService::RestoreOrResetSessionState(bool restore) {
+  // Schedule the restoring or deletion of the session cookies in the IO
+  // thread. |RestoreSessionCookies| (or |ForgetSessionCookies|) will be
+  // executed be done before any cookies are served, since the cookies will be
+  // served on the IO thread.
+  if (!session_cookie_store_.get())
+    return;
+
+  bool restore_session_cookies = enable_restore_session_data_ &&
+                                 restore &&
+                                 profile();
+  if (restore_session_cookies) {
+    content::BrowserThread::PostTask(
+        content::BrowserThread::IO, FROM_HERE,
+        base::Bind(
+            &SessionCookieStore::RestoreSessionCookies,
+            session_cookie_store_,
+            make_scoped_refptr(profile()->GetRequestContext())));
+  } else {
+    content::BrowserThread::PostTask(
+        content::BrowserThread::IO, FROM_HERE,
+        base::Bind(
+            &SessionCookieStore::ForgetSessionCookies,
+            session_cookie_store_));
+  }
+}
+
 void SessionService::Init() {
   // Register for the notifications we're interested in.
   registrar_.Add(this, content::NOTIFICATION_TAB_PARENTED,
@@ -482,6 +622,13 @@ void SessionService::Init() {
   registrar_.Add(
       this, chrome::NOTIFICATION_TAB_CONTENTS_APPLICATION_EXTENSION_CHANGED,
       content::NotificationService::AllSources());
+  if (enable_restore_session_data_ && profile()) {
+    registrar_.Add(
+        this, chrome::NOTIFICATION_COOKIE_CHANGED,
+        content::Source<Profile>(profile()));
+    session_cookie_store_ = new SessionCookieStore(
+        profile()->GetPath().Append(kSessionCookiesDatabaseName));
+  }
 }
 
 bool SessionService::ShouldNewWindowStartSession() {
@@ -650,6 +797,35 @@ void SessionService::Observe(int type,
       break;
     }
 
+    case chrome::NOTIFICATION_COOKIE_CHANGED: {
+      if (enable_restore_session_data_ && session_cookie_store_.get()) {
+        ChromeCookieDetails* cookie_details =
+            content::Details<ChromeCookieDetails>(details).ptr();
+        if (!cookie_details->cookie->IsPersistent()) {
+          scoped_refptr<SessionCookieStore::CanonicalCookieWrapper> cookie =
+              new SessionCookieStore::CanonicalCookieWrapper(
+                  cookie_details->cookie);
+          if (cookie_details->removed) {
+            content::BrowserThread::PostTask(
+                content::BrowserThread::IO, FROM_HERE,
+                base::Bind(
+                    &SessionCookieStore::RemoveSessionCookie,
+                    session_cookie_store_,
+                    cookie,
+                    SessionCookieStore::CookieOperationCallback()));
+          } else {
+            content::BrowserThread::PostTask(
+                content::BrowserThread::IO, FROM_HERE,
+                base::Bind(
+                    &SessionCookieStore::AddSessionCookie,
+                    session_cookie_store_,
+                    cookie));
+          }
+        }
+      }
+      break;
+    }
+
     default:
       NOTREACHED();
   }