Refactor how cookie blocking works

Tests started to fail for cookie blockin C71, this makes it work again

browser/brave_content_browser_client.cc

@@ -11,10 +11,13 @@
 #include "brave/browser/extensions/brave_tor_client_updater.h"
 #include "brave/browser/renderer_host/brave_navigation_ui_data.h"
 #include "brave/browser/tor/tor_profile_service_factory.h"
+#include "brave/common/brave_cookie_blocking.h"
 #include "brave/common/webui_url_constants.h"
 #include "brave/common/tor/tor_launcher.mojom.h"
 #include "brave/common/tor/switches.h"
+#include "brave/components/brave_shields/browser/brave_shields_util.h"
 #include "brave/components/brave_shields/browser/brave_shields_web_contents_observer.h"
+#include "brave/components/brave_shields/common/brave_shield_constants.h"
 #include "brave/components/brave_webtorrent/browser/content_browser_client_helper.h"
 #include "brave/components/content_settings/core/browser/brave_cookie_settings.h"
 #include "brave/grit/brave_generated_resources.h"
@@ -23,6 +26,7 @@
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/profiles/profile_io_data.h"
 #include "chrome/common/url_constants.h"
+#include "components/content_settings/core/browser/cookie_settings.h"
 #include "content/browser/frame_host/render_frame_host_impl.h"
 #include "content/public/browser/browser_task_traits.h"
 #include "content/public/browser/browser_thread.h"
@@ -33,7 +37,6 @@
 using content::BrowserThread;
 using content::RenderFrameHost;
 using content::WebContents;
-using content_settings::BraveCookieSettings;
 using brave_shields::BraveShieldsWebContentsObserver;
 
 #if BUILDFLAG(ENABLE_EXTENSIONS)
@@ -73,12 +76,6 @@ bool HandleURLReverseRewrite(GURL* url,
   return false;
 }
 
