Add support for headers and dynamic base URL (philss#65)

The :base_url attribute now accepts static headers or a custom function
that returns a URL and headers. This allows us to fetch NIFs from more
complicated sources, rather than just public GitHub releases.

Co-authored-by: Philip Sampaio <[email protected]>

lib/rustler_precompiled.ex

@@ -24,7 +24,20 @@ defmodule RustlerPrecompiled do
 
     * `:crate` - The name of Rust crate if different from the `:otp_app`. This is optional.
 
-    * `:base_url` - A valid URL that is used as base path for the NIF file.
+    * `:base_url` - Location where to find the NIFs from. This should be one of the following:
+
+      * A URL to a directory containing the NIFs. The name of the NIF will be appended to it
+        and a GET request will be made. Works well with public GitHub releases.
+
+      * A tuple of `{URL, headers}`. The headers should be a list of key-value pairs.
+        This is useful when the NIFs are hosted in a private server.
+
+      * A tuple of `{module, function}` where the `function` is an atom representing the function
+        name in that module. It's expected a function of arity 1, where the NIF file name is given,
+        and it should return a URL or a tuple of `{URL, headers}`.
+        This should be used for all cases not covered by the above.
+        For example when multiple requests have to be made, like when using a private GitHub release
+        through the GitHub API, or when the URLs don't resemble a simple directory.
 
     * `:version` - The version of precompiled assets (it is part of the NIF filename).
 
@@ -262,13 +275,19 @@ defmodule RustlerPrecompiled do
 
   @native_dir "priv/native"
 
+  @doc deprecated: "Use available_nifs/1 instead"
+  def available_nif_urls(nif_module) when is_atom(nif_module) do
+    available_nifs(nif_module)
+    |> Enum.map(fn {_lib_name, {url, _headers}} -> url end)
+  end
+
   @doc """
-  Returns URLs for NIFs based on its module name.
+  Returns URLs for NIFs based on its module name as a list of tuples: `[{lib_name, {url, headers}}]`.
 
   The module name is the one that defined the NIF and this information
   is stored in a metadata file.
   """
-  def available_nif_urls(nif_module) when is_atom(nif_module) do
+  def available_nifs(nif_module) when is_atom(nif_module) do
     nif_module
     |> metadata_file()
     |> read_map_from_file()
@@ -286,6 +305,17 @@ defmodule RustlerPrecompiled do
 
   @doc false
   def nif_urls_from_metadata(metadata) when is_map(metadata) do
+    case(nifs_from_metadata(metadata)) do
+      {:ok, nifs} ->
+        {:ok, Enum.map(nifs, fn {_lib_name, {url, _headers}} -> url end)}
+
+      {:error, wrong_meta} ->
+        {:error, wrong_meta}
+    end
+  end
+
+  @doc false
+  def nifs_from_metadata(metadata) when is_map(metadata) do
     case metadata do
       %{
         targets: targets,
@@ -320,22 +350,26 @@ defmodule RustlerPrecompiled do
     variants = Map.fetch!(variants, target_triple)
 
     for variant <- variants do
-      tar_gz_file_url(
-        base_url,
-        lib_name_with_ext(target_triple, lib_name <> "--" <> Atom.to_string(variant))
-      )
+      lib_name = lib_name_with_ext(target_triple, lib_name <> "--" <> Atom.to_string(variant))
+      {lib_name, tar_gz_file_url(base_url, lib_name)}
     end
   end
 
   defp maybe_variants_tar_gz_urls(_, _, _, _), do: []
 
+  @doc deprecated: "Use current_target_nifs/1 instead"
+  def current_target_nif_urls(nif_module) when is_atom(nif_module) do
+    current_target_nifs(nif_module)
+    |> Enum.map(fn {_lib_name, {url, _headers}} -> url end)
+  end
+
   @doc """
-  Returns the file URLs to be downloaded for current target.
+  Returns the file URLs to be downloaded for current target as a list of tuples: `[{lib_name, {url, headers}}]`.
 
   It is in the plural because a target may have some variants for it.
   It receives the NIF module.
   """
-  def current_target_nif_urls(nif_module) when is_atom(nif_module) do
+  def current_target_nifs(nif_module) when is_atom(nif_module) do
     metadata =
       nif_module
       |> metadata_file()
@@ -362,9 +396,10 @@ defmodule RustlerPrecompiled do
 
   defp tar_gz_urls(base_url, basename, version, nif_version, target_triple, variants) do
     lib_name = lib_name(basename, version, nif_version, target_triple)
+    lib_name_with_ext = lib_name_with_ext(target_triple, lib_name)
 
     [
-      tar_gz_file_url(base_url, lib_name_with_ext(target_triple, lib_name))
+      {lib_name_with_ext, tar_gz_file_url(base_url, lib_name_with_ext(target_triple, lib_name))}
       | maybe_variants_tar_gz_urls(variants, base_url, target_triple, lib_name)
     ]
   end
@@ -615,7 +650,7 @@ defmodule RustlerPrecompiled do
 
         # `cache_base_dir` is a "private" option used only in tests.
         cache_dir = cache_dir(config.base_cache_dir, "precompiled_nifs")
-        cached_tar_gz = Path.join(cache_dir, "#{file_name}.tar.gz")
+        cached_tar_gz = Path.join(cache_dir, file_name)
 
         {:ok,
          Map.merge(basic_metadata, %{
@@ -840,21 +875,34 @@ defmodule RustlerPrecompiled do
         "so"
       end
 
-    "#{lib_name}.#{ext}"
+    "#{lib_name}.#{ext}.tar.gz"
   end
 
-  defp tar_gz_file_url(base_url, file_name) do
+  defp tar_gz_file_url({module, function_name}, file_name)
+       when is_atom(module) and is_atom(function_name) do
+    apply(module, function_name, [file_name])
+  end
+
+  defp tar_gz_file_url({base_url, request_headers}, file_name) do
     uri = URI.parse(base_url)
 
     uri =
       Map.update!(uri, :path, fn path ->
-        Path.join(path || "", "#{file_name}.tar.gz")
+        Path.join(path || "", file_name)
       end)
 
-    to_string(uri)
+    {to_string(uri), request_headers}
   end
 
-  defp download_nif_artifact(url) do
+  defp tar_gz_file_url(base_url, file_name) do
+    tar_gz_file_url({base_url, []}, file_name)
+  end
+
+  defp download_nif_artifact(url) when is_binary(url) do
+    download_nif_artifact({url, []})
+  end
+
+  defp download_nif_artifact({url, request_headers}) do
     url = String.to_charlist(url)
     Logger.debug("Downloading NIF from #{url}")
 
@@ -895,7 +943,10 @@ defmodule RustlerPrecompiled do
 
     options = [body_format: :binary]
 
-    case :httpc.request(:get, {url, []}, http_options, options) do
+    request_headers =
+      Enum.map(request_headers, fn {k, v} when is_binary(k) -> {String.to_charlist(k), v} end)
+
+    case :httpc.request(:get, {url, request_headers}, http_options, options) do
       {:ok, {{_, 200, _}, _headers, body}} ->
         {:ok, body}
 
@@ -912,16 +963,17 @@ defmodule RustlerPrecompiled do
     attempts = max_retries(options)
 
     download_results =
-      for url <- urls, do: {url, with_retry(fn -> download_nif_artifact(url) end, attempts)}
+      for {lib_name, url} <- urls,
+          do: {lib_name, with_retry(fn -> download_nif_artifact(url) end, attempts)}
 
     cache_dir = cache_dir("precompiled_nifs")
     :ok = File.mkdir_p(cache_dir)
 
     Enum.flat_map(download_results, fn result ->
-      with {:download, {url, download_result}} <- {:download, result},
+      with {:download, {lib_name, download_result}} <- {:download, result},
            {:download_result, {:ok, body}} <- {:download_result, download_result},
            hash <- :crypto.hash(@checksum_algo, body),
-           path <- Path.join(cache_dir, basename_from_url(url)),
+           path <- Path.join(cache_dir, lib_name),
            {:file, :ok} <- {:file, File.write(path, body)} do
         checksum = Base.encode16(hash, case: :lower)
 
@@ -931,7 +983,7 @@ defmodule RustlerPrecompiled do
 
         [
           %{
-            url: url,
+            lib_name: lib_name,
             path: path,
             checksum: checksum,
             checksum_algo: @checksum_algo
@@ -985,14 +1037,6 @@ defmodule RustlerPrecompiled do
     end)
   end
 
-  defp basename_from_url(url) do
-    uri = URI.parse(url)
-
-    uri.path
-    |> String.split("/")
-    |> List.last()
-  end
-
   defp read_map_from_file(file) do
     with {:ok, contents} <- File.read(file),
          {%{} = contents, _} <- Code.eval_string(contents) do

lib/rustler_precompiled/config.ex

@@ -83,16 +83,34 @@ defmodule RustlerPrecompiled.Config do
 
   defp validate_base_url!(nil), do: raise_for_nil_field_value(:base_url)
 
-  defp validate_base_url!(base_url) do
+  defp validate_base_url!(base_url) when is_binary(base_url) do
+    validate_base_url!({base_url, []})
+  end
+
+  defp validate_base_url!({base_url, headers}) when is_binary(base_url) and is_list(headers) do
     case :uri_string.parse(base_url) do
       %{} ->
-        base_url
+        if Enum.all?(headers, &match?({key, value} when is_list(key) and is_binary(value), &1)) do
+          {base_url, headers}
+        else
+          raise "`:base_url` for `RustlerPrecompiled` must be a list of `{charlist(),binary()}`"
+        end
 
       {:error, :invalid_uri, error} ->
         raise "`:base_url` for `RustlerPrecompiled` is invalid: #{inspect(to_string(error))}"
     end
   end
 
+  defp validate_base_url!({module, function}) when is_atom(module) and is_atom(function) do
+    Code.ensure_compiled!(module)
+
+    if Kernel.function_exported?(module, function, 1) do
+      {module, function}
+    else
+      raise "`:base_url` for `RustlerPrecompiled` is a function that does not exist: `#{inspect(module)}.#{function}/1`"
+    end
+  end
+
   defp validate_list!(nil, option, _valid_values), do: raise_for_nil_field_value(option)
 
   defp validate_list!([_ | _] = values, option, valid_values) do

test/rustler_precompiled_test.exs

@@ -500,6 +500,111 @@ defmodule RustlerPrecompiledTest do
       end)
     end
 
+    @tag :tmp_dir
+    test "a project downloading precompiled NIFs with custom header", %{
+      tmp_dir: tmp_dir,
+      checksum_sample: checksum_sample,
+      nif_fixtures_dir: nif_fixtures_dir
+    } do
+      bypass = Bypass.open()
+
+      in_tmp(tmp_dir, fn ->
+        File.write!("checksum-Elixir.RustlerPrecompilationExample.Native.exs", checksum_sample)
+
+        Bypass.expect_once(bypass, fn conn ->
+          file_name = List.last(conn.path_info)
+          file = File.read!(Path.join([nif_fixtures_dir, "precompiled_nifs", file_name]))
+
+          if Plug.Conn.get_req_header(conn, "authorization") == ["Token 123"] do
+            Plug.Conn.resp(conn, 200, file)
+          else
+            Plug.Conn.resp(conn, 401, "Unauthorized")
+          end
+        end)
+
+        result =
+          capture_log(fn ->
+            config = %RustlerPrecompiled.Config{
+              otp_app: :rustler_precompiled,
+              module: RustlerPrecompilationExample.Native,
+              base_cache_dir: tmp_dir,
+              base_url:
+                {"http://localhost:#{bypass.port}/download", [{"authorization", "Token 123"}]},
+              version: "0.2.0",
+              crate: "example",
+              targets: @available_targets,
+              nif_versions: @default_nif_versions
+            }
+
+            {:ok, metadata} = RustlerPrecompiled.build_metadata(config)
+
+            assert {:ok, result} = RustlerPrecompiled.download_or_reuse_nif_file(config, metadata)
+
+            assert result.load?
+            assert {:rustler_precompiled, path} = result.load_from
+
+            assert path =~ "priv/native"
+            assert path =~ "example-v0.2.0-nif"
+          end)
+
+        assert result =~ "Downloading"
+        assert result =~ "http://localhost:#{bypass.port}/download"
+        assert result =~ "NIF cached at"
+      end)
+    end
+
+    @tag :tmp_dir
+    test "a project downloading precompiled NIFs with custom handler", %{
+      tmp_dir: tmp_dir,
+      checksum_sample: checksum_sample,
+      nif_fixtures_dir: nif_fixtures_dir
+    } do
+      bypass = Bypass.open(port: 1234)
+
+      in_tmp(tmp_dir, fn ->
+        File.write!("checksum-Elixir.RustlerPrecompilationExample.Native.exs", checksum_sample)
+
+        Bypass.expect_once(bypass, fn conn ->
+          %{"file_name" => file_name} = URI.decode_query(conn.query_string)
+          file = File.read!(Path.join([nif_fixtures_dir, "precompiled_nifs", file_name]))
+
+          if Plug.Conn.get_req_header(conn, "authorization") == ["Token 123"] do
+            Plug.Conn.resp(conn, 200, file)
+          else
+            Plug.Conn.resp(conn, 401, "Unauthorized")
+          end
+        end)
+
+        result =
+          capture_log(fn ->
+            config = %RustlerPrecompiled.Config{
+              otp_app: :rustler_precompiled,
+              module: RustlerPrecompilationExample.Native,
+              base_cache_dir: tmp_dir,
+              base_url: {__MODULE__, :url_with_headers},
+              version: "0.2.0",
+              crate: "example",
+              targets: @available_targets,
+              nif_versions: @default_nif_versions
+            }
+
+            {:ok, metadata} = RustlerPrecompiled.build_metadata(config)
+
+            assert {:ok, result} = RustlerPrecompiled.download_or_reuse_nif_file(config, metadata)
+
+            assert result.load?
+            assert {:rustler_precompiled, path} = result.load_from
+
+            assert path =~ "priv/native"
+            assert path =~ "example-v0.2.0-nif"
+          end)
+
+        assert result =~ "Downloading"
+        assert result =~ "http://localhost:#{bypass.port}/download"
+        assert result =~ "NIF cached at"
+      end)
+    end
+
     @tag :tmp_dir
     test "a project downloading precompiled NIFs with retry", %{
       tmp_dir: tmp_dir,
@@ -793,7 +898,7 @@ defmodule RustlerPrecompiledTest do
       # We need this guard because not every one is running the tests in the same OS/Arch.
       if metadata.lib_name =~ "x86_64-unknown-linux-gnu" do
         assert String.ends_with?(metadata.lib_name, "--old_glibc")
-        assert String.ends_with?(metadata.file_name, "--old_glibc.so")
+        assert String.ends_with?(metadata.file_name, "--old_glibc.so.tar.gz")
       end
     end
 
@@ -821,7 +926,7 @@ defmodule RustlerPrecompiledTest do
 
       if metadata.lib_name =~ "x86_64-unknown-linux-gnu" do
         assert String.ends_with?(metadata.lib_name, "--legacy_cpus")
-        assert String.ends_with?(metadata.file_name, "--legacy_cpus.so")
+        assert String.ends_with?(metadata.file_name, "--legacy_cpus.so.tar.gz")
       end
     end
   end
@@ -955,4 +1060,9 @@ defmodule RustlerPrecompiledTest do
     |> Base.encode64()
     |> binary_part(0, len)
   end
+
+  def url_with_headers(file_name) do
+    {"http://localhost:1234/download?file_name=#{file_name}&foo=bar",
+     [{"authorization", "Token 123"}]}
+  end
 end