electron/atom/browser/ui/webui/pdf_viewer_ui.cc

// Copyright (c) 2017 GitHub, Inc.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

#include "atom/browser/ui/webui/pdf_viewer_ui.h"

#include <map>

#include "atom/browser/atom_browser_context.h"
#include "atom/browser/ui/webui/pdf_viewer_handler.h"
#include "atom/common/atom_constants.h"
#include "components/pdf/common/pdf_messages.h"
#include "content/browser/loader/resource_dispatcher_host_impl.h"
#include "content/browser/loader/resource_request_info_impl.h"
#include "content/browser/loader/stream_resource_handler.h"
#include "content/browser/resource_context_impl.h"
#include "content/browser/streams/stream.h"
#include "content/browser/streams/stream_context.h"
#include "content/public/browser/browser_thread.h"
#include "content/public/browser/render_frame_host.h"
#include "content/public/browser/render_process_host.h"
#include "content/public/browser/render_view_host.h"
#include "content/public/browser/resource_context.h"
#include "content/public/browser/stream_handle.h"
#include "content/public/browser/stream_info.h"
#include "content/public/browser/url_data_source.h"
#include "content/public/browser/web_contents.h"
#include "grit/pdf_viewer_resources_map.h"
#include "net/base/load_flags.h"
#include "net/base/mime_util.h"
#include "net/url_request/url_request.h"
#include "net/url_request/url_request_context.h"
#include "ui/base/resource/resource_bundle.h"

using content::BrowserThread;

namespace atom {

namespace {

// Extracts the path value from the URL without the leading '/',
// which follows the mapping of names in pdf_viewer_resources_map.
std::string PathWithoutParams(const std::string& path) {
  return GURL(kPdfViewerUIOrigin + path).path().substr(1);
}

class BundledDataSource : public content::URLDataSource {
 public:
  BundledDataSource() {
    for (size_t i = 0; i < kPdfViewerResourcesSize; ++i) {
      std::string resource_path = kPdfViewerResources[i].name;
      DCHECK(path_to_resource_id_.find(resource_path) ==
             path_to_resource_id_.end());
      path_to_resource_id_[resource_path] = kPdfViewerResources[i].value;
    }
  }

  // content::URLDataSource implementation.
  std::string GetSource() const override { return kPdfViewerUIHost; }

  void StartDataRequest(
      const std::string& path,
      const content::ResourceRequestInfo::WebContentsGetter& wc_getter,
      const GotDataCallback& callback) override {
    std::string filename = PathWithoutParams(path);
    auto entry = path_to_resource_id_.find(filename);

    if (entry != path_to_resource_id_.end()) {
      int resource_id = entry->second;
      const ui::ResourceBundle& rb = ui::ResourceBundle::GetSharedInstance();
      callback.Run(rb.LoadDataResourceBytes(resource_id));
    } else {
      LOG(ERROR) << "Unable to find: " << path;
      callback.Run(new base::RefCountedString());
    }
  }

  std::string GetMimeType(const std::string& path) const override {
    std::string filename = PathWithoutParams(path);
    std::string mime_type;
    net::GetMimeTypeFromFile(
        base::FilePath::FromUTF8Unsafe(filename), &mime_type);
    return mime_type;
  }

  bool ShouldAddContentSecurityPolicy() const override { return false; }

  bool ShouldDenyXFrameOptions() const override { return false; }

  bool ShouldServeMimeTypeAsContentTypeHeader() const override { return true; }

 private:
  ~BundledDataSource() override {}

  // A map from a resource path to the resource ID.
  std::map<std::string, int> path_to_resource_id_;