-WebContents* GetWebContents(int render_process_id, int render_frame_id) {
-  RenderFrameHost* rfh =
-      RenderFrameHost::FromID(render_process_id, render_frame_id);
-  return WebContents::FromRenderFrameHost(rfh);
-}
-
 }
 
 BraveContentBrowserClient::BraveContentBrowserClient(std::unique_ptr<ui::DataPack> data_pack,
@@ -109,30 +106,51 @@ void BraveContentBrowserClient::BrowserURLHandlerCreated(
   ChromeContentBrowserClient::BrowserURLHandlerCreated(handler);
 }
 
-bool BraveContentBrowserClient::AllowGetCookie(
-    const GURL& url,
-    const GURL& first_party,
-    const net::CookieList& cookie_list,
-    content::ResourceContext* context,
-    int render_process_id,
-    int render_frame_id) {
-  DCHECK_CURRENTLY_ON(BrowserThread::IO);
+bool BraveContentBrowserClient::AllowAccessCookie(const GURL& url, const GURL& first_party,
+    content::ResourceContext* context, int render_process_id, int render_frame_id) {
+  GURL tab_origin =
+      BraveShieldsWebContentsObserver::GetTabURLFromRenderFrameInfo(
+          render_process_id, render_frame_id).GetOrigin();
   ProfileIOData* io_data = ProfileIOData::FromResourceContext(context);
+  bool allow_brave_shields = brave_shields::IsAllowContentSettingWithIOData(
+      io_data, tab_origin, tab_origin, CONTENT_SETTINGS_TYPE_PLUGINS,
+      brave_shields::kBraveShields);
+  bool allow_1p_cookies = brave_shields::IsAllowContentSettingWithIOData(
+      io_data, tab_origin, GURL("https://firstParty/"),
+      CONTENT_SETTINGS_TYPE_PLUGINS, brave_shields::kCookies);
+  bool allow_3p_cookies = brave_shields::IsAllowContentSettingWithIOData(
+      io_data, tab_origin, GURL(), CONTENT_SETTINGS_TYPE_PLUGINS,
+      brave_shields::kCookies);
+  content_settings::BraveCookieSettings* cookie_settings =
+      (content_settings::BraveCookieSettings*)io_data->GetCookieSettings();
+  bool allow = !ShouldBlockCookie(allow_brave_shields, allow_1p_cookies,
+                   allow_3p_cookies, first_party, url) &&
+      cookie_settings->IsCookieAccessAllowed(url, first_party, tab_origin);
+  return allow;
+}
 
-  GURL tab_url = BraveShieldsWebContentsObserver::GetTabURLFromRenderFrameInfo(
-      render_process_id, render_frame_id);
-  BraveCookieSettings* cookie_settings =
-      (BraveCookieSettings*)io_data->GetCookieSettings();
+bool BraveContentBrowserClient::AllowGetCookie(const GURL& url,
+    const GURL& first_party, const net::CookieList& cookie_list,
+    content::ResourceContext* context, int render_process_id,
+    int render_frame_id) {
+  DCHECK_CURRENTLY_ON(BrowserThread::IO);
+  bool allow = AllowAccessCookie(url, first_party, context, render_process_id,
+                                 render_frame_id);
+  OnCookiesRead(render_process_id, render_frame_id, url, first_party,
+                cookie_list, !allow);
 
-  bool allow =
-      cookie_settings->IsCookieAccessAllowed(url, first_party, tab_url);
+  return allow;
+}
 
-  base::Callback<WebContents*(void)> wc_getter =
-      base::Bind(&GetWebContents, render_process_id, render_frame_id);
-  base::PostTaskWithTraits(
-      FROM_HERE, {BrowserThread::UI},
-      base::BindOnce(&TabSpecificContentSettings::CookiesRead, wc_getter, url,
-                     first_party, cookie_list, !allow));
+bool BraveContentBrowserClient::AllowSetCookie(const GURL& url,
+    const GURL& first_party, const net::CanonicalCookie& cookie,
+    content::ResourceContext* context, int render_process_id,
+    int render_frame_id) {
+  DCHECK_CURRENTLY_ON(BrowserThread::IO);
+  bool allow = AllowAccessCookie(url, first_party, context, render_process_id,
+                                 render_frame_id);
+  OnCookieChange(render_process_id, render_frame_id, url, first_party, cookie,
+                 !allow);
   return allow;
 }
 
@@ -163,34 +181,6 @@ bool BraveContentBrowserClient::HandleExternalProtocol(
       page_transition, has_user_gesture);
 }
 
-bool BraveContentBrowserClient::AllowSetCookie(
-    const GURL& url,
-    const GURL& first_party,
-    const net::CanonicalCookie& cookie,
-    content::ResourceContext* context,
-    int render_process_id,
-    int render_frame_id) {
-  DCHECK_CURRENTLY_ON(BrowserThread::IO);
-  ProfileIOData* io_data = ProfileIOData::FromResourceContext(context);
-
-  BraveCookieSettings* cookie_settings =
-      (BraveCookieSettings*)io_data->GetCookieSettings();
-
-  GURL tab_url = BraveShieldsWebContentsObserver::GetTabURLFromRenderFrameInfo(
-      render_process_id, render_frame_id);
-
-  bool allow =
-    cookie_settings->IsCookieAccessAllowed(url, first_party, tab_url);
-
-  base::Callback<WebContents*(void)> wc_getter =
-      base::Bind(&GetWebContents, render_process_id, render_frame_id);
-  base::PostTaskWithTraits(
-      FROM_HERE, {BrowserThread::UI},
-      base::BindOnce(&TabSpecificContentSettings::CookieChanged, wc_getter, url,
-                     first_party, cookie, !allow));
-  return allow;
-}
-
 void BraveContentBrowserClient::RegisterOutOfProcessServices(
       OutOfProcessServiceMap* services) {
   ChromeContentBrowserClient::RegisterOutOfProcessServices(services);

browser/brave_content_browser_client.h

@@ -60,6 +60,9 @@ class BraveContentBrowserClient : public ChromeContentBrowserClient {
                        const GURL& url) override;
 
  private:
+  bool AllowAccessCookie(const GURL& url, const GURL& first_party,
+      content::ResourceContext* context, int render_process_id,
+      int render_frame_id);
   DISALLOW_COPY_AND_ASSIGN(BraveContentBrowserClient);
 };
 

browser/net/BUILD.gn

@@ -7,6 +7,8 @@ source_set("net") {
     "brave_ad_block_tp_network_delegate_helper.h",
     "brave_common_static_redirect_network_delegate_helper.cc",
     "brave_common_static_redirect_network_delegate_helper.h",
+    "cookie_network_delegate_helper.cc",
+    "cookie_network_delegate_helper.h",
     "brave_httpse_network_delegate_helper.cc",
     "brave_httpse_network_delegate_helper.h",
     "brave_network_delegate_base.cc",

browser/net/brave_network_delegate_base.cc

@@ -8,14 +8,34 @@
 
 #include "base/task/post_task.h"
 #include "brave/common/pref_names.h"
+#include "brave/components/brave_shields/browser/brave_shields_util.h"
 #include "brave/components/brave_shields/common/brave_shield_constants.h"
 #include "chrome/browser/browser_process.h"
+#include "chrome/browser/content_settings/tab_specific_content_settings.h"
 #include "components/prefs/pref_service.h"
 #include "content/public/browser/browser_task_traits.h"
 #include "content/public/browser/browser_thread.h"
+#include "content/public/browser/render_frame_host.h"
+#include "content/public/browser/web_contents.h"
 #include "net/url_request/url_request.h"
 
 using content::BrowserThread;
+using net::URLRequest;
+
+
+namespace {
+
+  content::WebContents* GetWebContentsFromProcessAndFrameId(
+      int render_process_id, int render_frame_id) {
+  if (render_process_id) {
+    content::RenderFrameHost* rfh =
+        content::RenderFrameHost::FromID(render_process_id, render_frame_id);
+    return content::WebContents::FromRenderFrameHost(rfh);
+  }
+  return content::WebContents::FromFrameTreeNodeId(render_frame_id);
+}
+
+}  // namespace
 
 BraveNetworkDelegateBase::BraveNetworkDelegateBase(
     extensions::EventRouterForwarder* event_router)
@@ -37,7 +57,7 @@ void BraveNetworkDelegateBase::GetReferralHeaders() {
     referral_headers_list_ = referral_headers->CreateDeepCopy();
 }
 
-int BraveNetworkDelegateBase::OnBeforeURLRequest(net::URLRequest* request,
+int BraveNetworkDelegateBase::OnBeforeURLRequest(URLRequest* request,
     net::CompletionOnceCallback callback,
     GURL* new_url) {
   if (before_url_request_callbacks_.empty() || !request) {
@@ -53,7 +73,7 @@ int BraveNetworkDelegateBase::OnBeforeURLRequest(net::URLRequest* request,
   return net::ERR_IO_PENDING;
 }
 
-int BraveNetworkDelegateBase::OnBeforeStartTransaction(net::URLRequest* request,
+int BraveNetworkDelegateBase::OnBeforeStartTransaction(URLRequest* request,
     net::CompletionOnceCallback callback,
     net::HttpRequestHeaders* headers) {
   if (before_start_transaction_callbacks_.empty() || !request) {
@@ -71,7 +91,7 @@ int BraveNetworkDelegateBase::OnBeforeStartTransaction(net::URLRequest* request,
   return net::ERR_IO_PENDING;
 }
 
-int BraveNetworkDelegateBase::OnHeadersReceived(net::URLRequest* request,
+int BraveNetworkDelegateBase::OnHeadersReceived(URLRequest* request,
       net::CompletionOnceCallback callback,
       const net::HttpResponseHeaders* original_response_headers,
       scoped_refptr<net::HttpResponseHeaders>* override_response_headers,
@@ -102,14 +122,73 @@ int BraveNetworkDelegateBase::OnHeadersReceived(net::URLRequest* request,
   return net::ERR_IO_PENDING;
 }
 
+bool BraveNetworkDelegateBase::OnCanGetCookies(const URLRequest& request,
+    const net::CookieList& cookie_list,
+    bool allowed_from_caller) {
+  std::shared_ptr<brave::BraveRequestInfo> ctx(
+      new brave::BraveRequestInfo());
+  brave::BraveRequestInfo::FillCTXFromRequest(&request, ctx);
+  ctx->event_type = brave::kOnCanGetCookies;
+  bool allow = std::all_of(can_get_cookies_callbacks_.begin(), can_get_cookies_callbacks_.end(),
+      [&ctx](brave::OnCanGetCookiesCallback callback){
+        return callback.Run(ctx);
+      });
+
+  int frame_id;
+  int process_id;
+  int frame_tree_node_id;
+  brave_shields::GetRenderFrameInfo(&request, &frame_id, &process_id,
+      &frame_tree_node_id);
+  base::RepeatingCallback<content::WebContents*(void)> wc_getter =
+      base::BindRepeating(&GetWebContentsFromProcessAndFrameId, process_id,
+                          frame_id);
+  base::PostTaskWithTraits(
+      FROM_HERE, {BrowserThread::UI},
+      base::BindOnce(&TabSpecificContentSettings::CookiesRead, wc_getter,
+          request.url(), request.site_for_cookies(), cookie_list,
+          !allow));
+
+  return allow;
+}
+
+bool BraveNetworkDelegateBase::OnCanSetCookie(const URLRequest& request,
+    const net::CanonicalCookie& cookie,
+    net::CookieOptions* options,
+    bool allowed_from_caller) {
+  std::shared_ptr<brave::BraveRequestInfo> ctx(
+      new brave::BraveRequestInfo());
+  brave::BraveRequestInfo::FillCTXFromRequest(&request, ctx);
+  ctx->event_type = brave::kOnCanSetCookies;
+  bool allow = std::all_of(can_set_cookies_callbacks_.begin(), can_set_cookies_callbacks_.end(),
+      [&ctx](brave::OnCanSetCookiesCallback callback){
+        return callback.Run(ctx);
+      });
+
+  int frame_id;
+  int process_id;
+  int frame_tree_node_id;
+  brave_shields::GetRenderFrameInfo(&request, &frame_id, &process_id,
+      &frame_tree_node_id);
+  base::RepeatingCallback<content::WebContents*(void)> wc_getter =
+      base::BindRepeating(&GetWebContentsFromProcessAndFrameId, process_id,
+                          frame_id);
+  base::PostTaskWithTraits(
+      FROM_HERE, {BrowserThread::UI},
+      base::BindOnce(&TabSpecificContentSettings::CookieChanged, wc_getter,
+          request.url(), request.site_for_cookies(), cookie,
+          !allow));
+
+  return allow;
+}
+
 void BraveNetworkDelegateBase::RunCallbackForRequestIdentifier(uint64_t request_identifier, int rv) {
   std::map<uint64_t, net::CompletionOnceCallback>::iterator it =
       callbacks_.find(request_identifier);
   std::move(it->second).Run(rv);
 }
 
 void BraveNetworkDelegateBase::RunNextCallback(
-    net::URLRequest* request,
+    URLRequest* request,
     std::shared_ptr<brave::BraveRequestInfo> ctx) {
   DCHECK_CURRENTLY_ON(BrowserThread::IO);
 
@@ -206,7 +285,7 @@ void BraveNetworkDelegateBase::RunNextCallback(
   }
 }
 
-void BraveNetworkDelegateBase::OnURLRequestDestroyed(net::URLRequest* request) {
+void BraveNetworkDelegateBase::OnURLRequestDestroyed(URLRequest* request) {
   if (ContainsKey(callbacks_, request->identifier())) {
     callbacks_.erase(request->identifier());
   }

browser/net/brave_network_delegate_base.h

@@ -45,6 +45,15 @@ class BraveNetworkDelegateBase : public ChromeNetworkDelegate {
       scoped_refptr<net::HttpResponseHeaders>* override_response_headers,
       GURL* allowed_unsafe_redirect_url) override;
 
+  bool OnCanGetCookies(const net::URLRequest& request,
+                       const net::CookieList& cookie_list,
+                       bool allowed_from_caller) override;
+
+  bool OnCanSetCookie(const net::URLRequest& request,
+                      const net::CanonicalCookie& cookie,
+                      net::CookieOptions* options,
+                      bool allowed_from_caller) override;
+
   void OnURLRequestDestroyed(net::URLRequest* request) override;
   void RunCallbackForRequestIdentifier(uint64_t request_identifier, int rv);
 
@@ -58,6 +67,10 @@ class BraveNetworkDelegateBase : public ChromeNetworkDelegate {
       before_start_transaction_callbacks_;
   std::vector<brave::OnHeadersReceivedCallback>
       headers_received_callbacks_;
+  std::vector<brave::OnCanGetCookiesCallback>
+      can_get_cookies_callbacks_;
+  std::vector<brave::OnCanSetCookiesCallback>
+      can_set_cookies_callbacks_;
 
  private:
   void GetReferralHeaders();

browser/net/brave_profile_network_delegate.cc

@@ -6,6 +6,7 @@
 
 #include "brave/browser/net/brave_ad_block_tp_network_delegate_helper.h"
 #include "brave/browser/net/brave_common_static_redirect_network_delegate_helper.h"
+#include "brave/browser/net/cookie_network_delegate_helper.h"
 #include "brave/browser/net/brave_httpse_network_delegate_helper.h"
 #include "brave/browser/net/brave_referrals_network_delegate_helper.h"
 #include "brave/browser/net/brave_site_hacks_network_delegate_helper.h"
@@ -56,6 +57,14 @@ BraveProfileNetworkDelegate::BraveProfileNetworkDelegate(
       base::Bind(
           webtorrent::OnHeadersReceived_TorrentRedirectWork);
   headers_received_callbacks_.push_back(headers_received_callback);
+
+  brave::OnCanGetCookiesCallback get_cookies_callback =
+      base::Bind(brave::OnCanGetCookiesForBraveShields);
+  can_get_cookies_callbacks_.push_back(get_cookies_callback);
+
+  brave::OnCanSetCookiesCallback set_cookies_callback =
+      base::Bind(brave::OnCanSetCookiesForBraveShields);
+  can_set_cookies_callbacks_.push_back(set_cookies_callback);
 }
 
 BraveProfileNetworkDelegate::~BraveProfileNetworkDelegate() {

browser/net/brave_site_hacks_network_delegate_helper.cc

@@ -54,7 +54,7 @@ int OnBeforeURLRequest_SiteHacksWork(
     const ResponseCallback& next_callback,
     std::shared_ptr<BraveRequestInfo> ctx) {
 
-  if (ApplyPotentialReferrerBlock(ctx->request)) {
+  if (ApplyPotentialReferrerBlock(const_cast<net::URLRequest*>(ctx->request))) {
     ctx->new_url_spec = ctx->request_url.spec();
     ctx->referrer_changed = true;
   }

browser/net/cookie_network_delegate_helper.cc

@@ -0,0 +1,21 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this file,
+ * You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "brave/browser/net/cookie_network_delegate_helper.h"
+
+#include "brave/common/brave_cookie_blocking.h"
+
+namespace brave {
+
+bool OnCanGetCookiesForBraveShields(std::shared_ptr<BraveRequestInfo> ctx) {
+  return !ShouldBlockCookie(ctx->allow_brave_shields, ctx->allow_1p_cookies,
+    ctx->allow_3p_cookies, ctx->tab_origin, ctx->request_url);
+}
+
+bool OnCanSetCookiesForBraveShields(std::shared_ptr<BraveRequestInfo> ctx) {
+  return !ShouldBlockCookie(ctx->allow_brave_shields, ctx->allow_1p_cookies,
+    ctx->allow_3p_cookies, ctx->tab_origin, ctx->request_url);
+}
+
+}  // namespace brave