From c3d6769ba628f4da20d9907dc44dfb6821036d03 Mon Sep 17 00:00:00 2001
From: Anthony Tseng <darkdh@gmail.com>
Date: Mon, 5 Feb 2018 22:41:11 -0800
Subject: [PATCH] Fix map index key by SiteInstance and add callback

---
 atom/browser/api/atom_api_session.cc   |  5 +++--
 atom/browser/api/atom_api_session.h    |  3 ++-
 brave/browser/brave_browser_context.cc | 16 +++++++++++-----
 brave/browser/brave_browser_context.h  |  5 +++--
 4 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/atom/browser/api/atom_api_session.cc b/atom/browser/api/atom_api_session.cc
index 1ed0481f0f..a93c79a74a 100644
--- a/atom/browser/api/atom_api_session.cc
+++ b/atom/browser/api/atom_api_session.cc
@@ -602,12 +602,13 @@ bool Session::IsOffTheRecord() const {
   return false;
 }
 
-void Session::SetTorNewIdentity(const GURL& origin) const {
+void Session::SetTorNewIdentity(const GURL& url,
+                                const base::Closure& callback) const {
   brave::BraveBrowserContext* brave_browser_context =
    brave::BraveBrowserContext::FromBrowserContext(profile_);
   if (!brave_browser_context->IsIsolatedStorage())
     return;
-  brave_browser_context->SetTorNewIdentity(origin);
+  brave_browser_context->SetTorNewIdentity(url, callback);
 }
 
 // static
diff --git a/atom/browser/api/atom_api_session.h b/atom/browser/api/atom_api_session.h
index aa682f0389..467cb137ac 100644
--- a/atom/browser/api/atom_api_session.h
+++ b/atom/browser/api/atom_api_session.h
@@ -93,7 +93,8 @@ class Session: public mate::TrackableObject<Session>,
   v8::Local<v8::Value> Extensions(v8::Isolate* isolate);
   bool Equal(Session* session) const;
   bool IsOffTheRecord() const;
-  void SetTorNewIdentity(const GURL& origin) const;
+  void SetTorNewIdentity(const GURL& url,
+                         const base::Closure& callback) const;
 
  protected:
   Session(v8::Isolate* isolate, Profile* browser_context);
diff --git a/brave/browser/brave_browser_context.cc b/brave/browser/brave_browser_context.cc
index 75f5a32f7c..5da7c04971 100644
--- a/brave/browser/brave_browser_context.cc
+++ b/brave/browser/brave_browser_context.cc
@@ -46,6 +46,7 @@
 #include "content/public/browser/notification_source.h"
 #include "content/public/browser/browser_thread.h"
 #include "content/public/browser/dom_storage_context.h"
+#include "content/public/browser/site_instance.h"
 #include "content/public/browser/storage_partition.h"
 #include "crypto/random.h"
 #include "extensions/browser/pref_names.h"
@@ -345,7 +346,8 @@ BraveBrowserContext::CreateRequestContextForStoragePartition(
         std::move(request_interceptors));
     StoragePartitionDescriptor descriptor(partition_path, in_memory);
     url_request_context_getter_map_[descriptor] = url_request_context_getter;
-    TorSetProxy(url_request_context_getter, partition_path);
+    TorSetProxy(url_request_context_getter, partition_path,
+                base::Bind(&base::DoNothing));
     return url_request_context_getter;
   } else {
     return nullptr;
@@ -429,7 +431,7 @@ void BraveBrowserContext::UpdateDefaultZoomLevel() {
 
 void BraveBrowserContext::TorSetProxy(
     brightray::URLRequestContextGetter* url_request_context_getter,
-    const base::FilePath partition_path) {
+    const base::FilePath partition_path, const base::Closure& callback) {
   if (!url_request_context_getter || !isolated_storage_)
     return;
   if (!tor_proxy_.empty() && GURL(tor_proxy_).is_valid()) {
@@ -452,6 +454,8 @@ void BraveBrowserContext::TorSetProxy(
     proxy_service->ResetConfigService(base::WrapUnique(
       new net::ProxyConfigServiceFixed(config)));
     proxy_service->ForceReloadProxyConfig();
+    if (callback)
+      callback.Run();
   }
 }
 
@@ -740,8 +744,10 @@ void BraveBrowserContext::SetExitType(ExitType exit_type) {
   }
 }
 
-void BraveBrowserContext::SetTorNewIdentity(const GURL& origin) {
-  const std::string host = origin.host();
+void BraveBrowserContext::SetTorNewIdentity(const GURL& url,
+                                            const base::Closure& callback) {
+  GURL site_url(content::SiteInstance::GetSiteForURL(this, url));
+  const std::string host = site_url.host();
   base::FilePath partition_path = this->GetPath().Append(
     content::StoragePartitionImplMap::GetStoragePartitionPath(host, host));
   brightray::URLRequestContextGetter* url_request_context_getter;
@@ -752,7 +758,7 @@ void BraveBrowserContext::SetTorNewIdentity(const GURL& origin) {
     url_request_context_getter = (iter->second).get();
   else
     return;
-  TorSetProxy(url_request_context_getter, partition_path);
+  TorSetProxy(url_request_context_getter, partition_path, callback);
 }
 
 scoped_refptr<base::SequencedTaskRunner>
diff --git a/brave/browser/brave_browser_context.h b/brave/browser/brave_browser_context.h
index 663026406a..80391d63cf 100644
--- a/brave/browser/brave_browser_context.h
+++ b/brave/browser/brave_browser_context.h
@@ -141,7 +141,7 @@ class BraveBrowserContext : public Profile {
 
   bool IsIsolatedStorage() const { return isolated_storage_; }
 
-  void SetTorNewIdentity(const GURL& origin);
+  void SetTorNewIdentity(const GURL& url, const base::Closure& callback);
 
  private:
     typedef std::map<StoragePartitionDescriptor,
@@ -156,7 +156,8 @@ class BraveBrowserContext : public Profile {
 
   void TorSetProxy(
     brightray::URLRequestContextGetter* url_request_context_getter,
-    const base::FilePath partition_path);
+    const base::FilePath partition_path,
+    const base::Closure& callback);
 
   scoped_refptr<user_prefs::PrefRegistrySyncable> pref_registry_;
   std::unique_ptr<sync_preferences::PrefServiceSyncable> user_prefs_;