Elixir simple seeded AES-CTR rng

simple_rng.ex

defmodule Foo.SimpleRNG do
  def new(), do: new(:crypto.strong_rand_bytes(16), 0)
  def new(<<seed::binary-size(16)>>), do: new(:aes_128_ctr, seed, 0)
  def new(<<seed::binary-size(32)>>), do: new(:aes_256_ctr, seed, 0)
  def new(cipher, <<seed::binary>>) when is_atom(cipher), do: new(cipher, seed, 0)
  def new(<<seed::binary-size(16)>>, ctr) when is_integer(ctr), do: new(:aes_128_ctr, seed, ctr)
  def new(<<seed::binary-size(32)>>, ctr) when is_integer(ctr), do: new(:aes_256_ctr, seed, ctr)

  @spec new(cipher :: :crypto.cipher_iv, crypto_seed :: binary, ctr_start :: non_neg_integer) :: :rand.state
  def new(cipher, seed, ctr)

  def new(:aes_128_ctr, crypto_seed, ctr) do
    {impl(:aes_128_ctr, block_size: 128), {crypto_seed, ctr}}
  end

  def new(:aes_256_ctr, crypto_seed, ctr) do
    {impl(:aes_256_ctr, block_size: 128), {crypto_seed, ctr}}
  end

  def new(:chacha20, crypto_seed, ctr) do
    {impl(:chacha20, block_size: 512, iv_size: 128, jump_size: 32), {crypto_seed, ctr}}
  end

  def new(:sm4_ctr, crypto_seed, ctr) do
    {impl(:sm4_ctr, block_size: 128), {crypto_seed, ctr}}
  end

  defp impl(cipher, options) do
    block_size = Keyword.fetch!(options, :block_size)
    iv_size = Keyword.get(options, :iv_size) || block_size
    jump_size = 2 ** (Keyword.get(options, :jump_size) || div(iv_size, 2))
    null_plaintext = <<0::size(block_size)>>

    %{
      type: cipher,
      bits: block_size,
      next: fn {crypto_seed, ctr} ->
        case :crypto.crypto_one_time(cipher, crypto_seed, <<ctr::size(iv_size)>>, null_plaintext, encrypt: true) do
          <<result::binary>> ->
            {:binary.decode_unsigned(result, :big), {crypto_seed, ctr + 1}}
        end
      end,
      jump: fn {crypto_seed, ctr} -> {crypto_seed, ctr + jump_size} end
    }
  end
end