Writing Rust NIF code to convert encryption code from Elixir to Rust

In this article, we will go through what a NIF is, how to write safe NIF code using Rust and Rustler

What's a NIF?

A NIF of Native Implemented Function, is a  function usually implemented in C, which can be called from Elixir. NIFs are usually used to run a small piece of native code for faster performance.

Writing a NIF using Rust and Rustler

From Rustler's documentation:

Rustler is a library for writing Erlang NIFs in safe Rust code. That means there should be no ways to crash the BEAM (Erlang VM). The library provides facilities for generating the boilerplate for interacting with the BEAM, handles encoding and decoding of Erlang terms, and catches rust panics before they unwind into C.

One of the big caveats of using NIF in an Elixir code base, is that it could potentially bring down the entire BEAM VM, if the NIF code panic. Thus writing NIF code in Rust gives the advantage that Rust code can catch any panic before.

Getting started with Rustler

The first step is adding rustler as a dependency to the project.

{:rustler, "~> 0.26.0"}

and running mix deps.get

Once the dependencies are installed, a new rustler project can be created by running:

mix rustler.new

Writing your first NIF

The following code, takes an encrypted string, decrypts it using AES-256-GCM mode.

defmodule Encryption do
  @moduledoc """
  Decrypts a given encrypted string using AES-256-GCM mode decryption.
  """
  @key_size 32
  @iv_size 16
  @mode :aes_256_gcm
  @tag_length 16

  @doc """
  Decrypts a given encrypted string using AES-256-GCM mode decryption.
  """
  @spec decrypt(binary) :: binary
  def decrypt(content) do
    data_bin = Base.url_decode64!(content)
    size = byte_size(data_bin)
    data_size = size - @key_size - 2

    <<_version::binary-size(2), data::binary-size(data_size), key::binary-size(@key_size)>> = data_bin
    <<iv::binary-size(@iv_size), tag::binary-size(@tag_length), aad::binary-size(@iv_size), cipher::binary>> = payload

    :crypto.crypto_one_time_aead(@mode, key, iv, cipher, aad, tag, false)
  end
end

Now, lets take the above code and convert that to a NIF using Rust.

Let's start with writing an elixir module

defmodule NIFEncryption do
  @moduledoc """
  NIF module decrypting data.
  """
  use Rustler, otp_app: :otp_app_name, crate: "encryption"

  @doc """
  Decodes an encrypted token.
  """
  @spec decrypt(binary()) :: {:ok, binary()} | {:error, binary()}
  def decrypt(_token), do: error()

  defp error, do: :erlang.nif_error(:nif_not_loaded)
end

The code will allow us to utilise the Rust crate encryption

Now, lets write the Rust code to implement our decrypt function.

// file: native/src/encryption/lib.rs

use openssl::symm::Cipher;

const KEY_SIZE: usize = 32;
const IV_SIZE: usize = 16;
const TAG_LENGTH: usize = 16;

#[rustler::nif]
pub fn decrypt(token: &str) -> String {
    let data_bin: Vec<u8> = base64_url::decode(token).unwrap();
    let size: usize = data_bin.len();
    let (iv, tag_and_aad): (&[u8], &[u8]) = data.split_at(IV_SIZE);
    let (tag, aad_and_cipher): (&[u8], &[u8]) = tag_and_aad.split_at(TAG_LENGTH);
    let (aad, cipher): (&[u8], &[u8]) = aad_and_cipher.split_at(IV_SIZE);

    let content: Vec<u8> =
        openssl::symm::decrypt_aead(Cipher::aes_256_gcm(), key, Some(iv), aad, cipher, tag)
            .unwrap();

    return content.iter().map(|e| *e as char).collect::<String>();
}

rustler::init!("Elixir.Encryption", [decode]);

the [rustler::nif] macro exposes the decrypt function as a nif, and we initialize the NIF using rustler::init! macro which exports the function as  Elixir.Encryption.decrypt/1

Benchmarks

While its cool to port over the Elixir code to Rust code, its also important to see the performance comparison of the Elixir implementation over the Rust implementation. Most often synthetic benchmarks does not bring in a lot of value add to the code, but it gives good insights into some of the stats.

Running a synthetic benchmark using benchee gives the following results.

TL;DR version

Elixir implementation is

• 9.51x slower
• uses 2.05x more memory

Operating System: macOS
CPU Information: Apple M1 Max
Number of Available Cores: 10
Available memory: 64 GB
Elixir 1.14.2
Erlang 25.2

Benchmark suite executing with the following configuration:
warmup: 2 s
time: 10 s
memory time: 2 s
reduction time: 0 ns
parallel: 10
inputs: none specified
Estimated total run time: 28 s

Benchmarking Elixir ...
Benchmarking Rust ...

Name             ips        average  deviation         median         99th %
Rust         22.77 K       43.93 μs   ±223.31%       35.33 μs      173.80 μs
Elixir        2.39 K      417.70 μs    ±90.41%      327.21 μs     2004.40 μs

Comparison: 
Rust         22.77 K
Elixir        2.39 K - 9.51x slower +373.78 μs

Memory usage statistics:

Name      Memory usage
Rust           0.59 KB
Elixir         1.20 KB - 2.05x memory usage +0.62 KB

Closing thoughts

While NIFs are cool and often gives better performance, there are caveats to consider.

• A panic in NIF can bring down the entire BEAM VM.
• Learning curve to implement native code in C or Rust.
• Maintenance overhead of adding another language to the stack.
• NIF is generically recommended to be used for short running operations - usually under 1s. If the operation takes more time, look into dirty schedulers.

Hope this helps someone.