chore: change rust kzg nif to direct elixir bindings

Dockerfile

@@ -30,18 +30,6 @@ RUN cargo build --release && \
     mv target/release/libbls_nif.so ./libbls_nif.so && \
     rm -rf target/
 
-# kzg nif
-FROM rust:1.81.0 AS kzg_nif_builder
-LABEL stage=builder
-
-RUN mkdir /kzg_nif
-WORKDIR /kzg_nif
-
-COPY ./native/kzg_nif /kzg_nif
-RUN cargo build --release && \
-    mv target/release/libkzg_nif.so ./libkzg_nif.so && \
-    rm -rf target/
-
 # snappy nif
 FROM rust:1.81.0 AS snappy_nif_builder
 LABEL stage=builder
@@ -106,7 +94,6 @@ COPY --from=libp2p_builder /libp2p_port/libp2p_port /consensus/priv/native/libp2
 # TODO: only copy artifacts
 # Copy precompiled rust crates. Rustler stores targets under _build
 COPY --from=bls_nif_builder /bls_nif/libbls_nif.so /consensus/priv/native/libbls_nif.so
-COPY --from=kzg_nif_builder /kzg_nif/libkzg_nif.so /consensus/priv/native/libkzg_nif.so
 COPY --from=snappy_nif_builder /snappy_nif/libsnappy_nif.so /consensus/priv/native/libsnappy_nif.so
 COPY --from=ssz_nif_builder /ssz_nif/libssz_nif.so /consensus/priv/native/libssz_nif.so
 

bench/kzg.exs

@@ -0,0 +1,61 @@
+alias ChainSpec
+alias Kzg
+
+# BLS12-381 Fr modulus
+modulus =
+  0x73EDA753299D7D483339D80809A1D80553BDA402FFFE5BFEFFFFFFFF00000001
+
+rand_field_element = fn ->
+  int =
+    :crypto.strong_rand_bytes(64)
+    |> :binary.decode_unsigned(:big)
+    |> rem(modulus)
+
+  <<int::unsigned-size(256)>>
+end
+
+ok! = fn
+  {:ok, value} -> value
+  other -> raise "expected {:ok, _}, got #{inspect(other)}"
+end
+
+fe_per_blob = ChainSpec.get("FIELD_ELEMENTS_PER_BLOB")
+
+rand_blob = fn ->
+  for _ <- 1..fe_per_blob, into: <<>> do
+    rand_field_element.()
+  end
+end
+
+z = rand_field_element.()
+blob = rand_blob.()
+
+commitment = blob |> Kzg.blob_to_kzg_commitment() |> ok!.()
+proof = blob |> Kzg.compute_blob_kzg_proof(commitment) |> ok!.()
+
+# Only did 4 blobs, we don't have all day...
+batch =
+  Enum.map(1..4, fn _ ->
+    b = rand_blob.()
+    c = b |> Kzg.blob_to_kzg_commitment() |> ok!.()
+    p = b |> Kzg.compute_blob_kzg_proof(c) |> ok!.()
+    {b, c, p}
+  end)
+
+blobs = Enum.map(batch, fn {b, _, _} -> b end)
+commitments = Enum.map(batch, fn {_, c, _} -> c end)
+proofs = Enum.map(batch, fn {_, _, p} -> p end)
+
+Benchee.run(
+  %{
+    "blob_to_kzg_commitment" => fn -> Kzg.blob_to_kzg_commitment(blob) end,
+    "compute_blob_kzg_proof" => fn -> Kzg.compute_blob_kzg_proof(blob, commitment) end,
+    "verify_blob_kzg_proof" => fn -> Kzg.verify_blob_kzg_proof(blob, commitment, proof) end,
+    "verify_blob_kzg_proof_batch (4)" => fn ->
+      Kzg.verify_blob_kzg_proof_batch(blobs, commitments, proofs)
+    end,
+    "compute_kzg_proof (z)" => fn -> Kzg.compute_kzg_proof(blob, z) end
+  },
+  warmup: 2,
+  time: 5
+)

config/config.exs

