content/browser/prerender/prerender_host_registry.h

// Copyright 2020 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 CONTENT_BROWSER_PRERENDER_PRERENDER_HOST_REGISTRY_H_
#define CONTENT_BROWSER_PRERENDER_PRERENDER_HOST_REGISTRY_H_

#include <vector>

#include "base/containers/flat_map.h"
#include "base/memory/weak_ptr.h"
#include "base/observer_list_types.h"
#include "base/types/pass_key.h"
#include "content/browser/prerender/prerender_host.h"
#include "content/browser/renderer_host/back_forward_cache_impl.h"
#include "content/common/content_export.h"
#include "third_party/blink/public/common/tokens/tokens.h"
#include "third_party/blink/public/mojom/prerender/prerender.mojom.h"
#include "url/gurl.h"
#include "url/origin.h"

namespace content {

class RenderFrameHostImpl;

// PrerenderHostRegistry creates and retains a prerender host, and reserves it
// for NavigationRequest to activate the prerendered page. This is created and
// owned by WebContentsImpl.
//
// The APIs of this class are categorized into two: APIs for triggers and APIs
// for activators.
//
// - Triggers (e.g., PrerenderProcessor) start prerendering by
//   CreateAndStartHost() and notify the registry of destruction of the trigger
//   by OnTriggerDestroyed().
// - Activators (i.e., NavigationRequest) can reserve the prerender host on
//   activation start by ReserveHostToActivate(), activate it by
//   ActivateReservedHost(), and notify the registry of completion of the
//   activation by OnActivationFinished().
class CONTENT_EXPORT PrerenderHostRegistry {
 public:
  using PassKey = base::PassKey<PrerenderHostRegistry>;

  PrerenderHostRegistry();
  ~PrerenderHostRegistry();

  PrerenderHostRegistry(const PrerenderHostRegistry&) = delete;
  PrerenderHostRegistry& operator=(const PrerenderHostRegistry&) = delete;
  PrerenderHostRegistry(PrerenderHostRegistry&&) = delete;
  PrerenderHostRegistry& operator=(PrerenderHostRegistry&&) = delete;

  class Observer : public base::CheckedObserver {
   public:
    // Called once per CreateAndStartHost() call. Does not necessarily
    // mean a host was created.
    virtual void OnTrigger(const GURL& url) {}

    // Called from the registry's destructor. The observer
    // should drop any reference to the registry.
    virtual void OnRegistryDestroyed() {}
  };

  void AddObserver(Observer* observer);
  void RemoveObserver(Observer* observer);

  // For triggers.
  // Creates and starts a host. Returns the root frame tree node id of the
  // prerendered page, which can be used as the id of the host.
  // TODO(https://crbug.com/1217045): Flatten the params and do not rely on
  // PrerenderAttributesPtr.
  int CreateAndStartHost(blink::mojom::PrerenderAttributesPtr attributes,
                         RenderFrameHostImpl& initiator_render_frame_host);

  // Cancels the host registered for `frame_tree_node_id`. The host is
  // immediately removed from the map of non-reserved or reserved hosts but
  // asynchronously destroyed so that prerendered pages can cancel themselves
  // without concern for self destruction.
  void CancelHost(int frame_tree_node_id,
                  PrerenderHost::FinalStatus final_status);

  // For activators.
  // Reserves the host to activate for a navigation for the given
  // NavigationRequest. Returns the root frame tree node id of the prerendered
  // page, which can be used as the id of the host. Returns
  // RenderFrameHost::kNoFrameTreeNodeId if it's not found or not ready for
  // activation yet. The caller is responsible for calling
  // OnActivationFinished() with the id to release the reserved host.
  int ReserveHostToActivate(NavigationRequest& navigation_request);

  // For activators.
  // Activates the host reserved by ReserveHostToActivate() and returns the
  // BackForwardCacheImpl::Entry containing the page that was activated on
  // success, or nullptr on failure.
  std::unique_ptr<BackForwardCacheImpl::Entry> ActivateReservedHost(
      int frame_tree_node_id,
      NavigationRequest& navigation_request);

  RenderFrameHostImpl* GetRenderFrameHostForReservedHost(
      int frame_tree_node_id);

  // For triggers.
  // Called from the triggers (e.g., PrerenderProcessor) when they are
  // destroyed. `frame_tree_node_id` should be the id returned by
  // CreateAndStartHost().
  void OnTriggerDestroyed(int frame_tree_node_id);

  // For activators.
  // Called from the destructor of NavigationRequest that reserved the host.
  // `frame_tree_node_id` should be the id returned by ReserveHostToActivate().
  void OnActivationFinished(int frame_tree_node_id);

  // Returns the non-reserved host with the given id. Returns nullptr if the id
  // does not match any non-reserved host.
  PrerenderHost* FindNonReservedHostById(int frame_tree_node_id);

  // Returns the reserved host with the given id. Returns nullptr if the id
  // does not match any reserved host.
  PrerenderHost* FindReservedHostById(int frame_tree_node_id);

  // Returns the main frames of FrameTrees owned by this registry's prerender
  // hosts.
  std::vector<RenderFrameHostImpl*> GetPrerenderedMainFrames();

  // Returns the non-reserved host for `prerendering_url`. Returns nullptr if
  // the URL doesn't match any non-reserved host.
  PrerenderHost* FindHostByUrlForTesting(const GURL& prerendering_url);

  base::WeakPtr<PrerenderHostRegistry> GetWeakPtr();

 private:
  void ScheduleToDeleteAbandonedHost(
      std::unique_ptr<PrerenderHost> prerender_host,
      PrerenderHost::FinalStatus final_status);
  void DeleteAbandonedHosts();

  void NotifyTrigger(const GURL& url);

  // Hosts that are not reserved for activation yet.
  // TODO(https://crbug.com/1132746): Expire prerendered contents if they are
  // not used for a while.
  base::flat_map<int, std::unique_ptr<PrerenderHost>>
      prerender_host_by_frame_tree_node_id_;

  // Hosts that are reserved for activation.
  base::flat_map<int, std::unique_ptr<PrerenderHost>>
      reserved_prerender_host_by_frame_tree_node_id_;

  // Hosts that are scheduled to be deleted asynchronously.
  // Design note: PrerenderHostRegistry should explicitly manage the hosts to be
  // deleted instead of depending on the deletion helpers like DeleteSoon() to
  // asynchronously destruct them before this instance is deleted. The helpers
  // could let the hosts and their FrameTrees outlive WebContentsImpl (the owner
  // of the registry) and results in UAF.
  std::vector<std::unique_ptr<PrerenderHost>> to_be_deleted_hosts_;

  base::ObserverList<Observer> observers_;

  base::WeakPtrFactory<PrerenderHostRegistry> weak_factory_{this};
};

}  // namespace content

#endif  // CONTENT_BROWSER_PRERENDER_PRERENDER_HOST_REGISTRY_H_