Refactor PSBT module & add new PSBT fields

lib/extendedkey.ex

@@ -42,8 +42,12 @@ defmodule Bitcoinex.ExtendedKey do
 
     def new(), do: %__MODULE__{child_nums: []}
 
-    @spec to_string(t()) :: {:ok, String.t()} | {:error, String.t()}
-    def to_string(%__MODULE__{child_nums: path}), do: tto_string(path, "")
+    @spec serialize(t(), atom) :: {:ok, String.t()} | {:ok, binary} | {:error, String.t()}
+    def serialize(dp = %__MODULE__{}, :to_string), do: path_to_string(dp)
+    def serialize(dp = %__MODULE__{}, :to_bin), do: to_bin(dp)
+
+    @spec path_to_string(t()) :: {:ok, String.t()} | {:error, String.t()}
+    def path_to_string(%__MODULE__{child_nums: path}), do: tto_string(path, "")
 
     defp tto_string([], path_acc), do: {:ok, path_acc}
 
@@ -78,27 +82,105 @@ defmodule Bitcoinex.ExtendedKey do
       end
     end
 
-    @spec from_string(String.t()) :: {:ok, t()} | {:error, String.t()}
-    def from_string(pathstr) do
+    @spec to_bin(t()) :: {:ok, binary} | {:error, String.t()}
+    def to_bin(%__MODULE__{child_nums: child_nums}) do
       try do
-        {:ok, %__MODULE__{child_nums: tfrom_string(String.split(pathstr, "/"))}}
+        {:ok, tto_bin(child_nums, <<>>)}
       rescue
         e in ArgumentError -> {:error, e.message}
       end
     end
 
-    defp tfrom_string(path_list) do
+    defp tto_bin([], path_acc), do: path_acc
+
+    defp tto_bin([lvl | rest], path_acc) do
+      cond do
+        lvl == :any or lvl == :anyh ->
+          raise(ArgumentError,
+            message: "Derivation Path with wildcard cannot be encoded to binary."
+          )
+
+        lvl > @max_hardened_child_num ->
+          raise(ArgumentError, message: "index cannot be greater than #{@max_hardened_child_num}")
+
+        lvl < @min_non_hardened_child_num ->
+          raise(ArgumentError, message: "index cannot be less than #{@min_non_hardened_child_num}")
+
+        true ->
+          lvlbin =
+            lvl
+            |> :binary.encode_unsigned(:little)
+            |> Bitcoinex.Utils.pad(4, :trailing)
+
+          tto_bin(rest, path_acc <> lvlbin)
+      end
+    end
+
+    @spec parse(binary, atom) :: {:ok, t()} | {:error, String.t()}
+    def parse(dp, :from_bin), do: from_bin(dp)
+    def parse(dp, :from_string), do: path_from_string(dp)
+
+    @spec path_from_string(String.t()) :: {:ok, t()} | {:error, String.t()}
+    def path_from_string(pathstr) do
+      try do
+        {:ok,
+         %__MODULE__{
+           child_nums:
+             pathstr
+             |> String.split("/")
+             |> tfrom_string([])
+             |> Enum.reverse()
+         }}
+      rescue
+        e in ArgumentError -> {:error, e.message}
+      end
+    end
+
+    defp tfrom_string(path_list, child_nums) do
       case path_list do
-        [] -> []
-        [""] -> []
-        ["m" | rest] -> tfrom_string(rest)
-        ["*" | rest] -> [:any | tfrom_string(rest)]
-        ["*'" | rest] -> [:anyh | tfrom_string(rest)]
-        ["*h" | rest] -> [:anyh | tfrom_string(rest)]
-        [i | rest] -> [str_to_level(i) | tfrom_string(rest)]
+        [] ->
+          child_nums
+
+        [""] ->
+          child_nums
+
+        ["m" | rest] ->
+          if child_nums != [] do
+            raise(ArgumentError,
+              message: "m can only be present at the begining of a derivation path."
+            )
+          else
+            tfrom_string(rest, child_nums)
+          end
+
+        ["*" | rest] ->
+          tfrom_string(rest, [:any | child_nums])
+
+        ["*'" | rest] ->
+          tfrom_string(rest, [:anyh | child_nums])
+
+        ["*h" | rest] ->
+          tfrom_string(rest, [:anyh | child_nums])
+
+        [i | rest] ->
+          tfrom_string(rest, [str_to_level(i) | child_nums])
+      end
+    end
+
+    @spec from_bin(binary) :: {:ok, t()} | {:error, String.t()}
+    def from_bin(bin) do
+      try do
+        {:ok, %__MODULE__{child_nums: Enum.reverse(tfrom_bin(bin, []))}}
+      rescue
+        _e in ArgumentError -> {:error, "invalid binary encoding of derivation path"}
       end
     end
 
+    defp tfrom_bin(<<>>, child_nums), do: child_nums
+
+    defp tfrom_bin(<<level::little-unsigned-32, bin::binary>>, child_nums),
+      do: tfrom_bin(bin, [level | child_nums])
+
     defp str_to_level(level) do
       {num, is_hardened} =
         case String.split(level, ["'", "h"]) do
@@ -111,6 +193,7 @@ defmodule Bitcoinex.ExtendedKey do
 
       nnum = String.to_integer(num)
 
