Support http proxies

lib/bumblebee/application.ex

@@ -0,0 +1,19 @@
+defmodule Bumblebee.Application do
+  @moduledoc false
+
+  use Application
+
+  @impl true
+  def start(_type, _args) do
+    Bumblebee.Utils.HTTP.start_inets_profile()
+
+    children = []
+    opts = [strategy: :one_for_one, name: Bumblebee.Supervisor]
+    Supervisor.start_link(children, opts)
+  end
+
+  @impl true
+  def stop(_state) do
+    Bumblebee.Utils.HTTP.stop_inets_profile()
+  end
+end

lib/bumblebee/utils/http.ex

@@ -42,13 +42,13 @@ defmodule Bumblebee.Utils.HTTP do
             if Process.alive?(caller) do
               send(caller, {:http, reply_info})
             else
-              :httpc.cancel_request(request_id)
+              :httpc.cancel_request(request_id, :bumblebee)
             end
           end
 
           opts = [stream: :self, sync: false, receiver: receiver]
 
-          {:ok, request_id} = :httpc.request(:get, request, http_opts, opts)
+          {:ok, request_id} = :httpc.request(:get, request, http_opts, opts, :bumblebee)
           download_loop(%{request_id: request_id, file: file, total_size: nil, size: nil})
         after
           File.close(file)
@@ -103,7 +103,7 @@ defmodule Bumblebee.Utils.HTTP do
         download_loop(state)
 
       {:error, error} ->
-        :httpc.cancel_request(state.request_id)
+        :httpc.cancel_request(state.request_id, :bumblebee)
         {:error, "failed to write to file, reason: #{:file.format_error(error)}"}
     end
   end
@@ -158,7 +158,7 @@ defmodule Bumblebee.Utils.HTTP do
       body_format: :binary
     ]
 
-    case :httpc.request(method, request, http_opts, opts) do
+    case :httpc.request(method, request, http_opts, opts, :bumblebee) do
       {:ok, {{_, status, _}, headers, body}} ->
         {:ok, %{status: status, headers: parse_headers(headers), body: body}}
 
@@ -183,13 +183,59 @@ defmodule Bumblebee.Utils.HTTP do
   end
 
   defp http_ssl_opts() do
+    # Allow a user-specified CA certs to support, for example, HTTPS proxies
+    cacert_opt =
+      case System.get_env("BUMBLEBEE_CACERTS_PATH") do
+        nil -> {:cacerts, :public_key.cacerts_get()}
+        file -> {:cacertfile, file}
+      end
+
     # Use secure options, see https://gist.github.com/jonatanklosko/5e20ca84127f6b31bbe3906498e1a1d7
     [
+      cacert_opt,
       verify: :verify_peer,
-      cacertfile: CAStore.file_path(),
       customize_hostname_check: [
         match_fun: :public_key.pkix_verify_hostname_match_fun(:https)
       ]
     ]
   end
+
+  @doc false
+  def start_inets_profile() do
+    # Starting an HTTP client profile allows us to scope the httpc
+    # configuration options, such as proxy options
+    {:ok, _pid} = :inets.start(:httpc, profile: :bumblebee)
+    set_proxy_options()
+  end
+
+  @doc false
+  def stop_inets_profile() do
+    :inets.stop(:httpc, :bumblebee)
+  end
+
+  defp set_proxy_options() do
+    http_proxy = System.get_env("HTTP_PROXY") || System.get_env("http_proxy")
+    https_proxy = System.get_env("HTTPS_PROXY") || System.get_env("https_proxy")
+
+    no_proxy =
+      if no_proxy = System.get_env("NO_PROXY") || System.get_env("no_proxy") do
+        no_proxy
+        |> String.split(",")
+        |> Enum.map(&String.to_charlist/1)
+      else
+        []
+      end
+
+    set_proxy_option(:proxy, http_proxy, no_proxy)
+    set_proxy_option(:https_proxy, https_proxy, no_proxy)
+  end
+
+  defp set_proxy_option(proxy_scheme, proxy, no_proxy) do
+    uri = URI.parse(proxy || "")
+
+    if uri.host && uri.port do
+      host = String.to_charlist(uri.host)
+      :httpc.set_options([{proxy_scheme, {{host, uri.port}, no_proxy}}], :bumblebee)
+    end
+  end
 end