@@ -49,7 +49,6 @@ config :lambda_ethereum_consensus, :logger, [
 # Avoid compiling Rustler NIFs when `RUSTLER_SKIP_COMPILE` is set
 if System.get_env("RUSTLER_SKIP_COMPILE") do
   config :lambda_ethereum_consensus, Bls, skip_compilation?: true
-  config :lambda_ethereum_consensus, Kzg, skip_compilation?: true
   config :lambda_ethereum_consensus, Snappy, skip_compilation?: true
   config :lambda_ethereum_consensus, Ssz, skip_compilation?: true
 end

lib/kzg.ex

@@ -2,53 +2,86 @@ defmodule Kzg do
   @moduledoc """
   KZG functions
   """
-  use Rustler, otp_app: :lambda_ethereum_consensus, crate: "kzg_nif"
+
+  alias KZG, as: CKZG
+
+  require Logger
+
+  @bytes_per_blob 4096 * 32
+  @bytes_per_commitment 48
+  @bytes_per_proof 48
+  @trusted_setup_path Path.expand("config/kzg/official_trusted_setup.txt", File.cwd!())
 
   @type commitment :: <<_::384>>
-  @type proof :: <<_::768>>
+  @type proof :: <<_::384>>
 
-  @spec blob_to_kzg_commitment(Types.blob()) :: {:ok, commitment()} | {:error, binary()}
-  def blob_to_kzg_commitment(_blob) do
-    :erlang.nif_error(:nif_not_loaded)
+  @spec blob_to_kzg_commitment(Types.blob()) :: {:ok, commitment()} | {:error, atom()}
+  def blob_to_kzg_commitment(blob) do
+    with :ok <- ensure_blob_size(blob) do
+      with_settings(fn s -> CKZG.blob_to_kzg_commitment(blob, s) end)
+    end
   end
 
   @spec compute_kzg_proof(Types.blob(), Types.bytes32()) ::
-          {:ok, {proof(), Types.bytes32()}} | {:error, binary()}
-  def compute_kzg_proof(_blob, _z) do
-    :erlang.nif_error(:nif_not_loaded)
+          {:ok, {proof(), Types.bytes32()}} | {:error, atom()}
+  def compute_kzg_proof(blob, z) do
+    with :ok <- ensure_blob_size(blob),
+         :ok <- ensure_bytes32(z) do
+      with_settings(fn s -> CKZG.compute_kzg_proof(blob, z, s) end)
+    end
   end
 
-  @spec verify_kzg_proof(
-          commitment(),
-          Types.bytes32(),
-          Types.bytes32(),
-          proof()
-        ) ::
-          {:ok, boolean} | {:error, binary()}
-  def verify_kzg_proof(_kzg_commitment, _z, _y, _kzg_proof) do
-    :erlang.nif_error(:nif_not_loaded)
+  @spec verify_kzg_proof(commitment(), Types.bytes32(), Types.bytes32(), proof()) ::
+          {:ok, boolean} | {:error, atom()}
+  def verify_kzg_proof(commitment, z, y, proof) do
+    with :ok <- ensure_commitment_size(commitment),
+         :ok <- ensure_bytes32(z),
+         :ok <- ensure_bytes32(y),
+         :ok <- ensure_proof_size(proof) do
+      with_settings(fn s -> CKZG.verify_kzg_proof(commitment, z, y, proof, s) end)
+    end
   end
 
   @spec compute_blob_kzg_proof(Types.blob(), commitment()) ::
-          {:ok, proof()} | {:error, binary()}
-  def compute_blob_kzg_proof(_blob, _kzg_commitment) do
-    :erlang.nif_error(:nif_not_loaded)
+          {:ok, proof()} | {:error, atom()}
+  def compute_blob_kzg_proof(blob, commitment) do
+    with :ok <- ensure_blob_size(blob),
+         :ok <- ensure_commitment_size(commitment) do
+      with_settings(fn s -> CKZG.compute_blob_kzg_proof(blob, commitment, s) end)
+    end
   end
 
   @spec verify_blob_kzg_proof(Types.blob(), commitment(), proof()) ::
-          {:ok, boolean} | {:error, binary()}
-  def verify_blob_kzg_proof(_blob, _kzg_commitment, _kzg_proof) do
-    :erlang.nif_error(:nif_not_loaded)
+          {:ok, boolean} | {:error, atom()}
+  def verify_blob_kzg_proof(blob, commitment, proof) do
+    with :ok <- ensure_blob_size(blob),
+         :ok <- ensure_commitment_size(commitment),
+         :ok <- ensure_proof_size(proof) do
+      with_settings(fn s ->
+        CKZG.verify_blob_kzg_proof(blob, commitment, proof, s)
+      end)
+    end
   end
 
   @spec verify_blob_kzg_proof_batch(
-          list(Types.blob()),
-          list(commitment()),
-          list(proof())
+          [Types.blob()],
+          [commitment()],
+          [proof()]
         ) ::
-          {:ok, boolean} | {:error, binary()}
-  def verify_blob_kzg_proof_batch(_blobs, _kzg_commitments, _kzg_proofs) do
-    :erlang.nif_error(:nif_not_loaded)
+          {:ok, boolean} | {:error, atom()}
+  def verify_blob_kzg_proof_batch(blobs, kzg_commitments, kzg_proofs) do
+    with :ok <- ensure_all_blob_sizes(blobs),
+         :ok <- ensure_all_commitment_sizes(kzg_commitments),
+         :ok <- ensure_all_proof_sizes(kzg_proofs),
+         :ok <- ensure_batch_lengths_match(blobs, kzg_commitments, kzg_proofs) do
+      blobs_bin = IO.iodata_to_binary(blobs)
+      commitments_bin = IO.iodata_to_binary(kzg_commitments)
+      proofs_bin = IO.iodata_to_binary(kzg_proofs)
+
+      with_settings(fn settings ->
+        CKZG.verify_blob_kzg_proof_batch(blobs_bin, commitments_bin, proofs_bin, settings)
+      end)
+    end
   end
 
   ################
@@ -66,4 +99,64 @@ defmodule Kzg do
       {:error, _} -> false
     end
   end
+
+  defp settings() do
+    :persistent_term.get({__MODULE__, :settings}, fn -> load_settings() end)
+  end
+
+  defp load_settings() do
+    {:ok, settings} = CKZG.load_trusted_setup(@trusted_setup_path, 0)
+    :persistent_term.put({__MODULE__, :settings}, settings)
+    settings
+  end
+
+  defp reset_settings() do
+    Logger.warning("[KZG] Resetting settings resource")
+    :persistent_term.erase({__MODULE__, :settings})
+    load_settings()
+  end
+
+  defp with_settings(fun) do
+    case fun.(settings()) do
+      {:error, :failed_get_settings_resource} -> fun.(reset_settings())
+      other -> other
+    end
+  end
+
+  defp ensure_blob_size(<<_::binary-size(@bytes_per_blob)>>), do: :ok
+  defp ensure_blob_size(_), do: {:error, :invalid_blob_length}
+
+  defp ensure_bytes32(<<_::binary-size(32)>>), do: :ok
+  defp ensure_bytes32(_), do: {:error, :invalid_field_element_length}
+
+  defp ensure_commitment_size(<<_::binary-size(@bytes_per_commitment)>>), do: :ok
+  defp ensure_commitment_size(_), do: {:error, :invalid_commitment_length}
+
+  defp ensure_proof_size(<<_::binary-size(@bytes_per_proof)>>), do: :ok
+  defp ensure_proof_size(_), do: {:error, :invalid_proof_length}
+
+  defp ensure_all_blob_sizes(list) when is_list(list) do
+    if Enum.all?(list, &match?(<<_::binary-size(@bytes_per_blob)>>, &1)),
+      do: :ok,
+      else: {:error, :invalid_blob_length}
+  end
+
+  defp ensure_all_commitment_sizes(list) when is_list(list) do
+    if Enum.all?(list, &match?(<<_::binary-size(@bytes_per_commitment)>>, &1)),
+      do: :ok,
+      else: {:error, :invalid_commitment_length}
+  end
+
+  defp ensure_all_proof_sizes(list) when is_list(list) do
+    if Enum.all?(list, &match?(<<_::binary-size(@bytes_per_proof)>>, &1)),
+      do: :ok,
+      else: {:error, :invalid_proof_length}
+  end
+
+  defp ensure_batch_lengths_match(blobs, commitments, proofs) do
+    case {length(blobs), length(commitments), length(proofs)} do
+      {same, same, same} -> :ok
+      _ -> {:error, :invalid_batch_length}
+    end
+  end
 end

mix.exs

@@ -67,7 +67,8 @@ defmodule LambdaEthereumConsensus.MixProject do
       {:uuid, "~> 1.1"},
       # TODO: (#1368) We might want to use phoenix_pubsub instead and do our implementation of SSE.
       {:sse, "~> 0.4"},
-      {:event_bus, ">= 1.6.0"}
+      {:event_bus, ">= 1.6.0"},
+      {:ckzg, github: "ethereum/c-kzg-4844", ref: "cfa04bb"}
     ]
   end
 

