feat: enable ssz generic spec tests (#683)

Co-authored-by: Godspower Eze <[email protected]>
Co-authored-by: Godspower Eze <[email protected]>

Author: f3r10

lib/ssz_ex.ex

@@ -3,8 +3,8 @@ defmodule LambdaEthereumConsensus.SszEx do
     SSZ library in Elixir
   """
   alias LambdaEthereumConsensus.Utils.BitList
-  alias LambdaEthereumConsensus.Utils.BitVector
   import alias LambdaEthereumConsensus.Utils.BitVector
+  import Bitwise
   alias LambdaEthereumConsensus.Utils.ZeroHashes
 
   #################
@@ -34,8 +34,11 @@ defmodule LambdaEthereumConsensus.SszEx do
       else: encode_fixed_size_list(list, basic_type, size)
   end
 
-  def encode(vector, {:vector, basic_type, size}),
-    do: encode_fixed_size_list(vector, basic_type, size)
+  def encode(vector, {:vector, basic_type, size}) do
+    if variable_size?(basic_type),
+      do: encode_variable_size_list(vector, basic_type, size),
+      else: encode_fixed_size_list(vector, basic_type, size)
+  end
 
   def encode(value, {:bitlist, max_size}) when is_bitstring(value),
     do: encode_bitlist(value, max_size)
@@ -56,18 +59,26 @@ defmodule LambdaEthereumConsensus.SszEx do
   def decode(binary, {:list, basic_type, size}) do
     if variable_size?(basic_type),
       do: decode_variable_list(binary, basic_type, size),
-      else: decode_list(binary, basic_type, size)
+      else: decode_fixed_list(binary, basic_type, size)
   end
 
-  def decode(binary, {:vector, basic_type, size}), do: decode_list(binary, basic_type, size)
+  def decode(binary, {:vector, basic_type, size}) do
+    if variable_size?(basic_type),
+      do: decode_variable_list(binary, basic_type, size),
+      else: decode_fixed_vector(binary, basic_type, size)
+  end
 
   def decode(value, {:bitlist, max_size}) when is_bitstring(value),
     do: decode_bitlist(value, max_size)
 
   def decode(value, {:bitvector, size}) when is_bitstring(value),
     do: decode_bitvector(value, size)
 
-  def decode(binary, module) when is_atom(module), do: decode_container(binary, module)
+  def decode(binary, module) when is_atom(module) do
+    if variable_size?(module),
+      do: decode_variable_container(binary, module),
+      else: decode_fixed_container(binary, module)
+  end
 
   @spec hash_tree_root!(boolean, atom) :: Types.root()
   def hash_tree_root!(value, :bool), do: pack(value, :bool)
@@ -274,13 +285,16 @@ defmodule LambdaEthereumConsensus.SszEx do
   defp encode_bool(true), do: {:ok, "\x01"}
   defp encode_bool(false), do: {:ok, "\x00"}
 
-  defp decode_uint(binary, size) do
+  defp decode_uint(binary, size) when bit_size(binary) == size do
     <<element::integer-size(size)-little, _rest::bitstring>> = binary
     {:ok, element}
   end
 
+  defp decode_uint(_binary, size), do: {:error, "invalid byte size #{inspect(size)}"}
+
   defp decode_bool("\x01"), do: {:ok, true}
   defp decode_bool("\x00"), do: {:ok, false}
+  defp decode_bool(_), do: {:error, "invalid bool value"}
 
   defp encode_fixed_size_list(list, _basic_type, max_size) when length(list) > max_size,
     do: {:error, "invalid max_size of list"}
@@ -348,7 +362,7 @@ defmodule LambdaEthereumConsensus.SszEx do
     end
   end
 
-  defp decode_bitlist(bit_list, max_size) do
+  defp decode_bitlist(bit_list, max_size) when bit_size(bit_list) > 0 do
     num_bytes = byte_size(bit_list)
     {decoded, len} = BitList.new(bit_list)
 
@@ -367,26 +381,70 @@ defmodule LambdaEthereumConsensus.SszEx do
     end
   end
 
-  defp decode_bitvector(bit_vector, size) when bit_size(bit_vector) == size,
-    do: {:ok, BitVector.new(bit_vector, size)}
+  defp decode_bitlist(_bit_list, _max_size), do: {:error, "invalid bitlist"}
 
-  defp decode_bitvector(_bit_vector, _size), do: {:error, "invalid bit_vector length"}
+  defp decode_bitvector(bit_vector, size) do
+    num_bytes = byte_size(bit_vector)
 
-  defp decode_list(binary, basic_type, size) do
+    cond do
+      bit_size(bit_vector) == 0 ->
+        {:error, "ExcessBits"}
+
+      num_bytes != max(1, div(size + 7, 8)) ->
+        {:error, "InvalidByteCount"}
+
+      true ->
+        case bit_vector do
+          # Padding bits are clear
+          <<_first::binary-size(num_bytes - 1), 0::size(8 - rem(size, 8) &&& 7),
+            _rest::bitstring>> ->
+            {:ok, BitVector.new(bit_vector, size)}
+
+          _else ->
+            {:error, "ExcessBits"}
+        end
+    end
+  end
+
+  defp decode_fixed_list(binary, basic_type, size) do
     fixed_size = get_fixed_size(basic_type)
 
-    with {:ok, decoded_list} = result <-
-           binary
-           |> decode_chunk(fixed_size, basic_type)
-           |> flatten_results() do
-      if length(decoded_list) > size do
-        {:error, "invalid max_size of list"}
-      else
-        result
-      end
+    with {:ok, decoded_list} = result <- decode_fixed_collection(binary, fixed_size, basic_type),
+         :ok <- check_valid_list_size_after_decode(size, length(decoded_list)) do
+      result
+    end
+  end
+
+  defp decode_fixed_vector(binary, basic_type, size) do
+    fixed_size = get_fixed_size(basic_type)
+
+    with :ok <- check_valid_vector_size_prev_decode(fixed_size, size, binary),
+         {:ok, decoded_vector} = result <-
+           decode_fixed_collection(binary, fixed_size, basic_type),
+         :ok <- check_valid_vector_size_after_decode(size, length(decoded_vector)) do
+      result
     end
   end
 
+  def check_valid_vector_size_prev_decode(fixed_size, size, binary)
+      when fixed_size * size == byte_size(binary),
+      do: :ok
+
+  def check_valid_vector_size_prev_decode(_fixed_size, _size, _binary),
+    do: {:error, "Invalid vector size"}
+
+  def check_valid_vector_size_after_decode(size, decoded_size)
+      when decoded_size == size and decoded_size > 0,
+      do: :ok
+
+  def check_valid_vector_size_after_decode(_size, _decoded_size),
+    do: {:error, "invalid vector decoded size"}
+
+  def check_valid_list_size_after_decode(size, decoded_size) when decoded_size <= size, do: :ok
+
+  def check_valid_list_size_after_decode(_size, _decoded_size),
+    do: {:error, "invalid max_size of list"}
+
   defp decode_variable_list(binary, _, _) when byte_size(binary) == 0 do
     {:ok, []}
   end
@@ -407,7 +465,7 @@ defmodule LambdaEthereumConsensus.SszEx do
          first_offset < @bytes_per_length_offset do
       {:error, "InvalidListFixedBytesLen"}
     else
-      with {:ok, first_offset} <-
+      with :ok <-
              sanitize_offset(first_offset, nil, byte_size(binary), first_offset) do
         decode_variable_list_elements(
           num_elements,
@@ -448,7 +506,7 @@ defmodule LambdaEthereumConsensus.SszEx do
        ) do
     <<next_offset::integer-32-little, rest_bytes::bitstring>> = acc_rest_bytes
 
-    with {:ok, next_offset} <-
+    with :ok <-
            sanitize_offset(next_offset, offset, byte_size(binary), first_offset) do
       part = :binary.part(binary, offset, next_offset - offset)
 
@@ -467,7 +525,7 @@ defmodule LambdaEthereumConsensus.SszEx do
   defp encode_container(container, schemas) do
     {fixed_size_values, fixed_length, variable_values} = analyze_schemas(container, schemas)
 
-    with {:ok, variable_parts} <- encode_schemas(variable_values),
+    with {:ok, variable_parts} <- encode_schemas(Enum.reverse(variable_values)),
          offsets = calculate_offsets(variable_parts, fixed_length),
          variable_length =
            Enum.reduce(variable_parts, 0, fn part, acc -> byte_size(part) + acc end),
@@ -524,48 +582,87 @@ defmodule LambdaEthereumConsensus.SszEx do
   defp replace_offset(element, {acc_fixed_list, acc_offsets_list}),
     do: {[element | acc_fixed_list], acc_offsets_list}
 
-  defp decode_container(binary, module) do
+  defp decode_variable_container(binary, module) do
     schemas = module.schema()
     fixed_length = get_fixed_length(schemas)
-    <<fixed_binary::binary-size(fixed_length), variable_binary::bitstring>> = binary
 
-    with {:ok, fixed_parts, offsets} <- decode_fixed_section(fixed_binary, schemas, fixed_length),
-         {:ok, variable_parts} <- decode_variable_section(variable_binary, offsets) do
+    with :ok <- sanitize_offset(fixed_length, nil, byte_size(binary), nil),
+         <<fixed_binary::binary-size(fixed_length), variable_binary::bitstring>> = binary,
+         {:ok, fixed_parts, offsets, items_index} <-
+           decode_fixed_section(fixed_binary, schemas, fixed_length),
+         :ok <- check_first_offset(offsets, items_index, byte_size(binary)),
+         {:ok, variable_parts} <- decode_variable_section(binary, variable_binary, offsets) do
       {:ok, struct!(module, fixed_parts ++ variable_parts)}
     end
   end
 
-  defp decode_variable_section(binary, offsets) do
+  defp decode_fixed_container(binary, module) do
+    schemas = module.schema()
+    fixed_length = get_fixed_length(schemas)
+
+    with {:ok, fixed_parts, _offsets, items_index} <-
+           decode_fixed_section(binary, schemas, fixed_length),
+         :ok <- check_byte_len(items_index, byte_size(binary)) do
+      {:ok, struct!(module, fixed_parts)}
+    end
+  end
+
+  defp check_first_offset([{offset, _} | _rest], items_index, _binary_size) do
+    cond do
+      offset < items_index -> {:error, "OffsetIntoFixedPortion"}
+      offset > items_index -> {:error, "OffsetSkipsVariableBytes"}
+      true -> :ok
+    end
+  end
+
+  defp check_byte_len(items_index, binary_size)
+       when items_index == binary_size,
+       do: :ok
+
+  defp check_byte_len(_items_index, _binary_size),
+    do: {:error, "InvalidByteLength"}
+
+  defp decode_variable_section(full_binary, binary, offsets) do
     offsets
     |> Enum.chunk_every(2, 1)
-    |> Enum.reduce({binary, []}, fn
+    |> Enum.reduce_while({binary, []}, fn
       [{offset, {key, schema}}, {next_offset, _}], {rest_bytes, acc_variable_parts} ->
-        size = next_offset - offset
-        <<chunk::binary-size(size), rest::bitstring>> = rest_bytes
-        {rest, [{key, decode(chunk, schema)} | acc_variable_parts]}
+        case sanitize_offset(next_offset, offset, byte_size(full_binary), nil) do
+          :ok ->
+            size = next_offset - offset
+            <<chunk::binary-size(size), rest::bitstring>> = rest_bytes
+            {:cont, {rest, [{key, decode(chunk, schema)} | acc_variable_parts]}}
+
+          error ->
+            {:halt, {<<>>, [{key, error} | acc_variable_parts]}}
+        end
 
       [{_offset, {key, schema}}], {rest_bytes, acc_variable_parts} ->
-        {<<>>, [{key, decode(rest_bytes, schema)} | acc_variable_parts]}
+        {:cont, {<<>>, [{key, decode(rest_bytes, schema)} | acc_variable_parts]}}
     end)
     |> then(fn {<<>>, variable_parts} ->
       flatten_container_results(variable_parts)
     end)
   end
 
-  defp decode_fixed_section(binary, schemas, fixed_length) do
+  defp decode_fixed_section(binary, schemas, _fixed_length) do
     schemas
-    |> Enum.reduce({binary, [], []}, fn {key, schema}, {binary, fixed_parts, offsets} ->
+    |> Enum.reduce({binary, [], [], 0}, fn {key, schema},
+                                           {binary, fixed_parts, offsets, items_index} ->
       if variable_size?(schema) do
         <<offset::integer-size(@offset_bits)-little, rest::bitstring>> = binary
-        {rest, fixed_parts, [{offset - fixed_length, {key, schema}} | offsets]}
+
+        {rest, fixed_parts, [{offset, {key, schema}} | offsets],
+         items_index + @bytes_per_length_offset}
       else
         ssz_fixed_len = get_fixed_size(schema)
         <<chunk::binary-size(ssz_fixed_len), rest::bitstring>> = binary
-        {rest, [{key, decode(chunk, schema)} | fixed_parts], offsets}
+        {rest, [{key, decode(chunk, schema)} | fixed_parts], offsets, items_index + ssz_fixed_len}
       end
     end)
-    |> then(fn {_rest_bytes, fixed_parts, offsets} ->
+    |> then(fn {_rest_bytes, fixed_parts, offsets, items_index} ->
       Tuple.append(flatten_container_results(fixed_parts), Enum.reverse(offsets))
+      |> Tuple.append(items_index)
     end)
   end
 
@@ -582,35 +679,47 @@ defmodule LambdaEthereumConsensus.SszEx do
   end
 
   # https://notes.ethereum.org/ruKvDXl6QOW3gnqVYb8ezA?view
-  defp sanitize_offset(offset, previous_offset, num_bytes, num_fixed_bytes) do
+  defp sanitize_offset(offset, previous_offset, num_bytes, nil) do
     cond do
-      offset < num_fixed_bytes ->
-        {:error, "OffsetIntoFixedPortion"}
-
-      previous_offset == nil && offset != num_fixed_bytes ->
-        {:error, "OffsetSkipsVariableBytes"}
-
       offset > num_bytes ->
         {:error, "OffsetOutOfBounds"}
 
       previous_offset != nil && previous_offset > offset ->
         {:error, "OffsetsAreDecreasing"}
 
       true ->
-        {:ok, offset}
+        :ok
     end
   end
 
-  defp decode_chunk(binary, chunk_size, basic_type) do
-    decode_chunk(binary, chunk_size, basic_type, [])
+  defp sanitize_offset(offset, previous_offset, _num_bytes, num_fixed_bytes) do
+    cond do
+      offset < num_fixed_bytes ->
+        {:error, "OffsetIntoFixedPortion"}
+
+      previous_offset == nil && offset != num_fixed_bytes ->
+        {:error, "OffsetSkipsVariableBytes"}
+
+      true ->
+        :ok
+    end
+  end
+
+  defp decode_fixed_collection(binary, chunk_size, basic_type) do
+    decode_fixed_collection(binary, chunk_size, basic_type, [])
     |> Enum.reverse()
+    |> flatten_results()
   end
 
-  defp decode_chunk(<<>>, _chunk_size, _basic_type, results), do: results
+  defp decode_fixed_collection(<<>>, _chunk_size, _basic_type, results), do: results
+
+  defp decode_fixed_collection(binary, chunk_size, _basic_type, results)
+       when byte_size(binary) < chunk_size,
+       do: [{:error, "InvalidByteLength"} | results]
 
-  defp decode_chunk(binary, chunk_size, basic_type, results) do
+  defp decode_fixed_collection(binary, chunk_size, basic_type, results) do
     <<element::binary-size(chunk_size), rest::bitstring>> = binary
-    decode_chunk(rest, chunk_size, basic_type, [decode(element, basic_type) | results])
+    decode_fixed_collection(rest, chunk_size, basic_type, [decode(element, basic_type) | results])
   end
 
   defp flatten_results(results) do
@@ -645,7 +754,8 @@ defmodule LambdaEthereumConsensus.SszEx do
   defp get_fixed_size(:bool), do: 1
   defp get_fixed_size({:int, size}), do: div(size, @bits_per_byte)
   defp get_fixed_size({:bytes, size}), do: size
-  defp get_fixed_size({:vector, _, size}), do: size
+  defp get_fixed_size({:vector, basic_type, size}), do: size * get_fixed_size(basic_type)
+  defp get_fixed_size({:bitvector, _}), do: 1
 
   defp get_fixed_size(module) when is_atom(module) do
     schemas = module.schema()
@@ -656,10 +766,12 @@ defmodule LambdaEthereumConsensus.SszEx do
   end
 
   defp variable_size?({:list, _, _}), do: true
-  defp variable_size?({:vector, _, _}), do: false
   defp variable_size?(:bool), do: false
   defp variable_size?({:int, _}), do: false
   defp variable_size?({:bytes, _}), do: false
+  defp variable_size?({:bitlist, _}), do: true
+  defp variable_size?({:bitvector, _}), do: false
+  defp variable_size?({:vector, basic_type, _}), do: variable_size?(basic_type)
 
   defp variable_size?(module) when is_atom(module) do
     module.schema()

test/spec/runners/helpers/ssz_static_containers/complex_test_struct.ex

@@ -35,7 +35,7 @@ defmodule Helpers.SszStaticContainers.ComplexTestStruct do
       {:A, {:int, 16}},
       {:B, {:list, {:int, 16}, 1024}},
       {:C, {:int, 8}},
-      {:D, {:bitlist, 256}},
+      {:D, {:list, {:int, 8}, 256}},
       {:E, VarTestStruct},
       {:F, {:vector, FixedTestStruct, 4}},
       {:G, {:vector, VarTestStruct, 2}}