Headless: Support sending and receiving of raw protocol messages

Also adds support for finding the render tree node id, if any, for a given devtools
agent host id.

BUG=546953

Review-Url: https://codereview.chromium.org/2812253002
Cr-Commit-Position: refs/heads/master@{#465658}

headless/lib/browser/headless_devtools_client_impl.cc

@@ -31,7 +31,9 @@ HeadlessDevToolsClientImpl* HeadlessDevToolsClientImpl::From(
 
 HeadlessDevToolsClientImpl::HeadlessDevToolsClientImpl()
     : agent_host_(nullptr),
+      raw_protocol_listener_(nullptr),
       next_message_id_(0),
+      next_raw_message_id_(1),
       renderer_crashed_(false),
       accessibility_domain_(this),
       animation_domain_(this),
@@ -95,10 +97,47 @@ void HeadlessDevToolsClientImpl::DetachFromHost(
   pending_messages_.clear();
 }
 
+void HeadlessDevToolsClientImpl::SetRawProtocolListener(
+    RawProtocolListener* raw_protocol_listener) {
+  raw_protocol_listener_ = raw_protocol_listener;
+}
+
+int HeadlessDevToolsClientImpl::GetNextRawDevToolsMessageId() {
+  int id = next_raw_message_id_;
+  next_raw_message_id_ += 2;
+  return id;
+}
+
+void HeadlessDevToolsClientImpl::SendRawDevToolsMessage(
+    const std::string& json_message) {
+#ifndef NDEBUG
+  std::unique_ptr<base::Value> message =
+      base::JSONReader::Read(json_message, base::JSON_PARSE_RFC);
+  const base::DictionaryValue* message_dict;
+  int id = 0;
+  if (!message || !message->GetAsDictionary(&message_dict) ||
+      !message_dict->GetInteger("id", &id)) {
+    NOTREACHED() << "Badly formed message";
+    return;
+  }
+  DCHECK_EQ((id % 2), 1) << "Raw devtools messages must have an odd ID.";
+#endif
+
+  agent_host_->DispatchProtocolMessage(this, json_message);
+}
+
+void HeadlessDevToolsClientImpl::SendRawDevToolsMessage(
+    const base::DictionaryValue& message) {
+  std::string json_message;
+  base::JSONWriter::Write(message, &json_message);
+  SendRawDevToolsMessage(json_message);
+}
+
 void HeadlessDevToolsClientImpl::DispatchProtocolMessage(
     content::DevToolsAgentHost* agent_host,
     const std::string& json_message) {
   DCHECK_EQ(agent_host_, agent_host);
+
   std::unique_ptr<base::Value> message =
       base::JSONReader::Read(json_message, base::JSON_PARSE_RFC);
   const base::DictionaryValue* message_dict;
@@ -106,6 +145,13 @@ void HeadlessDevToolsClientImpl::DispatchProtocolMessage(
     NOTREACHED() << "Badly formed reply";
     return;
   }
+
+  if (raw_protocol_listener_ &&
+      raw_protocol_listener_->OnProtocolMessage(agent_host->GetId(),
+                                                json_message, *message_dict)) {
+    return;
+  }
+
   if (!DispatchMessageReply(*message_dict) &&
       !DispatchEvent(std::move(message), *message_dict)) {
     DLOG(ERROR) << "Unhandled protocol message: " << json_message;
@@ -314,7 +360,8 @@ void HeadlessDevToolsClientImpl::FinalizeAndSendMessage(
   if (renderer_crashed_)
     return;
   DCHECK(agent_host_);
-  int id = next_message_id_++;
+  int id = next_message_id_;
+  next_message_id_ += 2;  // We only send even numbered messages.
   message->SetInteger("id", id);
   std::string json_message;
   base::JSONWriter::Write(*message, &json_message);

headless/lib/browser/headless_devtools_client_impl.h

@@ -93,6 +93,11 @@ class HeadlessDevToolsClientImpl : public HeadlessDevToolsClient,
   service_worker::Domain* GetServiceWorker() override;
   target::Domain* GetTarget() override;
   tracing::Domain* GetTracing() override;
+  void SetRawProtocolListener(
+      RawProtocolListener* raw_protocol_listener) override;
+  int GetNextRawDevToolsMessageId() override;
+  void SendRawDevToolsMessage(const std::string& json_message) override;
+  void SendRawDevToolsMessage(const base::DictionaryValue& message) override;
 
   // content::DevToolstAgentHostClient implementation:
   void DispatchProtocolMessage(content::DevToolsAgentHost* agent_host,
@@ -151,8 +156,10 @@ class HeadlessDevToolsClientImpl : public HeadlessDevToolsClient,
                          const EventHandler* event_handler,
                          const base::DictionaryValue* result_dict);
 
-  content::DevToolsAgentHost* agent_host_;  // Not owned.
+  content::DevToolsAgentHost* agent_host_;      // Not owned.
+  RawProtocolListener* raw_protocol_listener_;  // Not owned.
   int next_message_id_;
+  int next_raw_message_id_;
   std::unordered_map<int, Callback> pending_messages_;
 
   EventHandlerMap event_handlers_;

headless/lib/browser/headless_web_contents_impl.cc

@@ -213,6 +213,35 @@ void HeadlessWebContentsImpl::RenderFrameCreated(
                                      service.service_factory,
                                      browser()->BrowserMainThread());
   }
+
+  std::string devtools_agent_host_id =
+      content::DevToolsAgentHost::GetOrCreateFor(render_frame_host)->GetId();
+  render_frame_host_to_devtools_agent_host_id_[render_frame_host] =
+      devtools_agent_host_id;
+  devtools_agent_id_to_frame_tree_node_id_[devtools_agent_host_id] =
+      render_frame_host->GetFrameTreeNodeId();
+}
+
+void HeadlessWebContentsImpl::RenderFrameDeleted(
+    content::RenderFrameHost* render_frame_host) {
+  auto find_it =
+      render_frame_host_to_devtools_agent_host_id_.find(render_frame_host);
+  if (find_it == render_frame_host_to_devtools_agent_host_id_.end())
+    return;
+
+  devtools_agent_id_to_frame_tree_node_id_.erase(find_it->second);
+  render_frame_host_to_devtools_agent_host_id_.erase(find_it);
+}
+
+bool HeadlessWebContentsImpl::GetFrameTreeNodeIdForDevToolsAgentHostId(
+    const std::string& devtools_agent_host_id,
+    int* frame_tree_node_id) const {
+  const auto& find_it =
+      devtools_agent_id_to_frame_tree_node_id_.find(devtools_agent_host_id);
+  if (find_it == devtools_agent_id_to_frame_tree_node_id_.end())
+    return false;
+  *frame_tree_node_id = find_it->second;
+  return true;
 }
 
 bool HeadlessWebContentsImpl::OpenURL(const GURL& url) {

headless/lib/browser/headless_web_contents_impl.h

@@ -55,6 +55,9 @@ class HEADLESS_EXPORT HeadlessWebContentsImpl
   void RemoveObserver(Observer* observer) override;
   HeadlessDevToolsTarget* GetDevToolsTarget() override;
   HeadlessTabSocket* GetHeadlessTabSocket() const override;
+  bool GetFrameTreeNodeIdForDevToolsAgentHostId(
+      const std::string& devtools_agent_host_id,
+      int* frame_tree_node_id) const override;
 
   // HeadlessDevToolsTarget implementation:
   bool AttachClient(HeadlessDevToolsClient* client) override;
@@ -70,6 +73,7 @@ class HEADLESS_EXPORT HeadlessWebContentsImpl
 
   // content::WebContentsObserver implementation:
   void RenderFrameCreated(content::RenderFrameHost* render_frame_host) override;
+  void RenderFrameDeleted(content::RenderFrameHost* render_frame_host) override;
 
   content::WebContents* web_contents() const;
   bool OpenURL(const GURL& url);
@@ -96,6 +100,10 @@ class HEADLESS_EXPORT HeadlessWebContentsImpl
   void InitializeScreen(const gfx::Size& initial_size);
   using MojoService = HeadlessWebContents::Builder::MojoService;
 
+  std::unordered_map<content::RenderFrameHost*, std::string>
+      render_frame_host_to_devtools_agent_host_id_;
+  std::unordered_map<std::string, int> devtools_agent_id_to_frame_tree_node_id_;
+
   class Delegate;
   std::unique_ptr<Delegate> web_contents_delegate_;
   std::unique_ptr<HeadlessWindowTreeHost> window_tree_host_;

headless/lib/headless_devtools_client_browsertest.cc

@@ -987,4 +987,39 @@ class DevToolsHeaderStrippingTest : public HeadlessAsyncDevTooledBrowserTest,
 
 HEADLESS_ASYNC_DEVTOOLED_TEST_F(DevToolsHeaderStrippingTest);
 
+class RawDevtoolsProtocolTest
+    : public HeadlessAsyncDevTooledBrowserTest,
+      public HeadlessDevToolsClient::RawProtocolListener {
+ public:
+  void RunDevTooledTest() override {
+    devtools_client_->SetRawProtocolListener(this);
+
+    base::DictionaryValue message;
+    message.SetInteger("id", devtools_client_->GetNextRawDevToolsMessageId());
+    message.SetString("method", "Runtime.evaluate");
+    std::unique_ptr<base::DictionaryValue> params(new base::DictionaryValue());
+    params->SetString("expression", "1+1");
+    message.Set("params", std::move(params));
+    devtools_client_->SendRawDevToolsMessage(message);
+  }
+
+  bool OnProtocolMessage(const std::string& devtools_agent_host_id,
+                         const std::string& json_message,
+                         const base::DictionaryValue& parsed_message) override {
+    EXPECT_EQ(
+        "{\"id\":1,\"result\":{\"result\":{\"type\":\"number\","
+        "\"value\":2,\"description\":\"2\"}}}",
+        json_message);
+
+    int frame_tree_node_id = 0;
+    EXPECT_TRUE(web_contents_->GetFrameTreeNodeIdForDevToolsAgentHostId(
+        devtools_agent_host_id, &frame_tree_node_id));
+    EXPECT_NE(0, frame_tree_node_id);
+    FinishAsynchronousTest();
+    return true;
+  }
+};
+
+HEADLESS_ASYNC_DEVTOOLED_TEST_F(RawDevtoolsProtocolTest);
+
 }  // namespace headless

headless/public/headless_devtools_client.h

@@ -146,6 +146,31 @@ class HEADLESS_EXPORT HeadlessDevToolsClient {
   virtual target::Domain* GetTarget() = 0;
   virtual tracing::Domain* GetTracing() = 0;
 
+  class HEADLESS_EXPORT RawProtocolListener {
+   public:
+    RawProtocolListener() {}
+    virtual ~RawProtocolListener() {}
+
+    // Returns true if the listener handled the message.
+    virtual bool OnProtocolMessage(
+        const std::string& devtools_agent_host_id,
+        const std::string& json_message,
+        const base::DictionaryValue& parsed_message) = 0;
+
+   private:
+    DISALLOW_COPY_AND_ASSIGN(RawProtocolListener);
+  };
+
+  virtual void SetRawProtocolListener(
+      RawProtocolListener* raw_protocol_listener) = 0;
+
+  // Generates an odd numbered ID.
+  virtual int GetNextRawDevToolsMessageId() = 0;
+
+  // The id within the message must be odd to prevent collisions.
+  virtual void SendRawDevToolsMessage(const std::string& json_message) = 0;
+  virtual void SendRawDevToolsMessage(const base::DictionaryValue& message) = 0;
+
   // TODO(skyostil): Add notification for disconnection.
 
  private:

headless/public/headless_web_contents.h

@@ -75,6 +75,12 @@ class HEADLESS_EXPORT HeadlessWebContents {
   // Returns the headless tab socket for JS -> C++ if one was created.
   virtual HeadlessTabSocket* GetHeadlessTabSocket() const = 0;
 
+  // Returns the frame tree node id associated with the |devtools_agent_host_id|
+  // if any.
+  virtual bool GetFrameTreeNodeIdForDevToolsAgentHostId(
+      const std::string& devtools_agent_host_id,
+      int* frame_tree_node_id) const = 0;
+
  private:
   friend class HeadlessWebContentsImpl;
   HeadlessWebContents() {}