ex_bert

defmodule BERT do
  @moduledoc """
  Elixir wrapper allowing safe use of binary erlang terms.
  """

  @doc """
  Encodes input into BERT.

  Returns {:ok, binary} on success or {:error, reason} on failure.

  ## Options

      :compression - sets the compression level of the generated
      BERT. Possible values are integers in the range of 0..9, where
      0 (the default) is no compression and 9 is maximum compression.
  """
  @spec encode(any) :: {:ok, binary} | {:error, any}
  @spec encode(any, compression: 0..9) :: {:ok, binary} | {:error, any}
  def encode(term, opts \\ []) do
    {:ok, encode!(term, opts)}
  rescue
    e in ArgumentError -> {:error, e}
  end

  @doc """
  Encodes input into BERT.

  Returns the BERT binary on success, raises on failure.

  ## Options

      :compression - sets the compression level of the generated
      BERT. Possible values are integers in the range of 0..9, where
      0 (the default) is no compression and 9 is maximum compression.
  """
  @spec encode!(any) :: binary
  @spec encode!(any, compression: 0..9) :: binary
  def encode!(term, opts \\ []) do
    term
    |> :erlang.term_to_binary(compressed: Keyword.get(opts, :compression, 0), minor_version: 1)
  end

  @doc """
  Decodes input from BERT.

  Returns {:ok, term} on success, {:error, reason} on failure.

  ## Options

      :safe_atoms - true (default) will prevent the creation of atoms
      which don't already exist. This should always be enabled when
      decoding untrusted data.

      :safe_values - true (default) will prevent the creation of functions,
      pid's, ports, and references. This should always be enabled when
      decoding untrusted data.

  """
  @spec decode(binary) :: {:ok, any} | {:error, any}
  @spec decode(binary, safe_atoms: boolean, safe_values: boolean) :: {:ok, any} | {:error, any}
  def decode(binary, opts \\ []) do
    {:ok, decode!(binary, opts)}
  rescue
    e in ArgumentError -> {:error, e}
  end

  @doc """
  Decodes input from BERT.

  Returns term on on success, raises on failure.

  ## Options

      :safe_atoms - true (default) will prevent the creation of atoms
      which don't already exist. This should always be enabled when
      decoding untrusted data.

      :safe_values - true (default) will prevent the creation of functions,
      pid's, ports, and references. This should always be enabled when
      decoding untrusted data.

  """
  @spec decode!(binary) :: any
  @spec decode!(binary, safe_atoms: boolean, safe_values: boolean) :: any
  def decode!(binary, opts \\ []) do
    binary
    |> bert_decode(
      Keyword.get(opts, :safe_atoms, true),
      Keyword.get(opts, :safe_values, true)
    )
  end

  @spec bert_decode(binary, false, false) :: any
  defp bert_decode(binary, false, false) do
    binary
    |> :erlang.binary_to_term()
  end

  @spec bert_decode(binary, true, false) :: any
  defp bert_decode(binary, true, false) do
    binary
    |> :erlang.binary_to_term([:safe])
  end

  @spec bert_decode(binary, false, true) :: atom | bitstring | list | number | tuple | map
  defp bert_decode(binary, false, true) do
    binary
    |> safe_bert_decode()
  end

  @spec bert_decode(binary, true, true) :: atom | bitstring | list | number | tuple | map
  defp bert_decode(binary, true, true) do
    binary
    |> safe_bert_decode([:safe])
  end

  @typep safe :: atom | bitstring | number | tuple | map
  @typep safe_list :: safe | [safe]
  @typep safe_bert :: safe_list | [safe_list]

  @spec safe_bert_decode(binary, [:safe]) :: safe_bert
  defp safe_bert_decode(binary, opts \\ []) do
    binary
    |> :erlang.binary_to_term(opts)
    |> safe_terms()
  end

  @spec safe_terms(x) :: x when x: list
  defp safe_terms([h | t]) do
    [safe_terms(h) | safe_terms(t)]
  end

  defp safe_terms([]) do
    []
  end

  @spec safe_terms(x) :: x when x: tuple
  defp safe_terms(tuple) when is_tuple(tuple) do
    tuple
    |> Tuple.to_list()
    |> Enum.map(&safe_terms(&1))
    |> List.to_tuple()
  end

  @spec safe_terms(x) :: x when x: map
  defp safe_terms(map) when is_map(map) do
    map
    |> Map.to_list()
    |> safe_terms()
    |> Map.new()
  end

  @spec safe_terms(x) :: x when x: atom | number | bitstring
  defp safe_terms(value)
       when is_atom(value) or is_number(value) or is_bitstring(value) do
    value
  end

  @spec safe_terms(fun | port | pid | reference) :: none
  defp safe_terms(_unsafe) do
    raise ArgumentError, "binary contains unsafe terms"
  end
end