mix.lock

@@ -4,7 +4,9 @@
   "benchee": {:hex, :benchee, "1.4.0", "9f1f96a30ac80bab94faad644b39a9031d5632e517416a8ab0a6b0ac4df124ce", [:mix], [{:deep_merge, "~> 1.0", [hex: :deep_merge, repo: "hexpm", optional: false]}, {:statistex, "~> 1.0", [hex: :statistex, repo: "hexpm", optional: false]}, {:table, "~> 0.1.0", [hex: :table, repo: "hexpm", optional: true]}], "hexpm", "299cd10dd8ce51c9ea3ddb74bb150f93d25e968f93e4c1fa31698a8e4fa5d715"},
   "bunt": {:hex, :bunt, "1.0.0", "081c2c665f086849e6d57900292b3a161727ab40431219529f13c4ddcf3e7a44", [:mix], [], "hexpm", "dc5f86aa08a5f6fa6b8096f0735c4e76d54ae5c9fa2c143e5a1fc7c1cd9bb6b5"},
   "castore": {:hex, :castore, "1.0.12", "053f0e32700cbec356280c0e835df425a3be4bc1e0627b714330ad9d0f05497f", [:mix], [], "hexpm", "3dca286b2186055ba0c9449b4e95b97bf1b57b47c1f2644555879e659960c224"},
