From 57aeb9dfc6845fa32f65d59434af4b7beae0b8a4 Mon Sep 17 00:00:00 2001
From: "trop[bot]" <37223003+trop[bot]@users.noreply.github.com>
Date: Wed, 11 Sep 2024 17:56:23 +0200
Subject: [PATCH] fix: ensure `SetPluginCanSave` updated in PDFs (#43686)

fix: ensure SetPluginCanSave updated in PDFs

Co-authored-by: trop[bot] <37223003+trop[bot]@users.noreply.github.com>
Co-authored-by: Shelley Vohr <shelley.vohr@gmail.com>
---
 .../electron_pdf_document_helper_client.cc    | 36 +++++++++++++++++--
 .../electron_pdf_document_helper_client.h     |  3 +-
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/shell/browser/electron_pdf_document_helper_client.cc b/shell/browser/electron_pdf_document_helper_client.cc
index 6902f4c22f5..6a66f6d067e 100644
--- a/shell/browser/electron_pdf_document_helper_client.cc
+++ b/shell/browser/electron_pdf_document_helper_client.cc
@@ -4,8 +4,14 @@
 
 #include "shell/browser/electron_pdf_document_helper_client.h"
 
+#include "chrome/browser/pdf/pdf_viewer_stream_manager.h"
+#include "chrome/common/content_restriction.h"
+#include "components/pdf/browser/pdf_frame_util.h"
+#include "content/public/browser/render_frame_host.h"
 #include "content/public/browser/web_contents.h"
+#include "extensions/browser/guest_view/mime_handler_view/mime_handler_view_guest.h"
 #include "pdf/content_restriction.h"
+#include "pdf/pdf_features.h"
 #include "shell/browser/api/electron_api_web_contents.h"
 
 ElectronPDFDocumentHelperClient::ElectronPDFDocumentHelperClient() = default;
@@ -15,8 +21,8 @@ void ElectronPDFDocumentHelperClient::UpdateContentRestrictions(
     content::RenderFrameHost* render_frame_host,
     int content_restrictions) {
   // UpdateContentRestrictions potentially gets called twice from
-  // pdf/pdf_view_web_plugin.cc.  The first time it is potentially called is
-  // when loading starts and it is called with a restriction on printing.  The
+  // pdf/pdf_view_web_plugin.cc. The first time it is potentially called is
+  // when loading starts and it is called with a restriction on printing. The
   // second time it is called is when loading is finished and if printing is
   // allowed there won't be a printing restriction passed, so we can use this
   // second call to notify that the pdf document is ready to print.
@@ -31,3 +37,29 @@ void ElectronPDFDocumentHelperClient::UpdateContentRestrictions(
     }
   }
 }
+
+void ElectronPDFDocumentHelperClient::SetPluginCanSave(
+    content::RenderFrameHost* render_frame_host,
+    bool can_save) {
+  if (chrome_pdf::features::IsOopifPdfEnabled()) {
+    auto* pdf_viewer_stream_manager =
+        pdf::PdfViewerStreamManager::FromWebContents(
+            content::WebContents::FromRenderFrameHost(render_frame_host));
+    if (!pdf_viewer_stream_manager) {
+      return;
+    }
+
+    content::RenderFrameHost* embedder_host =
+        pdf_frame_util::GetEmbedderHost(render_frame_host);
+    CHECK(embedder_host);
+
+    pdf_viewer_stream_manager->SetPluginCanSave(embedder_host, can_save);
+    return;
+  }
+
+  auto* guest_view =
+      extensions::MimeHandlerViewGuest::FromRenderFrameHost(render_frame_host);
+  if (guest_view) {
+    guest_view->SetPluginCanSave(can_save);
+  }
+}
diff --git a/shell/browser/electron_pdf_document_helper_client.h b/shell/browser/electron_pdf_document_helper_client.h
index 73265544de3..c8552b33db8 100644
--- a/shell/browser/electron_pdf_document_helper_client.h
+++ b/shell/browser/electron_pdf_document_helper_client.h
@@ -17,13 +17,12 @@ class ElectronPDFDocumentHelperClient : public pdf::PDFDocumentHelperClient {
 
  private:
   // pdf::PDFDocumentHelperClient
-
   void UpdateContentRestrictions(content::RenderFrameHost* render_frame_host,
                                  int content_restrictions) override;
   void OnPDFHasUnsupportedFeature(content::WebContents* contents) override {}
   void OnSaveURL(content::WebContents* contents) override {}
   void SetPluginCanSave(content::RenderFrameHost* render_frame_host,
-                        bool can_save) override {}
+                        bool can_save) override;
 };
 
 #endif  // ELECTRON_SHELL_BROWSER_ELECTRON_PDF_DOCUMENT_HELPER_CLIENT_H_