[3] add a dns_module option that allows overriding the DNS resolver (useful for testing) (#6)

Author: bhan-slab

README.md

@@ -90,6 +90,7 @@ following options:
 
   * `:schemes` - List of allowed URL schemes. Defaults to `["http, "https"]`.
 
+  * `:dns_module` - Any module that implements DNSResolver. Defaults to DNS from the `dns` package.
 
 These options can be passed to the function directly or set globally in your `config.exs`
 file:

lib/dns_resolver.ex

@@ -0,0 +1,3 @@
+defmodule DNSResolver do
+  @callback resolve(String.t) :: {atom, list}
+end

lib/safeurl.ex

@@ -133,7 +133,7 @@ defmodule SafeURL do
   def allowed?(url, opts \\ []) do
     uri = URI.parse(url)
     opts = build_options(opts)
-    address = resolve_address(uri.host)
+    address = resolve_address(uri.host, opts.dns_module)
 
     cond do
       uri.scheme not in opts.schemes ->
@@ -230,6 +230,7 @@ defmodule SafeURL do
     schemes = get_option(opts, :schemes)
     allowlist = get_option(opts, :allowlist)
     blocklist = get_option(opts, :blocklist)
+    dns_module = get_option(opts, :dns_module)
 
     blocklist =
       if get_option(opts, :block_reserved) do
@@ -238,7 +239,7 @@ defmodule SafeURL do
         blocklist
       end
 
-    %{schemes: schemes, allowlist: allowlist, blocklist: blocklist}
+    %{schemes: schemes, allowlist: allowlist, blocklist: blocklist, dns_module: dns_module}
   end
 
 
@@ -254,7 +255,7 @@ defmodule SafeURL do
 
 
   # Resolve hostname in DNS to an IP address (if not already an IP)
-  defp resolve_address(hostname) do
+  defp resolve_address(hostname, dns_module) do
     hostname
     |> to_charlist()
     |> :inet.parse_address()
@@ -264,7 +265,7 @@ defmodule SafeURL do
 
       {:error, :einval} ->
         # TODO: safely handle multiple IPs/round-robin DNS
-        case DNS.resolve(hostname) do
+        case dns_module.resolve(hostname) do
           {:ok, ips} -> List.first(ips)
           {:error, _reason} -> nil
         end

mix.exs

@@ -49,6 +49,7 @@ defmodule SafeURL.MixProject do
       block_reserved: true,
       blocklist: [],
       allowlist: [],
+      dns_module: DNS,
     ]
   end
 

test/safeurl_test.exs

@@ -1,3 +1,10 @@
+defmodule TestDNSResolver do
+  @behaviour DNSResolver
+
+  @impl DNSResolver
+  def resolve(_domain), do: {:ok, [{192, 0, 78, 24}]}
+end
+
 defmodule SafeURLTest do
   use ExUnit.Case
 
@@ -10,12 +17,12 @@ defmodule SafeURLTest do
 
   describe "#allowed?" do
     test "returns true for only allowed schemes" do
-      assert SafeURL.allowed?("http://includesecurity.com")
-      assert SafeURL.allowed?("https://includesecurity.com")
-      refute SafeURL.allowed?("ftp://includesecurity.com")
+      assert SafeURL.allowed?("http://includesecurity.com", dns_module: TestDNSResolver)
+      assert SafeURL.allowed?("https://includesecurity.com", dns_module: TestDNSResolver)
+      refute SafeURL.allowed?("ftp://includesecurity.com", dns_module: TestDNSResolver)
 
-      assert SafeURL.allowed?("ftp://includesecurity.com", schemes: ~w[ftp])
-      refute SafeURL.allowed?("http://includesecurity.com", schemes: ~w[ftp])
+      assert SafeURL.allowed?("ftp://includesecurity.com", schemes: ~w[ftp], dns_module: TestDNSResolver)
+      refute SafeURL.allowed?("http://includesecurity.com", schemes: ~w[ftp], dns_module: TestDNSResolver)
     end
 
     test "returns false for reserved ranges" do
@@ -37,7 +44,7 @@ defmodule SafeURLTest do
     end
 
     test "blocking custom IP ranges" do
-      opts = [blocklist: ["5.5.0.0/16", "100.0.0.0/24"]]
+      opts = [blocklist: ["5.5.0.0/16", "100.0.0.0/24"], dns_module: TestDNSResolver]
 
       assert SafeURL.allowed?("http://includesecurity.com", opts)
       assert SafeURL.allowed?("http://3.3.3.3", opts)
@@ -46,7 +53,7 @@ defmodule SafeURLTest do
     end
 
     test "only allows IPs in the allowlist when present" do
-      opts = [allowlist: ["10.0.0.0/24"]]
+      opts = [allowlist: ["10.0.0.0/24"], dns_module: TestDNSResolver]
 
       assert SafeURL.allowed?("http://10.0.0.1/", opts)
       refute SafeURL.allowed?("http://72.254.45.178", opts)