Headless: Add an option to turn cookies on and off

Bug: 
Change-Id: I09a31dcf632c8150cd1fc6f5560af826aed89bd1
Reviewed-on: https://chromium-review.googlesource.com/721360Reviewed-by: Eric Seckler <eseckler@chromium.org>
Commit-Queue: Alex Clarke <alexclarke@chromium.org>
Cr-Commit-Position: refs/heads/master@{#509336}

headless/lib/browser/headless_browser_context_impl.cc

@@ -475,6 +475,12 @@ HeadlessBrowserContext::Builder::SetIncognitoMode(bool incognito_mode) {
   return *this;
 }
 
+HeadlessBrowserContext::Builder&
+HeadlessBrowserContext::Builder::SetAllowCookies(bool allow_cookies) {
+  options_->allow_cookies_ = allow_cookies;
+  return *this;
+}
+
 HeadlessBrowserContext::Builder&
 HeadlessBrowserContext::Builder::AddTabSocketMojoBindings() {
   std::string js_bindings =

headless/lib/browser/headless_browser_context_options.cc

@@ -74,6 +74,10 @@ bool HeadlessBrowserContextOptions::incognito_mode() const {
                                browser_options_->incognito_mode);
 }
 
+bool HeadlessBrowserContextOptions::allow_cookies() const {
+  return ReturnOverriddenValue(allow_cookies_, browser_options_->allow_cookies);
+}
+
 const base::Callback<void(WebPreferences*)>&
 HeadlessBrowserContextOptions::override_web_preferences_callback() const {
   return ReturnOverriddenValue(

headless/lib/browser/headless_browser_context_options.h

@@ -44,6 +44,8 @@ class HeadlessBrowserContextOptions {
   // Set HeadlessBrowser::Options::incognito_mode.
   bool incognito_mode() const;
 
+  bool allow_cookies() const;
+
   // Custom network protocol handlers. These can be used to override URL
   // fetching for different network schemes.
   const ProtocolHandlerMap& protocol_handlers() const;
@@ -70,6 +72,7 @@ class HeadlessBrowserContextOptions {
   base::Optional<gfx::Size> window_size_;
   base::Optional<base::FilePath> user_data_dir_;
   base::Optional<bool> incognito_mode_;
+  base::Optional<bool> allow_cookies_;
   base::Optional<base::Callback<void(WebPreferences*)>>
       override_web_preferences_callback_;
 

headless/lib/browser/headless_browser_impl.cc

@@ -103,7 +103,10 @@ void HeadlessBrowserImpl::Shutdown() {
   weak_ptr_factory_.InvalidateWeakPtrs();
 
   // Destroy all browser contexts.
-  browser_contexts_.clear();
+  {
+    base::AutoLock lock(browser_contexts_lock_);
+    browser_contexts_.clear();
+  }
 
   BrowserMainThread()->PostTask(FROM_HERE,
                                 base::MessageLoop::QuitWhenIdleClosure());
@@ -112,6 +115,7 @@ void HeadlessBrowserImpl::Shutdown() {
 std::vector<HeadlessBrowserContext*>
 HeadlessBrowserImpl::GetAllBrowserContexts() {
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+  base::AutoLock lock(browser_contexts_lock_);
 
   std::vector<HeadlessBrowserContext*> result;
   result.reserve(browser_contexts_.size());
@@ -162,6 +166,7 @@ HeadlessBrowserContext* HeadlessBrowserImpl::CreateBrowserContext(
 
   HeadlessBrowserContext* result = browser_context.get();
 
+  base::AutoLock lock(browser_contexts_lock_);
   browser_contexts_[browser_context->Id()] = std::move(browser_context);
 
   return result;
@@ -169,6 +174,7 @@ HeadlessBrowserContext* HeadlessBrowserImpl::CreateBrowserContext(
 
 void HeadlessBrowserImpl::DestroyBrowserContext(
     HeadlessBrowserContextImpl* browser_context) {
+  base::AutoLock lock(browser_contexts_lock_);
   auto it = browser_contexts_.find(browser_context->Id());
   DCHECK(it != browser_contexts_.end());
   browser_contexts_.erase(it);
@@ -218,12 +224,26 @@ HeadlessWebContentsImpl* HeadlessBrowserImpl::GetWebContentsForWindowId(
 
 HeadlessBrowserContext* HeadlessBrowserImpl::GetBrowserContextForId(
     const std::string& id) {
+  base::AutoLock lock(browser_contexts_lock_);
   auto find_it = browser_contexts_.find(id);
   if (find_it == browser_contexts_.end())
     return nullptr;
   return find_it->second.get();
 }
 
+LockedPtr<HeadlessBrowserContextImpl>
+HeadlessBrowserImpl::GetBrowserContextForRenderFrame(
+    int render_process_id,
+    int render_frame_id) const {
+  MoveableAutoLock lock(browser_contexts_lock_);
+  for (const auto& pair : browser_contexts_) {
+    if (pair.second->GetDevToolsFrameToken(render_process_id, render_frame_id))
+      return LockedPtr<HeadlessBrowserContextImpl>(std::move(lock),
+                                                   pair.second.get());
+  }
+  return LockedPtr<HeadlessBrowserContextImpl>(std::move(lock), nullptr);
+}
+
 HeadlessDevToolsTarget* HeadlessBrowserImpl::GetDevToolsTarget() {
   return agent_host_ ? this : nullptr;
 }

headless/lib/browser/headless_browser_impl.h

@@ -18,6 +18,7 @@
 #include "headless/lib/browser/headless_devtools_manager_delegate.h"
 #include "headless/lib/browser/headless_web_contents_impl.h"
 #include "headless/public/headless_export.h"
+#include "headless/public/util/moveable_auto_lock.h"
 
 namespace ui {
 class Compositor;
@@ -78,6 +79,12 @@ class HEADLESS_EXPORT HeadlessBrowserImpl : public HeadlessBrowser,
 
   base::WeakPtr<HeadlessBrowserImpl> GetWeakPtr();
 
+  // Returns the corresponding HeadlessBrowserContextImpl or null if one can't
+  // be found. Can be called on any thread.
+  LockedPtr<HeadlessBrowserContextImpl> GetBrowserContextForRenderFrame(
+      int render_process_id,
+      int render_frame_id) const;
+
   // All the methods that begin with Platform need to be implemented by the
   // platform specific headless implementation.
   // Helper for one time initialization of application
@@ -93,7 +100,8 @@ class HEADLESS_EXPORT HeadlessBrowserImpl : public HeadlessBrowser,
   HeadlessBrowser::Options options_;
   HeadlessBrowserMainParts* browser_main_parts_;  // Not owned.
 
-  std::unordered_map<std::string, std::unique_ptr<HeadlessBrowserContextImpl>>
+  mutable base::Lock browser_contexts_lock_;  // Protects |browser_contexts_|
+  base::flat_map<std::string, std::unique_ptr<HeadlessBrowserContextImpl>>
       browser_contexts_;
   HeadlessBrowserContext* default_browser_context_;  // Not owned.
 

headless/lib/browser/headless_content_browser_client.cc

@@ -14,6 +14,7 @@
 #include "base/memory/ptr_util.h"
 #include "base/path_service.h"
 #include "content/public/browser/browser_context.h"
+#include "content/public/browser/browser_thread.h"
 #include "content/public/browser/render_process_host.h"
 #include "content/public/browser/render_view_host.h"
 #include "content/public/browser/resource_dispatcher_host.h"
@@ -300,4 +301,37 @@ net::NetLog* HeadlessContentBrowserClient::GetNetLog() {
   return browser_->browser_main_parts()->net_log();
 }
 
+bool HeadlessContentBrowserClient::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(content::BrowserThread::IO);
+  LockedPtr<HeadlessBrowserContextImpl> browser_context =
+      browser_->GetBrowserContextForRenderFrame(render_process_id,
+                                                render_frame_id);
+  if (!browser_context)
+    return false;
+  return browser_context->options()->allow_cookies();
+}
+
+bool HeadlessContentBrowserClient::AllowSetCookie(
+    const GURL& url,
+    const GURL& first_party,
+    const std::string& cookie_line,
+    content::ResourceContext* context,
+    int render_process_id,
+    int render_frame_id,
+    const net::CookieOptions& options) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
+  LockedPtr<HeadlessBrowserContextImpl> browser_context =
+      browser_->GetBrowserContextForRenderFrame(render_process_id,
+                                                render_frame_id);
+  if (!browser_context)
+    return false;
+  return browser_context->options()->allow_cookies();
+}
+
 }  // namespace headless

headless/lib/browser/headless_content_browser_client.h

@@ -55,6 +55,21 @@ class HeadlessContentBrowserClient : public content::ContentBrowserClient {
 
   net::NetLog* GetNetLog() override;
 
+  bool AllowGetCookie(const GURL& url,
+                      const GURL& first_party,
+                      const net::CookieList& cookie_list,
+                      content::ResourceContext* context,
+                      int render_process_id,
+                      int render_frame_id) override;
+
+  bool AllowSetCookie(const GURL& url,
+                      const GURL& first_party,
+                      const std::string& cookie_line,
+                      content::ResourceContext* context,
+                      int render_process_id,
+                      int render_frame_id,
+                      const net::CookieOptions& options) override;
+
  private:
   std::unique_ptr<base::Value> GetBrowserServiceManifestOverlay();
   std::unique_ptr<base::Value> GetRendererServiceManifestOverlay();

headless/lib/headless_web_contents_browsertest.cc

@@ -1354,4 +1354,74 @@ HEADLESS_ASYNC_DEVTOOLED_TEST_F(
 
 #endif  // !defined(OS_MACOSX)
 
+class CookiesEnabled : public HeadlessAsyncDevTooledBrowserTest,
+                       page::Observer {
+ public:
+  void RunDevTooledTest() override {
+    devtools_client_->GetPage()->AddObserver(this);
+    devtools_client_->GetPage()->Enable();
+
+    EXPECT_TRUE(embedded_test_server()->Start());
+    devtools_client_->GetPage()->Navigate(
+        embedded_test_server()->GetURL("/cookie.html").spec());
+  }
+
+  // page::Observer implementation:
+  void OnLoadEventFired(const page::LoadEventFiredParams& params) override {
+    devtools_client_->GetRuntime()->Evaluate(
+        "window.test_result",
+        base::Bind(&CookiesEnabled::OnResult, base::Unretained(this)));
+  }
+
+  void OnResult(std::unique_ptr<runtime::EvaluateResult> result) {
+    std::string value;
+    EXPECT_TRUE(result->GetResult()->HasValue());
+    EXPECT_TRUE(result->GetResult()->GetValue()->GetAsString(&value));
+    EXPECT_EQ("0", value);
+    FinishAsynchronousTest();
+  }
+
+  void CustomizeHeadlessBrowserContext(
+      HeadlessBrowserContext::Builder& builder) override {
+    builder.SetAllowCookies(true);
+  }
+};
+
+HEADLESS_ASYNC_DEVTOOLED_TEST_F(CookiesEnabled);
+
+class CookiesDisabled : public HeadlessAsyncDevTooledBrowserTest,
+                        page::Observer {
+ public:
+  void RunDevTooledTest() override {
+    devtools_client_->GetPage()->AddObserver(this);
+    devtools_client_->GetPage()->Enable();
+
+    EXPECT_TRUE(embedded_test_server()->Start());
+    devtools_client_->GetPage()->Navigate(
+        embedded_test_server()->GetURL("/cookie.html").spec());
+  }
+
+  // page::Observer implementation:
+  void OnLoadEventFired(const page::LoadEventFiredParams& params) override {
+    devtools_client_->GetRuntime()->Evaluate(
+        "window.test_result",
+        base::Bind(&CookiesDisabled::OnResult, base::Unretained(this)));
+  }
+
+  void OnResult(std::unique_ptr<runtime::EvaluateResult> result) {
+    std::string value;
+    EXPECT_TRUE(result->GetResult()->HasValue());
+    EXPECT_TRUE(result->GetResult()->GetValue()->GetAsString(&value));
+    EXPECT_EQ("-1", value);
+    FinishAsynchronousTest();
+  }
+
+  void CustomizeHeadlessBrowserContext(
+      HeadlessBrowserContext::Builder& builder) override {
+    builder.SetAllowCookies(false);
+  }
+};
+
+HEADLESS_ASYNC_DEVTOOLED_TEST_F(CookiesDisabled);
+
 }  // namespace headless

headless/public/headless_browser.cc

@@ -56,6 +56,7 @@ Options::Options(int argc, const char** argv)
       user_agent(content::BuildUserAgentFromProduct(product_name_and_version)),
       window_size(kDefaultWindowSize),
       incognito_mode(true),
+      allow_cookies(true),
       enable_crash_reporter(false) {
 }
 

headless/public/headless_browser.h

@@ -173,6 +173,9 @@ struct HEADLESS_EXPORT HeadlessBrowser::Options {
   // Run a browser context in an incognito mode. Enabled by default.
   bool incognito_mode;
 
+  // Whether cookies are allowed. Enabled by default.
+  bool allow_cookies;
+
   // Set a callback that is invoked to override WebPreferences for RenderViews
   // created within the HeadlessBrowser. Called whenever the WebPreferences of a
   // RenderView change. Executed on the browser main thread.
@@ -227,6 +230,7 @@ class HEADLESS_EXPORT HeadlessBrowser::Options::Builder {
   Builder& SetWindowSize(const gfx::Size& window_size);
   Builder& SetUserDataDir(const base::FilePath& user_data_dir);
   Builder& SetIncognitoMode(bool incognito_mode);
+  Builder& SetAllowCookies(bool allow_cookies);
   Builder& SetOverrideWebPreferencesCallback(
       base::Callback<void(WebPreferences*)> callback);
   Builder& SetCrashReporterEnabled(bool enabled);

headless/public/headless_browser_context.h

@@ -132,6 +132,7 @@ class HEADLESS_EXPORT HeadlessBrowserContext::Builder {
   Builder& SetWindowSize(const gfx::Size& window_size);
   Builder& SetUserDataDir(const base::FilePath& user_data_dir);
   Builder& SetIncognitoMode(bool incognito_mode);
+  Builder& SetAllowCookies(bool incognito_mode);
   Builder& SetOverrideWebPreferencesCallback(
       base::Callback<void(WebPreferences*)> callback);
 

headless/public/util/moveable_auto_lock.h

+// Copyright 2017 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.
+
+#ifndef HEADLESS_PUBLIC_UTIL_MOVEABLE_AUTO_LOCK_H_
+#define HEADLESS_PUBLIC_UTIL_MOVEABLE_AUTO_LOCK_H_
+
+#include "base/synchronization/lock.h"
+
+namespace headless {
+
+class MoveableAutoLock {
+ public:
+  explicit MoveableAutoLock(base::Lock& lock) : lock_(lock), moved_(false) {
+    lock_.Acquire();
+  }
+
+  MoveableAutoLock(MoveableAutoLock&& other)
+      : lock_(other.lock_), moved_(other.moved_) {
+    lock_.AssertAcquired();
+    other.moved_ = true;
+  }
+
+  ~MoveableAutoLock() {
+    if (moved_)
+      return;
+    lock_.AssertAcquired();
+    lock_.Release();
+  }
+
+ private:
+  base::Lock& lock_;
+  bool moved_;
+  DISALLOW_COPY_AND_ASSIGN(MoveableAutoLock);
+};
+
+// RAII helper to allow threadsafe access to an object guarded by a lock.
+template <class T>
+class LockedPtr {
+ public:
+  LockedPtr(MoveableAutoLock&& lock, T* object)
+      : lock_(std::move(lock)), object_(object) {}
+
+  LockedPtr(LockedPtr&& other)
+      : lock_(std::move(other.lock_)), object_(other.object_) {}
+
+  T* operator->() { return object_; }
+
+  T& operator*() { return *object_; }
+
+  T* Get() { return object_; }
+
+  explicit operator bool() const { return object_; }
+
+ private:
+  MoveableAutoLock lock_;
+  T* object_;
+
+  DISALLOW_COPY_AND_ASSIGN(LockedPtr);
+};
+
+}  // namespace headless
+
+#endif  // HEADLESS_PUBLIC_UTIL_MOVEABLE_AUTO_LOCK_H_

headless/test/data/cookie.html

+<html>
+<script>
+document.cookie = "SessionID=1234567890";
+window.test_result = JSON.stringify(document.cookie.indexOf('SessionID'));
+</script>
+</html>