DevTools: Fix getCookies to report for all resources

content/browser/devtools/protocol/network_handler.cc

 // Copyright 2014 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "content/browser/devtools/protocol/network_handler.h"
 
 #include <stddef.h>
 
 #include "base/containers/hash_tables.h"
 #include "base/strings/stringprintf.h"
@@ skipping 16 matching lines @@
 
 namespace content {
 namespace protocol {
 namespace {
 
 using GetCookiesCallback = protocol::Network::Backend::GetCookiesCallback;
 using GetAllCookiesCallback = protocol::Network::Backend::GetAllCookiesCallback;
 using SetCookieCallback = protocol::Network::Backend::SetCookieCallback;
 using DeleteCookieCallback = protocol::Network::Backend::DeleteCookieCallback;
 
+class CookieReceiver {

[pfeldman, 2017/01/18 01:41:11]

nit: inherit from base::NonThreadSafe and DCHECK(CalledOnValidThread()) below.

nit: also, making it ref-counted-thread-safe would allow graceful threads
termination.

+  public:
+    CookieReceiver(const net::CookieStore::GetCookieListCallback& callback)
+        : callback_(callback), request_count_(0) {}
+
+    void increment(int count) {
+      request_count_ += count;

[pfeldman, 2017/01/18 01:41:11]

nit: use base::BarrierClosure instead - your callback would receive it, do the
work and call barrier's closure.

+    }
+
+    net::CookieStore::GetCookieListCallback bound() {
+      return base::Bind(&CookieReceiver::GotCookies, base::Unretained(this));
+    }
+
+  protected:
+    void GotCookies(const net::CookieList& cookie_list) {
+      DCHECK_CURRENTLY_ON(BrowserThread::IO);
+      for (const net::CanonicalCookie& cookie : cookie_list) {
+        std::string key = base::StringPrintf(
+            "%s::%s::%s::%d", cookie.Name().c_str(), cookie.Domain().c_str(),
+            cookie.Path().c_str(), cookie.IsSecure());
+        cookies_[key] = cookie;
+      }
+      --request_count_;
+      if (!request_count_) {
+        SendResponse();
+        delete this;
+      }
+    }
+
+    void SendResponse() {
+      net::CookieList master_cookie_list;
+      for (const auto& pair : cookies_) {
+        master_cookie_list.push_back(pair.second);
+      }
+
+      BrowserThread::PostTask(
+          BrowserThread::UI,
+          FROM_HERE,
+          base::Bind(callback_, master_cookie_list));
+    }
+
+    const net::CookieStore::GetCookieListCallback& callback_;
+    int request_count_;
+    base::hash_map<std::string, net::CanonicalCookie> cookies_;
+};
+
 net::URLRequestContext* GetRequestContextOnIO(
     ResourceContext* resource_context,
     net::URLRequestContextGetter* context_getter,
     const GURL& url) {
   DCHECK_CURRENTLY_ON(BrowserThread::IO);
   net::URLRequestContext* context =
       GetContentClient()->browser()->OverrideRequestContextForURL(
           url, resource_context);
   if (!context)
     context = context_getter->GetURLRequestContext();
   return context;
 }
 
-void GotCookiesOnIO(
-    const net::CookieStore::GetCookieListCallback& callback,
-    const net::CookieList& cookie_list) {
-  DCHECK_CURRENTLY_ON(BrowserThread::IO);
-  BrowserThread::PostTask(
-      BrowserThread::UI,
-      FROM_HERE,
-      base::Bind(callback, cookie_list));
-}
-
-void GetCookiesForURLOnIO(
+void GetCookiesForURLsOnIO(
     ResourceContext* resource_context,
     net::URLRequestContextGetter* context_getter,
-    const GURL& url,
-    const net::CookieStore::GetCookieListCallback& callback) {
+    const std::vector<GURL>& urls,
+    const net::CookieStore::GetCookieListCallback& finished_callback) {
   DCHECK_CURRENTLY_ON(BrowserThread::IO);
-  net::URLRequestContext* request_context =
-      GetRequestContextOnIO(resource_context, context_getter, url);
-  request_context->cookie_store()->GetAllCookiesForURLAsync(
-      url, base::Bind(&GotCookiesOnIO, callback));
+  CookieReceiver* receiver = new CookieReceiver(finished_callback);
+  receiver->increment(urls.size());
+
+  net::CookieStore::GetCookieListCallback callback = receiver->bound();
+
+  for (const GURL& url : urls) {
+    net::URLRequestContext* request_context =
+        GetRequestContextOnIO(resource_context, context_getter, url);
+    request_context->cookie_store()->GetAllCookiesForURLAsync(url, callback);
+  }
 }
 
 void GetAllCookiesOnIO(
     ResourceContext* resource_context,
     net::URLRequestContextGetter* context_getter,
-    const net::CookieStore::GetCookieListCallback& callback) {
+    const net::CookieStore::GetCookieListCallback& finished_callback) {
   DCHECK_CURRENTLY_ON(BrowserThread::IO);
+  CookieReceiver* receiver = new CookieReceiver(finished_callback);
+  receiver->increment(1);
+
+  net::CookieStore::GetCookieListCallback callback = receiver->bound();
   net::URLRequestContext* request_context =
       context_getter->GetURLRequestContext();
-  request_context->cookie_store()->GetAllCookiesAsync(
-      base::Bind(&GotCookiesOnIO, callback));
+  request_context->cookie_store()->GetAllCookiesAsync(callback);
 }
 
 void DeletedCookieOnIO(std::unique_ptr<DeleteCookieCallback> callback) {
   DCHECK_CURRENTLY_ON(BrowserThread::IO);
   BrowserThread::PostTask(
       BrowserThread::UI,
       FROM_HERE,
       base::Bind(&DeleteCookieCallback::sendSuccess,
                  base::Passed(std::move(callback))));
 }
@@ skipping 53 matching lines @@
       same_site,
       are_experimental_cookie_features_enabled,
       net::COOKIE_PRIORITY_DEFAULT,
       base::Bind(&CookieSetOnIO, base::Passed(std::move(callback))));
 }
 
 template <typename Callback>
 class GetCookiesCommandBase {
  public:
   GetCookiesCommandBase(std::unique_ptr<Callback> callback)
-      : callback_(std::move(callback)), request_count_(0) {}
+      : callback_(std::move(callback)) {}
 
  protected:
-  void GotCookiesForURL(const net::CookieList& cookie_list) {
-    DCHECK_CURRENTLY_ON(BrowserThread::UI);
-    for (const net::CanonicalCookie& cookie : cookie_list) {
-      std::string key = base::StringPrintf(
-          "%s::%s::%s::%d", cookie.Name().c_str(), cookie.Domain().c_str(),
-          cookie.Path().c_str(), cookie.IsSecure());
-      cookies_[key] = cookie;
-    }
-    --request_count_;
-    if (!request_count_) {
-      SendResponse();
-      delete this;
-    }
-  }
-
-  void SendResponse() {
+  void SendResponse(const net::CookieList& cookie_list) {
     std::unique_ptr<protocol::Array<Network::Cookie>> cookies =
         protocol::Array<Network::Cookie>::create();
-    for (const auto& pair : cookies_) {
-      const net::CanonicalCookie& cookie = pair.second;
+    for (const auto& cookie : cookie_list) {
       std::unique_ptr<Network::Cookie> devtools_cookie =
            Network::Cookie::Create()
            .SetName(cookie.Name())
            .SetValue(cookie.Value())
            .SetDomain(cookie.Domain())
            .SetPath(cookie.Path())
            .SetExpires(cookie.ExpiryDate().ToDoubleT() * 1000)
            .SetSize(cookie.Name().length() + cookie.Value().length())
            .SetHttpOnly(cookie.IsHttpOnly())
            .SetSecure(cookie.IsSecure())
            .SetSession(!cookie.IsPersistent())
            .Build();
 
       switch (cookie.SameSite()) {
         case net::CookieSameSite::STRICT_MODE:
           devtools_cookie->SetSameSite(Network::CookieSameSiteEnum::Strict);
           break;
         case net::CookieSameSite::LAX_MODE:
           devtools_cookie->SetSameSite(Network::CookieSameSiteEnum::Lax);
           break;
         case net::CookieSameSite::NO_RESTRICTION:
           break;
       }
       cookies->addItem(std::move(devtools_cookie));
     }
     callback_->sendSuccess(std::move(cookies));
   }
 
   std::unique_ptr<Callback> callback_;
-  int request_count_;
-  base::hash_map<std::string, net::CanonicalCookie> cookies_;
 };
 
 class GetCookiesCommand : public GetCookiesCommandBase<GetCookiesCallback> {
  public:
   GetCookiesCommand(RenderFrameHostImpl* frame_host,
+                    Maybe<protocol::Array<std::string>>& maybe_urls,
                     std::unique_ptr<GetCookiesCallback> callback)
       : GetCookiesCommandBase(std::move(callback)) {
+    std::vector<GURL> urls_to_fetch;
     net::CookieStore::GetCookieListCallback got_cookies_callback = base::Bind(
-        &GetCookiesCommand::GotCookiesForURL, base::Unretained(this));
+        &GetCookiesCommand::SendResponse, base::Unretained(this));

[pfeldman, 2017/01/18 01:41:11]

This is dangerous - when your callback triggers, this object can no longer be
there. You need to use WeakPtr factory here.

 
-    std::queue<FrameTreeNode*> queue;
-    queue.push(frame_host->frame_tree_node());
-    while (!queue.empty()) {
-      FrameTreeNode* node = queue.front();
-      queue.pop();
+    if (maybe_urls.isJust()) {
+      std::unique_ptr<protocol::Array<std::string>> urls =
+          maybe_urls.takeJust();
 
-      // Only traverse nodes with the same local root.
-      if (node->current_frame_host()->IsCrossProcessSubframe())
-        continue;
-      ++request_count_;
-      BrowserThread::PostTask(
-          BrowserThread::IO, FROM_HERE,
-          base::Bind(&GetCookiesForURLOnIO,
-                     base::Unretained(frame_host->GetSiteInstance()
-                                          ->GetBrowserContext()
-                                          ->GetResourceContext()),
-                     base::Unretained(frame_host->GetProcess()
-                                          ->GetStoragePartition()
-                                          ->GetURLRequestContext()),
-                     node->current_url(), got_cookies_callback));
+      for (size_t i = 0; i < urls->length(); i++)
+        urls_to_fetch.push_back(GURL(urls->get(i)));
+    } else {
+      std::queue<FrameTreeNode*> queue;
+      queue.push(frame_host->frame_tree_node());
+      while (!queue.empty()) {
+        FrameTreeNode* node = queue.front();
+        queue.pop();
 
-      for (size_t i = 0; i < node->child_count(); ++i)
-        queue.push(node->child_at(i));
+        // Only traverse nodes with the same local root.
+        if (node->current_frame_host()->IsCrossProcessSubframe())
+          continue;
+
+        urls_to_fetch.push_back(node->current_url());
+        for (size_t i = 0; i < node->child_count(); ++i)
+          queue.push(node->child_at(i));
+      }
     }
+
+    if (urls_to_fetch.size() == 0) {
+      got_cookies_callback();
+      return;
+    }
+
+    BrowserThread::PostTask(
+        BrowserThread::IO, FROM_HERE,
+        base::Bind(&GetCookiesForURLsOnIO,
+                   base::Unretained(frame_host->GetSiteInstance()
+                                        ->GetBrowserContext()
+                                        ->GetResourceContext()),
+                   base::Unretained(frame_host->GetProcess()
+                                        ->GetStoragePartition()
+                                        ->GetURLRequestContext()),
+                   urls_to_fetch, got_cookies_callback));
   }
 };
 
 class GetAllCookiesCommand
     : public GetCookiesCommandBase<GetAllCookiesCallback> {
  public:
   GetAllCookiesCommand(RenderFrameHostImpl* frame_host,
                        std::unique_ptr<GetAllCookiesCallback> callback)
       : GetCookiesCommandBase(std::move(callback)) {
     net::CookieStore::GetCookieListCallback got_cookies_callback = base::Bind(
-        &GetAllCookiesCommand::GotCookiesForURL, base::Unretained(this));
+        &GetAllCookiesCommand::SendResponse, base::Unretained(this));
 
-    request_count_ = 1;
     BrowserThread::PostTask(
         BrowserThread::IO, FROM_HERE,
         base::Bind(&GetAllCookiesOnIO,
                    base::Unretained(frame_host->GetSiteInstance()
                                         ->GetBrowserContext()
                                         ->GetResourceContext()),
                    base::Unretained(frame_host->GetProcess()
                                         ->GetStoragePartition()
                                         ->GetURLRequestContext()),
                    got_cookies_callback));
@@ skipping 43 matching lines @@
   return Response::OK();
 }
 
 Response NetworkHandler::ClearBrowserCookies() {
   if (host_)
     GetContentClient()->browser()->ClearCookies(host_);
   return Response::OK();
 }
 
 void NetworkHandler::GetCookies(
+    Maybe<protocol::Array<String>> urls,
     std::unique_ptr<GetCookiesCallback> callback) {
   if (!host_)
     callback->sendFailure(Response::InternalError());
   else
-    new GetCookiesCommand(host_, std::move(callback));
+    new GetCookiesCommand(host_, urls, std::move(callback));
 }
 
 void NetworkHandler::GetAllCookies(
     std::unique_ptr<GetAllCookiesCallback> callback) {
   if (!host_)
     callback->sendFailure(Response::InternalError());
   else
     new GetAllCookiesCommand(host_, std::move(callback));
 }
 
@@ skipping 72 matching lines @@
   *result = false;
   return Response::OK();
 }
 
 std::string NetworkHandler::UserAgentOverride() const {
   return enabled_ ? user_agent_ : std::string();
 }
 
 }  // namespace protocol
 }  // namespace content