Refactor PSBT module & add new PSBT fields

lib/extendedkey.ex

@@ -42,18 +42,21 @@ defmodule Bitcoinex.ExtendedKey do
 
     def new(), do: %__MODULE__{child_nums: []}
 
+    @spec serialize(t()) :: {:ok, binary} | {:error, String.t()}
+    def serialize(dp = %__MODULE__{}), do: to_bin(dp)
+
     @spec to_string(t()) :: {:ok, String.t()} | {:error, String.t()}
-    def to_string(%__MODULE__{child_nums: path}), do: tto_string(path, "")
+    def to_string(%__MODULE__{child_nums: path}), do: path_to_string(path, "")
 
-    defp tto_string([], path_acc), do: {:ok, path_acc}
+    defp path_to_string([], path_acc), do: {:ok, path_acc}
 
-    defp tto_string([l | rest], path_acc) do
+    defp path_to_string([l | rest], path_acc) do
       cond do
         l == :any ->
-          tto_string(rest, path_acc <> "*/")
+          path_to_string(rest, path_acc <> "*/")
 
         l == :anyh ->
-          tto_string(rest, path_acc <> "*'/")
+          path_to_string(rest, path_acc <> "*'/")
 
         l > @max_hardened_child_num ->
           {:error, "index cannot be greater than #{@max_hardened_child_num}"}
@@ -63,7 +66,7 @@ defmodule Bitcoinex.ExtendedKey do
 
         # hardened
         l >= @min_hardened_child_num ->
-          tto_string(
+          path_to_string(
             rest,
             path_acc <>
               (l
@@ -74,31 +77,108 @@ defmodule Bitcoinex.ExtendedKey do
 
         # unhardened
         true ->
-          tto_string(rest, path_acc <> Integer.to_string(l) <> "/")
+          path_to_string(rest, path_acc <> Integer.to_string(l) <> "/")
+      end
+    end
+
+    @spec to_bin(t()) :: {:ok, binary} | {:error, String.t()}
+    def to_bin(%__MODULE__{child_nums: child_nums}) do
+      try do
+        {:ok, to_bin(child_nums, <<>>)}
+      rescue
+        e in ArgumentError -> {:error, e.message}
+      end
+    end
+
+    defp to_bin([], path_acc), do: path_acc
+
+    defp to_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)
+
+          to_bin(rest, path_acc <> lvlbin)
       end
     end
 
+    @spec parse(binary) :: {:ok, t()} | {:error, String.t()}
+    def parse(dp), do: from_bin(dp)
+
     @spec from_string(String.t()) :: {:ok, t()} | {:error, String.t()}
     def from_string(pathstr) do
       try do
-        {:ok, %__MODULE__{child_nums: tfrom_string(String.split(pathstr, "/"))}}
+        {:ok,
+         %__MODULE__{
+           child_nums:
+             pathstr
+             |> String.split("/")
+             |> path_from_string([])
+             |> Enum.reverse()
+         }}
       rescue
         e in ArgumentError -> {:error, e.message}
       end
     end
 
-    defp tfrom_string(path_list) do
+    defp path_from_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
+            path_from_string(rest, child_nums)
+          end
+
+        ["*" | rest] ->
+          path_from_string(rest, [:any | child_nums])
+
+        ["*'" | rest] ->
+          path_from_string(rest, [:anyh | child_nums])
+
+        ["*h" | rest] ->
+          path_from_string(rest, [:anyh | child_nums])
+
+        [i | rest] ->
+          path_from_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(from_bin(bin, []))}}
+      rescue
+        _e in ArgumentError -> {:error, "invalid binary encoding of derivation path"}
+      end
+    end
+
+    defp from_bin(<<>>, child_nums), do: child_nums
+
+    defp from_bin(<<level::little-unsigned-32, bin::binary>>, child_nums),
+      do: from_bin(bin, [level | child_nums])
+
     defp str_to_level(level) do
       {num, is_hardened} =
         case String.split(level, ["'", "h"]) do
@@ -111,6 +191,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 +205,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 +328,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 +380,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 +411,33 @@ 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
-    (xkey.prefix <>
-       xkey.depth <> xkey.parent_fingerprint <> xkey.child_num <> xkey.chaincode <> xkey.key)
-    |> Base58.append_checksum()
+  @spec serialize(t(), list({:with_checksum?, boolean})) :: binary
+  def serialize(xkey, opts \\ []) do
+    with_checksum? = Keyword.get(opts, :with_checksum?, true)
+
+    extended_key_without_checksum_bin =
+      xkey.prefix <>
+        xkey.depth <> xkey.parent_fingerprint <> xkey.child_num <> xkey.chaincode <> xkey.key
+
+    case with_checksum? do
+      true ->
+        Base58.append_checksum(extended_key_without_checksum_bin)
+
+      false ->
+        extended_key_without_checksum_bin
+    end
   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 +457,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 +486,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 +593,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 +644,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