Add support for Plug.RewriteOn to rewrite nonstandard headers

lib/plug/rewrite_on.ex

@@ -10,11 +10,18 @@ defmodule Plug.RewriteOn do
 
   The supported values are:
 
-    * `:x_forwarded_for` - to override the remote ip based on the "x-forwarded-for" header
+    * `:x_forwarded_for` - to override the remote IP based on the "x-forwarded-for" header
     * `:x_forwarded_host` - to override the host based on the "x-forwarded-host" header
     * `:x_forwarded_port` - to override the port based on the "x-forwarded-port" header
     * `:x_forwarded_proto` - to override the protocol based on the "x-forwarded-proto" header
 
+  Some HTTPS proxies use nonstandard headers, which can be specified in the list via tuples:
+
+    * `{:remote_ip, header}` - to override the remote IP based on a custom header
+    * `{:host, header}` - to override the host based on a custom header
+    * `{:port, header}` - to override the port based on a custom header
+    * `{:scheme, header}` - to override the protocol based on a custom header
+
   A tuple representing a Module-Function-Args can also be given as argument
   instead of a list.
 
@@ -30,31 +37,47 @@ defmodule Plug.RewriteOn do
   import Plug.Conn, only: [get_req_header: 2]
 
   @impl true
-  def init(header) when is_tuple(header), do: header
+  def init({_m, _f, _a} = header), do: header
   def init(header), do: List.wrap(header)
 
   @impl true
   def call(conn, [:x_forwarded_for | rewrite_on]) do
+    call(conn, [{:remote_ip, "x-forwarded-for"} | rewrite_on])
+  end
+
+  def call(conn, [:x_forwarded_proto | rewrite_on]) do
+    call(conn, [{:scheme, "x-forwarded-proto"} | rewrite_on])
+  end
+
+  def call(conn, [:x_forwarded_port | rewrite_on]) do
+    call(conn, [{:port, "x-forwarded-port"} | rewrite_on])
+  end
+
+  def call(conn, [:x_forwarded_host | rewrite_on]) do
+    call(conn, [{:host, "x-forwarded-host"} | rewrite_on])
+  end
+
+  def call(conn, [{:remote_ip, header} | rewrite_on]) do
     conn
-    |> put_remote_ip(get_req_header(conn, "x-forwarded-for"))
+    |> put_remote_ip(get_req_header(conn, header))
     |> call(rewrite_on)
   end
 
-  def call(conn, [:x_forwarded_proto | rewrite_on]) do
+  def call(conn, [{:scheme, header} | rewrite_on]) do
     conn
-    |> put_scheme(get_req_header(conn, "x-forwarded-proto"))
+    |> put_scheme(get_req_header(conn, header))
     |> call(rewrite_on)
   end
 
-  def call(conn, [:x_forwarded_port | rewrite_on]) do
+  def call(conn, [{:port, header} | rewrite_on]) do
     conn
-    |> put_port(get_req_header(conn, "x-forwarded-port"))
+    |> put_port(get_req_header(conn, header))
     |> call(rewrite_on)
   end
 
-  def call(conn, [:x_forwarded_host | rewrite_on]) do
+  def call(conn, [{:host, header} | rewrite_on]) do
     conn
-    |> put_host(get_req_header(conn, "x-forwarded-host"))
+    |> put_host(get_req_header(conn, header))
     |> call(rewrite_on)
   end
 

test/plug/rewrite_on_test.exs

@@ -27,6 +27,16 @@ defmodule Plug.RewriteOnTest do
     assert conn.port == 443
   end
 
+  test "rewrites http to https based on a custom header" do
+    conn =
+      conn(:get, "http://example.com/")
+      |> put_req_header("custom-forwarded-proto", "https")
+      |> call({:scheme, "custom-forwarded-proto"})
+
+    assert conn.scheme == :https
+    assert conn.port == 443
+  end
+
   test "doesn't change the port when it doesn't match the scheme" do
     conn =
       conn(:get, "http://example.com:1234/")
@@ -46,6 +56,15 @@ defmodule Plug.RewriteOnTest do
     assert conn.host == "truessl.example.com"
   end
 
+  test "rewrites host with a custom header" do
+    conn =
+      conn(:get, "http://example.com/")
+      |> put_req_header("custom-forwarded-host", "truessl.example.com")
+      |> call({:host, "custom-forwarded-host"})
+
+    assert conn.host == "truessl.example.com"
+  end
+
   test "rewrites port with a x-forwarded-port header" do
     conn =
       conn(:get, "http://example.com/")
@@ -55,6 +74,15 @@ defmodule Plug.RewriteOnTest do
     assert conn.port == 3030
   end
 
+  test "rewrites port with a custom header" do
+    conn =
+      conn(:get, "http://example.com/")
+      |> put_req_header("custom-forwarded-port", "3030")
+      |> call({:port, "custom-forwarded-port"})
+
+    assert conn.port == 3030
+  end
+
   test "rewrites remote_ip with a x-forwarded-for header" do
     conn =
       conn(:get, "http://example.com/")
@@ -85,6 +113,38 @@ defmodule Plug.RewriteOnTest do
     assert conn.remote_ip == {0, 0, 0, 0, 0, 0, 0, 1}
   end
 
+
+  test "rewrites remote_ip with a custom header" do
+    conn =
+      conn(:get, "http://example.com/")
+      |> put_req_header("custom-forwarded-for", "bad")
+      |> call({:remote_ip, "custom-forwarded-for"})
+
+    assert conn.remote_ip == {127, 0, 0, 1}
+
+    conn =
+      conn(:get, "http://example.com/")
+      |> put_req_header("custom-forwarded-for", "4.3.2.1")
+      |> call({:remote_ip, "custom-forwarded-for"})
+
+    assert conn.remote_ip == {4, 3, 2, 1}
+
+    conn =
+      conn(:get, "http://example.com/")
+      |> put_req_header("custom-forwarded-for", "1.2.3.4,::1")
+      |> call({:remote_ip, "custom-forwarded-for"})
+
+    assert conn.remote_ip == {1, 2, 3, 4}
+
+    conn =
+      conn(:get, "http://example.com/")
+      |> put_req_header("custom-forwarded-for", "::1,1.2.3.4")
+      |> call({:remote_ip, "custom-forwarded-for"})
+
+    assert conn.remote_ip == {0, 0, 0, 0, 0, 0, 0, 1}
+  end
+
+
   test "rewrites the host, the port, and the protocol" do
     conn =
       conn(:get, "http://example.com/")
@@ -98,9 +158,28 @@ defmodule Plug.RewriteOnTest do
     assert conn.scheme == :https
   end
 
+
+  test "rewrites the host, the port, and the protocol with custom headers" do
+    conn =
+      conn(:get, "http://example.com/")
+      |> put_req_header("custom-forwarded-host", "truessl.example.com")
+      |> put_req_header("custom-forwarded-port", "3030")
+      |> put_req_header("custom-forwarded-proto", "https")
+      |> call([host: "custom-forwarded-host",
+               port: "custom-forwarded-port",
+               scheme: "custom-forwarded-proto"])
+
+    assert conn.host == "truessl.example.com"
+    assert conn.port == 3030
+    assert conn.scheme == :https
+  end
+
   test "raises when receiving an unknown rewrite" do
     assert_raise RuntimeError, "unknown rewrite: :x_forwarded_other", fn ->
       call(conn(:get, "http://example.com/"), :x_forwarded_other)
     end
+    assert_raise RuntimeError, "unknown rewrite: {:other, \"value\"}", fn ->
+      call(conn(:get, "http://example.com/"), [other: "value"])
+    end
   end
 end