diff --git a/README.md b/README.md
index d227118..a905050 100644
--- a/README.md
+++ b/README.md
@@ -56,7 +56,7 @@ defmodule MyPhoenixApp.MyOtherController do
     |> text "The API key used is: #{conn.assigns[:api_key]}"
   end
 
-  def handle_missing_header(conn, missing_header_key) do
+  def handle_missing_header(conn, {_connection_assignment_key, missing_header_key}) do
     conn
     |> send_resp(Status.code(:bad_request), "Missing header: #{missing_header_key}")
     |> halt
@@ -81,7 +81,7 @@ avoiding this is to have your callback function pattern match on the state of th
   plug PlugRequireHeader, headers: [api_key: "x-api-key", secret: "x-secret"], on_missing: {__MODULE__, :handle_missing_header}
 
   def handle_missing_header(%Plug.Conn{state: :sent} = conn, _), do: conn
-  def handle_missing_header(conn, missing_header_key) do
+  def handle_missing_header(conn, {_connection_assignment_key, missing_header_key}) do
     conn
     |> send_resp(Status.code(:bad_request), "Missing header: #{missing_header_key}")
     |> halt
diff --git a/lib/plug_require_header.ex b/lib/plug_require_header.ex
index ace5502..d324a4d 100644
--- a/lib/plug_require_header.ex
+++ b/lib/plug_require_header.ex
@@ -2,7 +2,7 @@ defmodule PlugRequireHeader do
   import Plug.Conn
   alias Plug.Conn.Status
 
-  @vsn "0.3.2"
+  @vsn "0.4.0"
   @doc false
   def version, do: @vsn
 
@@ -63,7 +63,7 @@ defmodule PlugRequireHeader do
   defp extract_header_key(conn, {connection_key, header_key}, callback) do
     case List.keyfind(conn.req_headers, header_key, 0) do
       {^header_key, value} -> assign_connection_key(conn, connection_key, value)
-      _ -> callback.(conn, header_key)
+      _ -> callback.(conn, {connection_key, header_key})
     end
   end
 
diff --git a/mix.exs b/mix.exs
index e9187cc..eebe983 100644
--- a/mix.exs
+++ b/mix.exs
@@ -4,7 +4,7 @@ defmodule PlugRequireHeader.Mixfile do
   def project do
     [
       app: :plug_require_header,
-      version: "0.3.2",
+      version: "0.4.0",
       name: "PlugRequireHeader",
       source_url: "https://github.com/DevL/plug_require_header",
       elixir: "~> 1.0",
diff --git a/test/test_helper.exs b/test/test_helper.exs
index 935b0f6..60792e4 100644
--- a/test/test_helper.exs
+++ b/test/test_helper.exs
@@ -28,7 +28,7 @@ defmodule TestAppWithCallback do
     send_resp(conn, Status.code(:ok), "#{conn.assigns[:api_key]}")
   end
 
-  def callback(conn, missing_header_key) do
+  def callback(conn, {_, missing_header_key}) do
     conn
     |> send_resp(Status.code(:precondition_failed), "Missing header: #{missing_header_key}")
     |> halt
@@ -50,11 +50,7 @@ defmodule TestAppWithCallbackAndMultipleRequiredHeaders do
     send_resp(conn, Status.code(:ok), "API key: #{conn.assigns[:api_key]} and the secret #{conn.assigns[:secret]}")
   end
 
-  def callback(conn, "x-api-key") do
-    conn |> assign :api_key, "not available"
-  end
-
-  def callback(conn, "x-secret") do
-    conn |> assign :secret, "is missing"
+  def callback(conn, {connection_key, _}) do
+    conn |> assign connection_key, "is missing"
   end
 end