  DISALLOW_COPY_AND_ASSIGN(BundledDataSource);
};

void RequestPdfResource(
    const GURL& url,
    const GURL& origin,
    int render_process_id,
    int render_view_id,
    int render_frame_id,
    content::ResourceContext* resource_context,
    base::Callback<void(std::unique_ptr<content::StreamInfo>)> callback,
    atom::LayeredResourceHandler::Delegate* delegate) {
  DCHECK_CURRENTLY_ON(BrowserThread::IO);

  const net::URLRequestContext* request_context =
      resource_context->GetRequestContext();
  std::unique_ptr<net::URLRequest> request(
      request_context->CreateRequest(url, net::DEFAULT_PRIORITY, nullptr));
  request->set_method("GET");

  content::ResourceDispatcherHostImpl::Get()->InitializeURLRequest(
      request.get(), content::Referrer(url, blink::WebReferrerPolicyDefault),
      false,  // download.
      render_process_id, render_view_id, render_frame_id, resource_context);

  content::ResourceRequestInfoImpl* info =
      content::ResourceRequestInfoImpl::ForRequest(request.get());
  content::StreamContext* stream_context =
      content::GetStreamContextForResourceContext(resource_context);

  std::unique_ptr<content::ResourceHandler> handler(
      new content::StreamResourceHandler(request.get(),
                                         stream_context->registry(), origin));
  info->set_is_stream(true);
  std::unique_ptr<content::StreamInfo> stream_info(new content::StreamInfo);
  stream_info->handle =
      static_cast<content::StreamResourceHandler*>(handler.get())
          ->stream()
          ->CreateHandle();
  stream_info->original_url = request->url();

  // Helper to fill stream response details.
  handler.reset(new atom::LayeredResourceHandler(request.get(),
                                                 std::move(handler), delegate));

  content::ResourceDispatcherHostImpl::Get()->BeginURLRequest(
      std::move(request), std::move(handler),
      false,  // download
      false,  // content_initiated (download specific)
      false,  // do_not_prompt_for_login (download specific)
      resource_context);

  callback.Run(std::move(stream_info));
}

}  // namespace

PdfViewerUI::PdfViewerUI(content::BrowserContext* browser_context,
                         content::WebUI* web_ui,
                         const std::string& src)
    : content::WebUIController(web_ui),
      content::WebContentsObserver(web_ui->GetWebContents()),
      src_(src) {
  pdf_handler_ = new PdfViewerHandler(src);
  web_ui->AddMessageHandler(pdf_handler_);
  content::URLDataSource::Add(browser_context, new BundledDataSource);
}

PdfViewerUI::~PdfViewerUI() {}

bool PdfViewerUI::OnMessageReceived(
    const IPC::Message& message,
    content::RenderFrameHost* render_frame_host) {
  bool handled = true;
  IPC_BEGIN_MESSAGE_MAP(PdfViewerUI, message)
    IPC_MESSAGE_HANDLER(PDFHostMsg_PDFSaveURLAs, OnSaveURLAs)
    IPC_MESSAGE_UNHANDLED(handled = false)
  IPC_END_MESSAGE_MAP()
  return handled;
}

void PdfViewerUI::OnPdfStreamResponseStarted(
    scoped_refptr<net::HttpResponseHeaders> headers,
    const std::string& mime_type) {
  if (headers.get())
    stream_->response_headers =
        new net::HttpResponseHeaders(headers->raw_headers());
  stream_->mime_type = mime_type;
  pdf_handler_->SetPdfResourceStream(stream_.get());
}

void PdfViewerUI::OnResponseStarted(content::ResourceResponse* response) {
  DCHECK_CURRENTLY_ON(BrowserThread::IO);

  auto resource_response_head = response->head;
  BrowserThread::PostTask(
      BrowserThread::UI, FROM_HERE,
      base::Bind(&PdfViewerUI::OnPdfStreamResponseStarted,
                 base::Unretained(this), resource_response_head.headers,
                 resource_response_head.mime_type));
}

void PdfViewerUI::OnPdfStreamCreated(
    std::unique_ptr<content::StreamInfo> stream) {
  stream_ = std::move(stream);
}

void PdfViewerUI::RenderFrameCreated(content::RenderFrameHost* rfh) {
  int render_process_id = rfh->GetProcess()->GetID();
  int render_frame_id = rfh->GetRoutingID();
  int render_view_id = rfh->GetRenderViewHost()->GetRoutingID();
  auto resource_context =
      web_contents()->GetBrowserContext()->GetResourceContext();
  BrowserThread::PostTask(
      BrowserThread::IO, FROM_HERE,
      base::Bind(
          &RequestPdfResource, GURL(src_), GURL(kPdfViewerUIOrigin),
          render_process_id, render_view_id, render_frame_id, resource_context,
          base::Bind(&PdfViewerUI::OnPdfStreamCreated, base::Unretained(this)),
          this));
}

void PdfViewerUI::OnSaveURLAs(const GURL& url,
                              const content::Referrer& referrer) {
  web_contents()->SaveFrame(url, referrer);
}

}  // namespace atom