1 files changed, 20 insertions, 9 deletions
diff --git a/lib/matrix_app_service_web/auth_plug.ex b/lib/matrix_app_service_web/auth_plug.ex
index 19f8c21..be38fb0 100644
--- a/lib/matrix_app_service_web/auth_plug.ex
+++ b/lib/matrix_app_service_web/auth_plug.ex
@@ -11,21 +11,18 @@ defmodule MatrixAppServiceWeb.AuthPlug do
 
   @doc false
   @impl Plug
-  def init(opts) do
-    opts
+  def init(homeserver_token) do
+    homeserver_token
   end
 
   @doc false
   @impl Plug
-  def call(%Plug.Conn{params: %{"access_token" => access_token}} = conn, homeserver_token)
-      when access_token == homeserver_token do
-    conn
+  def call(%Plug.Conn{params: %{"access_token" => access_token}} = conn, homeserver_token) when is_binary(homeserver_token) do
+    verify_access_token(conn, homeserver_token)
   end
 
-  def call(%Plug.Conn{params: %{"access_token" => _access_token}} = conn, _homeserver_token) do
-    Logger.warn("Received invalid homeserver token")
-
-    respond_error(conn, 403)
+  def call(%Plug.Conn{params: %{"access_token" => access_token}} = conn, homeserver_token) when is_function(homeserver_token, 0) do
+    verify_access_token(conn, homeserver_token.())
   end
 
   def call(conn, _opts) do
@@ -34,6 +31,20 @@ defmodule MatrixAppServiceWeb.AuthPlug do
     respond_error(conn, 401)
   end
 
+  defp verify_access_token(
+         %Plug.Conn{params: %{"access_token" => access_token}} = conn,
+         homeserver_token
+       )
+       when access_token == homeserver_token do
+    conn
+  end
+
+  defp verify_access_token(conn, _homeserver_token) do
+    Logger.warn("Received invalid homeserver token")
+
+    respond_error(conn, 403)
+  end
+
   defp respond_error(conn, error_code) do
     conn
     |> put_status(error_code)