Improve unpacker validation in Msgpax.PlugParser

lexmag · Jan 4, 2021 · 66c93b5 · 66c93b5
1 parent b3530ac
commit 66c93b5
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 34 deletions.
diff --git a/lib/msgpax/plug_parser.ex b/lib/msgpax/plug_parser.ex
@@ -1,4 +1,4 @@
-if Code.ensure_compiled?(Plug) do
+if Code.ensure_loaded?(Plug) do
   defmodule Msgpax.PlugParser do
     @moduledoc """
     A `Plug.Parsers` plug for parsing a MessagePack-encoded body.
@@ -39,7 +39,7 @@ if Code.ensure_compiled?(Plug) do
     def parse(%Plug.Conn{} = conn, "application", "msgpack", _params, {unpacker, options}) do
       case read_body(conn, options) do
         {:ok, body, conn} ->
-          {:ok, unpack_body(body, unpacker), conn}
+          {:ok, unpack(body, unpacker), conn}
 
         {:more, _partial_body, conn} ->
           {:error, :too_large, conn}
@@ -52,54 +52,43 @@ if Code.ensure_compiled?(Plug) do
 
     def init(options) do
       {unpacker, options} = Keyword.pop(options, :unpacker, Msgpax)
-
-      validate_unpacker!(unpacker)
-
-      {unpacker, options}
+      {validate_unpacker!(unpacker), options}
     end
 
-    defp unpack_body(body, unpacker) do
-      case apply_mfa_or_module(body, unpacker) do
-        data = %_{} -> %{"_msgpack" => data}
+    defp unpack(body, {module, function, extra_args}) do
+      try do
+        apply(module, function, [body | extra_args])
+      rescue
+        exception ->
+          raise Plug.Parsers.ParseError, exception: exception
+      else
+        %_{} = data -> %{"_msgpack" => data}
         data when is_map(data) -> data
         data -> %{"_msgpack" => data}
       end
-    rescue
-      exception ->
-        raise Plug.Parsers.ParseError, exception: exception
-    end
-
-    defp apply_mfa_or_module(body, {module, function, extra_args}) do
-      apply(module, function, [body | extra_args])
     end
 
-    defp apply_mfa_or_module(body, unpacker) do
-      unpacker.unpack!(body)
-    end
-
-    defp validate_unpacker!({module, function, extra_args})
+    defp validate_unpacker!({module, function, extra_args} = unpacker)
          when is_atom(module) and is_atom(function) and is_list(extra_args) do
       arity = length(extra_args) + 1
 
-      unless Code.ensure_compiled?(module) and function_exported?(module, function, arity) do
-        raise ArgumentError,
-              "invalid :unpacker option. Undefined function " <>
-                Exception.format_mfa(module, function, arity)
-      end
-    end
-
-    defp validate_unpacker!(unpacker) when is_atom(unpacker) do
-      unless Code.ensure_compiled?(unpacker) do
+      if Code.ensure_compiled(module) != {:module, module} do
         raise ArgumentError,
               "invalid :unpacker option. The module #{inspect(unpacker)} is not " <>
                 "loaded and could not be found"
       end
 
-      unless function_exported?(unpacker, :unpack!, 1) do
+      if not function_exported?(module, function, arity) do
         raise ArgumentError,
-              "invalid :unpacker option. The module #{inspect(unpacker)} must " <>
-                "implement unpack!/1"
+              "invalid :unpacker option. The module #{inspect(module)} must " <>
+                "implement #{function}/#{arity}"
       end
+
+      unpacker
+    end
+
+    defp validate_unpacker!(unpacker) when is_atom(unpacker) do
+      validate_unpacker!({unpacker, :unpack!, []})
     end
 
     defp validate_unpacker!(unpacker) do

diff --git a/mix.lock b/mix.lock
@@ -5,6 +5,6 @@
   "makeup_elixir": {:hex, :makeup_elixir, "0.15.0", "98312c9f0d3730fde4049985a1105da5155bfe5c11e47bdc7406d88e01e4219b", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.1", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "75ffa34ab1056b7e24844c90bfc62aaf6f3a37a15faa76b07bc5eba27e4a8b4a"},
   "mime": {:hex, :mime, "1.3.0", "5e8d45a39e95c650900d03f897fbf99ae04f60ab1daa4a34c7a20a5151b7a5fe", [:mix], [], "hexpm", "5e839994289d60326aa86020c4fbd9c6938af188ecddab2579f07b66cd665328"},
   "nimble_parsec": {:hex, :nimble_parsec, "1.1.0", "3a6fca1550363552e54c216debb6a9e95bd8d32348938e13de5eda962c0d7f89", [:mix], [], "hexpm", "08eb32d66b706e913ff748f11694b17981c0b04a33ef470e33e11b3d3ac8f54b"},
-  "plug": {:hex, :plug, "1.7.1", "8516d565fb84a6a8b2ca722e74e2cd25ca0fc9d64f364ec9dbec09d33eb78ccd", [:mix], [{:mime, "~> 1.0", [hex: :mime, repo: "hexpm", optional: false]}, {:plug_crypto, "~> 1.0", [hex: :plug_crypto, repo: "hexpm", optional: false]}], "hexpm", "daa5fee4209c12c3c48b05a96cf88c320b627c9575f987554dcdc1fdcdf2c15e"},
+  "plug": {:hex, :plug, "1.7.2", "d7b7db7fbd755e8283b6c0a50be71ec0a3d67d9213d74422d9372effc8e87fd1", [:mix], [{:mime, "~> 1.0", [hex: :mime, repo: "hexpm", optional: false]}, {:plug_crypto, "~> 1.0", [hex: :plug_crypto, repo: "hexpm", optional: false]}], "hexpm", "de9825f21c6fd6adfdeae8f9c80dcd88c1e58301f06bf13d659b7e606b88abe0"},
   "plug_crypto": {:hex, :plug_crypto, "1.0.0", "18e49317d3fa343f24620ed22795ec29d4a5e602d52d1513ccea0b07d8ea7d4d", [:mix], [], "hexpm", "73c1682f0e414cfb5d9b95c8e8cd6ffcfdae699e3b05e1db744e58b7be857759"},
 }