+      # TODO benchmark and make this two comparisons
       if nnum in @min_non_hardened_child_num..@max_non_hardened_child_num do
         if is_hardened do
           nnum + @min_hardened_child_num
@@ -124,6 +207,8 @@ defmodule Bitcoinex.ExtendedKey do
 
     def add(%__MODULE__{child_nums: path1}, %__MODULE__{child_nums: path2}),
       do: %__MODULE__{child_nums: path1 ++ path2}
+
+    def depth(%__MODULE__{child_nums: child_nums}), do: length(child_nums)
   end
 
   @type t :: %__MODULE__{
@@ -245,11 +330,25 @@ defmodule Bitcoinex.ExtendedKey do
   # PARSE & SERIALIZE 
 
   @doc """
-    parse_extended_key takes binary or string representation 
+    parse! calls parse, which takes binary or string representation 
+    of an extended key and parses it to an extended key object.
+    parse! raises ArgumentError on failure. 
+  """
+  @spec parse!(binary) :: t()
+  def parse!(xpub) do
+    case parse(xpub) do
+      {:ok, res} -> res
+      {:error, msg} -> raise(ArgumentError, message: msg)
+    end
+  end
+
+  @doc """
+    parse takes binary or string representation 
     of an extended key and parses it to an extended key object
+    returns {:error, msg} on failure
   """
-  @spec parse_extended_key(binary) :: {:ok, t()} | {:error, String.t()}
-  def parse_extended_key(
+  @spec parse(binary) :: {:ok, t()} | {:error, String.t()}
+  def parse(
         xkey =
           <<prefix::binary-size(4), depth::binary-size(1), parent_fingerprint::binary-size(4),
             child_num::binary-size(4), chaincode::binary-size(32), key::binary-size(33),
@@ -283,16 +382,27 @@ defmodule Bitcoinex.ExtendedKey do
     end
   end
 
+  # parse without checksum (used for PSBT encodings)
+  def parse(
+        xkey =
+          <<_prefix::binary-size(4), _depth::binary-size(1), _parent_fingerprint::binary-size(4),
+            _child_num::binary-size(4), _chaincode::binary-size(32), _key::binary-size(33)>>
+      ) do
+    xkey
+    |> Base58.append_checksum()
+    |> parse()
+  end
+
   # parse from string
-  def parse_extended_key(xkey) do
+  def parse(xkey) do
     case Base58.decode(xkey) do
       {:error, _} ->
         {:error, "error parsing key"}
 
       {:ok, xkey} ->
         xkey
         |> Base58.append_checksum()
-        |> parse_extended_key()
+        |> parse()
     end
   end
 
@@ -303,23 +413,34 @@ defmodule Bitcoinex.ExtendedKey do
   end
 
   @doc """
-    serialize_extended_key takes an extended key
+    serialize takes an extended key
     and returns the binary
   """
-  @spec serialize_extended_key(t()) :: binary
-  def serialize_extended_key(xkey) do
+  @spec serialize(t()) :: binary
+  def serialize(xkey) do
     (xkey.prefix <>
        xkey.depth <> xkey.parent_fingerprint <> xkey.child_num <> xkey.chaincode <> xkey.key)
     |> Base58.append_checksum()
   end
 
+  @doc """
+    serialize takes an extended key
+    and returns the binary without the checksum appended
+    (used for PSBT encoding)
+  """
+  @spec serialize(t(), atom) :: binary
+  def serialize(xkey = %__MODULE__{}, :no_checksum) do
+    xkey.prefix <>
+      xkey.depth <> xkey.parent_fingerprint <> xkey.child_num <> xkey.chaincode <> xkey.key
+  end
+
   @doc """
     display returns the extended key as a string
   """
-  @spec display_extended_key(t()) :: String.t()
-  def display_extended_key(xkey) do
+  @spec display(t()) :: String.t()
+  def display(xkey) do
     xkey
-    |> serialize_extended_key()
+    |> serialize()
     |> Base58.encode_base()
   end
 
@@ -339,7 +460,7 @@ defmodule Bitcoinex.ExtendedKey do
 
       (prefix <> depth_fingerprint_childnum <> chaincode <> <<0>> <> key)
       |> Base58.append_checksum()
-      |> parse_extended_key()
+      |> parse()
     else
       {:error, "invalid extended private key prefix"}
     end
@@ -368,7 +489,7 @@ defmodule Bitcoinex.ExtendedKey do
         |> Kernel.<>(xprv.chaincode)
         |> Kernel.<>(pubkey)
         |> Base58.append_checksum()
-        |> parse_extended_key()
+        |> parse()
       rescue
         _ in MatchError -> {:error, "invalid private key"}
       end
@@ -475,7 +596,7 @@ defmodule Bitcoinex.ExtendedKey do
             (xkey.prefix <>
                child_depth <> fingerprint <> i <> child_chaincode <> Point.sec(pubkey))
             |> Base58.append_checksum()
-            |> parse_extended_key()
+            |> parse()
           end
         end
     end
@@ -526,7 +647,7 @@ defmodule Bitcoinex.ExtendedKey do
 
       (xkey.prefix <> child_depth <> fingerprint <> i <> child_chaincode <> <<0>> <> child_key)
       |> Base58.append_checksum()
-      |> parse_extended_key()
+      |> parse()
     rescue
       _ in MatchError -> {:error, "invalid private key in extended private key"}
     end