From 340c10d3e1d6a6e3bd30a9b700352b77edf649ae Mon Sep 17 00:00:00 2001
From: Austin Ziegler <austin@zieglers.ca>
Date: Fri, 6 Oct 2023 00:47:03 -0400
Subject: [PATCH] Add `assign_as` option for Plug.RequestId

- `:assign_as` - The name of the key that will be used to store the
  discovered or generated request id in `conn.private`. If not provided,
  the request id will not be stored.

  ```elixir
  plug Plug.RequestId, assign_as: :plug_request_id
  ```

Resolves: #1171
---
 lib/plug/request_id.ex        | 22 +++++++++++++++++-----
 test/plug/request_id_test.exs | 16 ++++++++++++++++
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/lib/plug/request_id.ex b/lib/plug/request_id.ex
index bc52cb15..24e6a4c2 100644
--- a/lib/plug/request_id.ex
+++ b/lib/plug/request_id.ex
@@ -34,6 +34,12 @@ defmodule Plug.RequestId do
 
           plug Plug.RequestId, http_header: "custom-request-id"
 
+    * `:assign_as` - The name of the key that will be used to store the
+      discovered or generated request id in `conn.assigns`. If not provided,
+      the request id will not be stored.
+
+          plug Plug.RequestId, assign_as: :plug_request_id
+
   """
 
   require Logger
@@ -42,14 +48,17 @@ defmodule Plug.RequestId do
 
   @impl true
   def init(opts) do
-    Keyword.get(opts, :http_header, "x-request-id")
+    {
+      Keyword.get(opts, :http_header, "x-request-id"),
+      Keyword.get(opts, :assign_as)
+    }
   end
 
   @impl true
-  def call(conn, req_id_header) do
+  def call(conn, {header, assign_as}) do
     conn
-    |> get_request_id(req_id_header)
-    |> set_request_id(req_id_header)
+    |> get_request_id(header)
+    |> set_request_id(header, assign_as)
   end
 
   defp get_request_id(conn, header) do
@@ -59,8 +68,11 @@ defmodule Plug.RequestId do
     end
   end
 
-  defp set_request_id({conn, request_id}, header) do
+  defp set_request_id({conn, request_id}, header, assign_as) do
     Logger.metadata(request_id: request_id)
+
+    conn = if assign_as, do: Conn.assign(conn, assign_as, request_id), else: conn
+
     Conn.put_resp_header(conn, header, request_id)
   end
 
diff --git a/test/plug/request_id_test.exs b/test/plug/request_id_test.exs
index 54cc9510..d7cc9700 100644
--- a/test/plug/request_id_test.exs
+++ b/test/plug/request_id_test.exs
@@ -65,6 +65,22 @@ defmodule Plug.RequestIdTest do
     assert res_request_id == meta_request_id
   end
 
+  test "assigns the request id to conn.assigns when given an assignment key" do
+    request_id = "existingidthatislongenough"
+
+    conn =
+      conn(:get, "/")
+      |> put_req_header("x-request-id", request_id)
+      |> call(assign_as: :plug_request_id)
+
+    [res_request_id] = get_resp_header(conn, "x-request-id")
+    meta_request_id = Logger.metadata()[:request_id]
+    assigned_request_id = conn.assigns.plug_request_id
+    assert assigned_request_id == request_id
+    assert res_request_id == assigned_request_id
+    assert res_request_id == meta_request_id
+  end
+
   defp generated_request_id?(request_id) do
     Regex.match?(~r/\A[A-Za-z0-9-_]+\z/, request_id)
   end