From bd8a574fd825d2c0c203613aa0935d983df664dd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Valim?= <jose.valim@dashbit.co>
Date: Tue, 14 May 2024 09:00:51 +0200
Subject: [PATCH] Support MFArgs in rewrite_on

---
 lib/plug/rewrite_on.ex        |  8 ++++++++
 test/plug/rewrite_on_test.exs | 10 ++++++++++
 2 files changed, 18 insertions(+)

diff --git a/lib/plug/rewrite_on.ex b/lib/plug/rewrite_on.ex
index 2521c4e1..134bb383 100644
--- a/lib/plug/rewrite_on.ex
+++ b/lib/plug/rewrite_on.ex
@@ -14,6 +14,9 @@ defmodule Plug.RewriteOn do
     * `:x_forwarded_port` - to override the port based on on the "x-forwarded-port" header
     * `:x_forwarded_proto` - to override the protocol based on on the "x-forwarded-proto" header
 
+  A tuple representing a Module-Function-Args can also be given as argument
+  instead of a list.
+
   Since rewriting the scheme based on `x-forwarded-*` headers can open up
   security vulnerabilities, only use this plug if:
 
@@ -26,6 +29,7 @@ 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(header), do: List.wrap(header)
 
   @impl true
@@ -55,6 +59,10 @@ defmodule Plug.RewriteOn do
     conn
   end
 
+  def call(conn, {mod, fun, args}) do
+    call(conn, apply(mod, fun, args))
+  end
+
   defp put_scheme(%{scheme: :http, port: 80} = conn, ["https"]),
     do: %{conn | scheme: :https, port: 443}
 
diff --git a/test/plug/rewrite_on_test.exs b/test/plug/rewrite_on_test.exs
index 0230212c..766f5c68 100644
--- a/test/plug/rewrite_on_test.exs
+++ b/test/plug/rewrite_on_test.exs
@@ -6,6 +6,16 @@ defmodule Plug.RewriteOnTest do
     Plug.RewriteOn.call(conn, Plug.RewriteOn.init(rewrite))
   end
 
+  test "rewrites http to https based on MFArgs" do
+    conn =
+      conn(:get, "http://example.com/")
+      |> put_req_header("x-forwarded-proto", "https")
+      |> call({List, :flatten, [[:x_forwarded_proto]]})
+
+    assert conn.scheme == :https
+    assert conn.port == 443
+  end
+
   test "rewrites http to https based on x-forwarded-proto" do
     conn =
       conn(:get, "http://example.com/")