mix.exs

@@ -21,6 +21,7 @@ defmodule Bumblebee.MixProject do
 
   def application do
     [
+      mod: {Bumblebee.Application, []},
       extra_applications: [:logger, :inets, :ssl]
     ]
   end
@@ -42,7 +43,6 @@ defmodule Bumblebee.MixProject do
       {:nx_image, "~> 0.1.0"},
       {:unpickler, "~> 0.1.0"},
       {:safetensors, "~> 0.1.3"},
-      {:castore, "~> 0.1 or ~> 1.0"},
       {:jason, "~> 1.4.0"},
       {:unzip, "~> 0.10.0"},
       {:progress_bar, "~> 3.0"},

mix.lock

@@ -1,7 +1,7 @@
 %{
   "axon": {:git, "https://github.com/elixir-nx/axon.git", "054eb4c1c224582528e8d1603ad08e7c4088f21c", []},
   "bypass": {:hex, :bypass, "2.1.0", "909782781bf8e20ee86a9cabde36b259d44af8b9f38756173e8f5e2e1fabb9b1", [:mix], [{:plug, "~> 1.7", [hex: :plug, repo: "hexpm", optional: false]}, {:plug_cowboy, "~> 2.0", [hex: :plug_cowboy, repo: "hexpm", optional: false]}, {:ranch, "~> 1.3", [hex: :ranch, repo: "hexpm", optional: false]}], "hexpm", "d9b5df8fa5b7a6efa08384e9bbecfe4ce61c77d28a4282f79e02f1ef78d96b80"},
-  "castore": {:hex, :castore, "1.0.4", "ff4d0fb2e6411c0479b1d965a814ea6d00e51eb2f58697446e9c41a97d940b28", [:mix], [], "hexpm", "9418c1b8144e11656f0be99943db4caf04612e3eaecefb5dae9a2a87565584f8"},
+  "castore": {:hex, :castore, "1.0.8", "dedcf20ea746694647f883590b82d9e96014057aff1d44d03ec90f36a5c0dc6e", [:mix], [], "hexpm", "0b2b66d2ee742cb1d9cb8c8be3b43c3a70ee8651f37b75a8b982e036752983f1"},
   "cc_precompiler": {:hex, :cc_precompiler, "0.1.8", "933a5f4da3b19ee56539a076076ce4d7716d64efc8db46fd066996a7e46e2bfd", [:mix], [{:elixir_make, "~> 0.7.3", [hex: :elixir_make, repo: "hexpm", optional: false]}], "hexpm", "176bdf4366956e456bf761b54ad70bc4103d0269ca9558fd7cee93d1b3f116db"},
   "complex": {:hex, :complex, "0.5.0", "af2d2331ff6170b61bb738695e481b27a66780e18763e066ee2cd863d0b1dd92", [:mix], [], "hexpm", "2683bd3c184466cfb94fad74cbfddfaa94b860e27ad4ca1bffe3bff169d91ef1"},
   "cowboy": {:hex, :cowboy, "2.9.0", "865dd8b6607e14cf03282e10e934023a1bd8be6f6bacf921a7e2a96d800cd452", [:make, :rebar3], [{:cowlib, "2.11.0", [hex: :cowlib, repo: "hexpm", optional: false]}, {:ranch, "1.8.0", [hex: :ranch, repo: "hexpm", optional: false]}], "hexpm", "2c729f934b4e1aa149aff882f57c6372c15399a20d54f65c8d67bef583021bde"},