Support custom protocols (#320)

This patch makes the API more versatile by allowing to run custom protocols. - add :protocol option to `print_to_pdf/2` & `capture_screenshot/2` - refactor API & options handling around `chrome_export/2` - add new top-level but hidden `run_protocol/2` API For the time being, these features are not advertised in the documentation. The `ChromicPDF.Protocol` machinery can be quite tricky and the poor DSL in `ChromicPDF.ProtocolMacros` is subject to change. Nonetheless, this could be a way to avoid adding further options to change and extend the default protocols' behaviour. Relates to #319
bitcrowd · Aug 9, 2024 · dbb7577 · dbb7577
1 parent 9a4671c
commit dbb7577
Show file tree

Hide file tree

Showing 10 changed files with 204 additions and 90 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,9 +1,17 @@
 ## Unreleased
 
+### Changed
+
+- Some of the types have been renamed, e.g. `ChromicPDF.export_option` to `shared_option`.
+
 ### Fixed
 
 - Small fix `:chrome_version` config switch allowing to pass `Chrome x.y.z.zz` instead of just `x.y.z.zz`
 
+### Added
+
+- Support custom protocols through `:protocol` option on `print_to_pdf/2` and `capture_screenshot`, as well as new `ChromicPDF.run_protocol/2` function. These features are considered internal API.
+
 ## [1.16.1] - 2024-07-25
 
 ### Added

diff --git a/lib/chromic_pdf/api.ex b/lib/chromic_pdf/api.ex
@@ -10,15 +10,15 @@ defmodule ChromicPDF.API do
     CaptureScreenshot,
     ExportOptions,
     GhostscriptPool,
-    PDFOptions,
-    PrintToPDF
+    PrintToPDF,
+    ProtocolOptions
   }
 
   @spec print_to_pdf(
           ChromicPDF.Supervisor.services(),
           ChromicPDF.source() | [ChromicPDF.source()],
-          [ChromicPDF.pdf_option() | ChromicPDF.export_option()]
-        ) :: ChromicPDF.export_return()
+          [ChromicPDF.pdf_option() | ChromicPDF.shared_option()]
+        ) :: ChromicPDF.result()
   def print_to_pdf(services, sources, opts) when is_list(sources) and is_list(opts) do
     with_tmp_dir(fn tmp_dir ->
       paths =
@@ -43,41 +43,51 @@ defmodule ChromicPDF.API do
   end
 
   def print_to_pdf(services, source, opts) when is_tuple(source) and is_list(opts) do
-    chrome_export(services, :print_to_pdf, source, opts)
+    {protocol, opts} =
+      opts
+      |> ProtocolOptions.prepare_print_to_pdf_options(source)
+      |> Keyword.pop(:protocol, PrintToPDF)
+
+    chrome_export(services, :print_to_pdf, protocol, opts)
   end
 
   @spec capture_screenshot(ChromicPDF.Supervisor.services(), ChromicPDF.source(), [
-          ChromicPDF.capture_screenshot_option() | ChromicPDF.export_option()
+          ChromicPDF.capture_screenshot_option() | ChromicPDF.shared_option()
         ]) ::
-          ChromicPDF.export_return()
+          ChromicPDF.result()
   def capture_screenshot(services, %{source: source, opts: opts}, overrides)
       when is_tuple(source) and is_list(opts) and is_list(overrides) do
     capture_screenshot(services, source, Keyword.merge(opts, overrides))
   end
 
   def capture_screenshot(services, source, opts) when is_tuple(source) and is_list(opts) do
-    chrome_export(services, :capture_screenshot, source, opts)
-  end
+    {protocol, opts} =
+      opts
+      |> ProtocolOptions.prepare_capture_screenshot_options(source)
+      |> Keyword.pop(:protocol, CaptureScreenshot)
 
-  @export_protocols %{
-    capture_screenshot: CaptureScreenshot,
-    print_to_pdf: PrintToPDF
-  }
+    chrome_export(services, :capture_screenshot, protocol, opts)
+  end
 
-  defp chrome_export(services, protocol, source, opts) do
-    opts = PDFOptions.prepare_input_options(source, opts)
+  @spec run_protocol(ChromicPDF.Supervisor.services(), module(), [
+          ChromicPDF.shared_option() | ChromicPDF.protocol_option()
+        ]) :: ChromicPDF.result()
+  def run_protocol(services, protocol, opts) when is_atom(protocol) and is_list(opts) do
+    chrome_export(services, :run_protocol, protocol, opts)
+  end
 
-    with_telemetry(protocol, opts, fn ->
+  defp chrome_export(services, operation, protocol, opts) do
+    with_telemetry(operation, opts, fn ->
       services.browser
-      |> Browser.new_protocol(Map.fetch!(@export_protocols, protocol), opts)
+      |> Browser.new_protocol(protocol, opts)
       |> ExportOptions.feed_chrome_data_into_output(opts)
     end)
   end
 
   @spec convert_to_pdfa(ChromicPDF.Supervisor.services(), ChromicPDF.path(), [
-          ChromicPDF.pdfa_option() | ChromicPDF.export_option()
+          ChromicPDF.pdfa_option() | ChromicPDF.shared_option()
         ]) ::
-          ChromicPDF.export_return()
+          ChromicPDF.result()
   def convert_to_pdfa(services, pdf_path, opts) when is_binary(pdf_path) and is_list(opts) do
     with_tmp_dir(fn tmp_dir ->
       do_convert_to_pdfa(services, pdf_path, opts, tmp_dir)
@@ -88,10 +98,10 @@ defmodule ChromicPDF.API do
           ChromicPDF.Supervisor.services(),
           ChromicPDF.source() | [ChromicPDF.source()],
           [
-            ChromicPDF.pdf_option() | ChromicPDF.pdfa_option() | ChromicPDF.export_option()
+            ChromicPDF.pdf_option() | ChromicPDF.pdfa_option() | ChromicPDF.shared_option()
           ]
         ) ::
-          ChromicPDF.export_return()
+          ChromicPDF.result()
   def print_to_pdfa(services, source, opts) when is_list(opts) do
     with_tmp_dir(fn tmp_dir ->
       pdf_path = Path.join(tmp_dir, random_file_name(".pdf"))

diff --git a/lib/chromic_pdf/api/pdf_options.ex → lib/chromic_pdf/api/protocol_options.ex b/lib/chromic_pdf/api/pdf_options.ex → lib/chromic_pdf/api/protocol_options.ex
@@ -1,17 +1,30 @@
 # SPDX-License-Identifier: Apache-2.0
 
-defmodule ChromicPDF.PDFOptions do
+defmodule ChromicPDF.ProtocolOptions do
   @moduledoc false
 
   require EEx
   import ChromicPDF.Utils, only: [rendered_to_binary: 1]
 
-  def prepare_input_options(source, opts) do
+  def prepare_print_to_pdf_options(opts, source) do
+    opts
+    |> prepare_navigate_options(source)
+    |> stringify_map_keys(:print_to_pdf)
+    |> sanitize_binary_option([:print_to_pdf, "headerTemplate"])
+    |> sanitize_binary_option([:print_to_pdf, "footerTemplate"])
+  end
+
+  def prepare_capture_screenshot_options(opts, source) do
+    opts
+    |> prepare_navigate_options(source)
+    |> stringify_map_keys(:capture_screenshot)
+  end
+
+  defp prepare_navigate_options(opts, source) do
     opts
     |> put_source(source)
     |> replace_wait_for_with_evaluate()
-    |> stringify_map_keys()
-    |> sanitize_binaries()
+    |> sanitize_binary_option(:html)
   end
 
   defp put_source(opts, {:file, source}), do: put_source(opts, {:url, source})
@@ -91,30 +104,18 @@ defmodule ChromicPDF.PDFOptions do
     end)
   end
 
-  @map_options [:print_to_pdf, :capture_screenshot]
-
-  defp stringify_map_keys(opts) do
-    Enum.reduce(@map_options, opts, fn key, acc ->
-      Keyword.update(acc, key, %{}, &do_stringify_map_keys/1)
-    end)
+  def stringify_map_keys(opts, key) do
+    Keyword.update(opts, key, %{}, &do_stringify_map_keys/1)
   end
 
   defp do_stringify_map_keys(map) do
     Enum.into(map, %{}, fn {k, v} -> {to_string(k), v} end)
   end
 
-  @binary_options [
-    [:html],
-    [:print_to_pdf, "headerTemplate"],
-    [:print_to_pdf, "footerTemplate"]
-  ]
-
-  defp sanitize_binaries(opts) do
-    Enum.reduce(@binary_options, opts, fn path, acc ->
-      update_in(acc, path, fn
-        nil -> ""
-        other -> rendered_to_binary(other)
-      end)
+  defp sanitize_binary_option(opts, path) do
+    update_in(opts, List.wrap(path), fn
+      nil -> ""
+      other -> rendered_to_binary(other)
     end)
   end
 end
diff --git a/lib/chromic_pdf/pdf/browser/channel.ex b/lib/chromic_pdf/pdf/browser/channel.ex
@@ -93,6 +93,8 @@ defmodule ChromicPDF.Browser.Channel do
     :ok
   end
 
+  def terminate(_exception, _state), do: :ok
+
   defp warn_on_inspector_crash(msg) do
     if match?(%{"method" => "Inspector.targetCrashed"}, msg) do
       Logger.error("""

diff --git a/lib/chromic_pdf/pdf/protocol.ex b/lib/chromic_pdf/pdf/protocol.ex
@@ -19,8 +19,8 @@ defmodule ChromicPDF.Protocol do
   @type step :: call_step() | await_step() | output_step()
 
   # A protocol knows three types of steps: calls, awaits, and output.
-  # * The call step is a protocol call to send to the browser. Multiple call steps in sequence
-  #   are executed sequentially until the next await step is found.
+  # * The call step transforms the state and produces a protocol call to send to the browser.
+  #   Multiple call steps in sequence are executed sequentially until the next await step is found.
   # * Await steps are steps that try to match on messages received from the browser. When a
   #   message is matched, the await step can be removed from the queue (depending on the second
   #   element of the return tuple, `:keep | :remove`). Multiple await steps in sequence are

diff --git a/lib/chromic_pdf/supervisor.ex b/lib/chromic_pdf/supervisor.ex
@@ -146,17 +146,17 @@ defmodule ChromicPDF.Supervisor do
       @type source_and_options :: %{source: source_tuple(), opts: [pdf_option()]}
       @type source :: source() | source_and_options()
 
+      @type protocol_option :: {any(), any()}
+
       @type output_function_result :: any()
-      @type output_function :: (binary() -> output_function_result())
+      @type output_function :: (any() -> output_function_result())
       @type output_option :: {:output, binary()} | {:output, output_function()}
 
-      @type telemetry_metadata_option :: {:telemetry_metadata, map()}
-
-      @type export_option ::
+      @type shared_option ::
               output_option()
-              | telemetry_metadata_option()
+              | {:telemetry_metadata, map()}
 
-      @type export_return :: :ok | {:ok, binary()} | {:ok, output_function_result()}
+      @type result :: :ok | {:ok, any()} | {:ok, output_function_result()}
 
       @type info_option ::
               {:info,
@@ -187,6 +187,7 @@ defmodule ChromicPDF.Supervisor do
 
       @type pdf_option ::
               {:print_to_pdf, map()}
+              | {:protocol, module()}
               | navigate_option()
 
       @type pdfa_option ::
@@ -199,6 +200,7 @@ defmodule ChromicPDF.Supervisor do
       @type capture_screenshot_option ::
               {:capture_screenshot, map()}
               | {:full_page, boolean()}
+              | {:protocol, module()}
               | navigate_option()
 
       @type session_option ::
@@ -591,9 +593,9 @@ defmodule ChromicPDF.Supervisor do
 
           ChromicPDF.print_to_pdf({:url, "http:///example.net"}, wait_for: wait_for)
       '''
-      @spec print_to_pdf(source() | [source()]) :: export_return()
-      @spec print_to_pdf(source() | [source()], [pdf_option() | export_option()]) ::
-              export_return()
+      @spec print_to_pdf(source() | [source()]) :: result()
+      @spec print_to_pdf(source() | [source()], [pdf_option() | shared_option()]) ::
+              result()
       def print_to_pdf(source, opts \\ []) do
         with_services(&API.print_to_pdf(&1, source, opts))
       end
@@ -634,13 +636,20 @@ defmodule ChromicPDF.Supervisor do
             full_page: true
           )
       """
-      @spec capture_screenshot(source()) :: export_return()
-      @spec capture_screenshot(source(), [capture_screenshot_option() | export_option()]) ::
-              export_return()
+      @spec capture_screenshot(source()) :: result()
+      @spec capture_screenshot(source(), [capture_screenshot_option() | shared_option()]) ::
+              result()
       def capture_screenshot(source, opts \\ []) do
         with_services(&API.capture_screenshot(&1, source, opts))
       end
 
+      @doc false
+      @spec run_protocol(module()) :: result()
+      @spec run_protocol(module(), [shared_option() | protocol_option()]) :: result()
+      def run_protocol(protocol, opts \\ []) do
+        with_services(&API.run_protocol(&1, protocol, opts))
+      end
+
       @doc """
       Converts a PDF to PDF/A (either PDF/A-2b or PDF/A-3b).
 
@@ -724,8 +733,8 @@ defmodule ChromicPDF.Supervisor do
       ("Tags") and hence disables accessibility features of assistive technologies. See
       [On Accessibility / PDF/UA](#module-on-accessibility-pdf-ua) section for details.
       """
-      @spec convert_to_pdfa(path()) :: export_return()
-      @spec convert_to_pdfa(path(), [pdfa_option()]) :: export_return()
+      @spec convert_to_pdfa(path()) :: result()
+      @spec convert_to_pdfa(path(), [pdfa_option()]) :: result()
       def convert_to_pdfa(pdf_path, opts \\ []) do
         with_services(&API.convert_to_pdfa(&1, pdf_path, opts))
       end
@@ -739,9 +748,9 @@ defmodule ChromicPDF.Supervisor do
 
           ChromicPDF.print_to_pdfa({:url, "https://example.net"})
       """
-      @spec print_to_pdfa(source() | [source()]) :: export_return()
-      @spec print_to_pdfa(source() | [source()], [pdf_option() | pdfa_option() | export_option()]) ::
-              export_return()
+      @spec print_to_pdfa(source() | [source()]) :: result()
+      @spec print_to_pdfa(source() | [source()], [pdf_option() | pdfa_option() | shared_option()]) ::
+              result()
       def print_to_pdfa(source, opts \\ []) do
         with_services(&API.print_to_pdfa(&1, source, opts))
       end

diff --git a/test/integration/custom_protocol_test.exs b/test/integration/custom_protocol_test.exs
@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+
+defmodule ChromicPDF.CustomProtocolTest do
+  use ChromicPDF.Case, async: false
+  import ChromicPDF.TestAPI
+
+  defmodule BypassCSP do
+    import ChromicPDF.ProtocolMacros
+
+    steps do
+      call(:set_bypass_csp, "Page.setBypassCSP", [], %{"enabled" => true})
+      await_response(:bypass_csp_set, [])
+
+      include_protocol(ChromicPDF.PrintToPDF)
+    end
+  end
+
+  defmodule FixedScreenMetrics do
+    import ChromicPDF.ProtocolMacros
+
+    steps do
+      call(
+        :device_metrics,
+        "Emulation.setDeviceMetricsOverride",
+        [],
+        %{"width" => 200, "height" => 200, "mobile" => false, "deviceScaleFactor" => 1}
+      )
+
+      await_response(:device_metrics_response, [])
+
+      include_protocol(ChromicPDF.CaptureScreenshot)
+    end
+  end
+
+  defmodule GetUserAgent do
+    import ChromicPDF.ProtocolMacros
+
+    steps do
+      call(:get_version, "Browser.getVersion", [], %{})
+      await_response(:version, ["userAgent"])
+
+      output("userAgent")
+    end
+  end
+
+  setup do
+    start_supervised!(ChromicPDF)
+    :ok
+  end
+
+  describe ":protocol option to print_to_pdf/2" do
+    @html_with_csp """
+    <html>
+      <head>
+        <meta http-equiv="Content-Security-Policy" content="default-src 'none'">
+      </head>
+      <body>
+        <iframe src="data:text/html;charset=utf-8,%3Chtml%3E%3Cbody%3Efrom iframe%3C/body%3E%3C/html%3E">
+      </body>
+    </html>
+    """
+
+    @tag :pdftotext
+    test "allows to override the default protocol" do
+      print_to_pdf({:html, @html_with_csp}, fn text ->
+        refute String.contains?(text, "from iframe")
+      end)
+
+      print_to_pdf({:html, @html_with_csp}, [protocol: BypassCSP], fn text ->
+        assert String.contains?(text, "from iframe")
+      end)
+    end
+  end
+
+  describe ":protocol option to capture_screenshot/2" do
+    test "allows to set a custom protocol for capture_screenshot/2" do
+      assert {_, 200, 200} = capture_screenshot_and_identify(protocol: FixedScreenMetrics)
+    end
+  end
+
+  describe "run_protocol/2" do
+    test "allows to run custom protocols and get their output" do
+      assert {:ok, user_agent} = ChromicPDF.run_protocol(GetUserAgent)
+      assert user_agent =~ ~r/chrom/i
+    end
+  end
+end