+  "cc_precompiler": {:hex, :cc_precompiler, "0.1.11", "8c844d0b9fb98a3edea067f94f616b3f6b29b959b6b3bf25fee94ffe34364768", [:mix], [{:elixir_make, "~> 0.7", [hex: :elixir_make, repo: "hexpm", optional: false]}], "hexpm", "3427232caf0835f94680e5bcf082408a70b48ad68a5f5c0b02a3bea9f3a075b9"},
   "certifi": {:hex, :certifi, "2.15.0", "0e6e882fcdaaa0a5a9f2b3db55b1394dba07e8d6d9bcad08318fb604c6839712", [:rebar3], [], "hexpm", "b147ed22ce71d72eafdad94f055165c1c182f61a2ff49df28bcc71d1d5b94a60"},
+  "ckzg": {:git, "https://github.com/ethereum/c-kzg-4844.git", "cfa04bb15610189e2969801bec79595d83ddc40a", [ref: "cfa04bb"]},
   "combine": {:hex, :combine, "0.10.0", "eff8224eeb56498a2af13011d142c5e7997a80c8f5b97c499f84c841032e429f", [:mix], [], "hexpm", "1b1dbc1790073076580d0d1d64e42eae2366583e7aecd455d1215b0d16f2451b"},
   "cowboy": {:hex, :cowboy, "2.13.0", "09d770dd5f6a22cc60c071f432cd7cb87776164527f205c5a6b0f24ff6b38990", [:make, :rebar3], [{:cowlib, ">= 2.14.0 and < 3.0.0", [hex: :cowlib, repo: "hexpm", optional: false]}, {:ranch, ">= 1.8.0 and < 3.0.0", [hex: :ranch, repo: "hexpm", optional: false]}], "hexpm", "e724d3a70995025d654c1992c7b11dbfea95205c047d86ff9bf1cda92ddc5614"},
   "cowboy_telemetry": {:hex, :cowboy_telemetry, "0.4.0", "f239f68b588efa7707abce16a84d0d2acf3a0f50571f8bb7f56a15865aae820c", [:rebar3], [{:cowboy, "~> 2.7", [hex: :cowboy, repo: "hexpm", optional: false]}, {:telemetry, "~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "7d98bac1ee4565d31b62d59f8823dfd8356a169e7fcbb83831b8a5397404c9de"},