Martin PAUCOT

Posted on Oct 28, 2023

Roll your own auth with Rust and Protobuf

Introduction

Before starting, take a deep breath, maybe 5 minutes of meditation and prepare a drink because this tutorial imply Rust code. Not complex, but Rust code.

But first, let's talk about Protobuf, what is it?

1. What is Protobuf?

According to the documentation:

"Protocol Buffers are language-neutral, platform-neutral extensible mechanisms for serializing structured data."

JSON is really flexible when you want to share data through services you can decode JSON without knowing it's structure first.
But it is unstructured, takes a lot of space and bandwith.

With Protocol Buffers you define a message and its structure that must be known by both the server and the client to encode and decode.

// user.proto

syntax = "proto3";

message User {
    string firstname = 1;
    string lastname = 2;
    string email = 3;
}

You can then use generators to create the SDKs for your favorites languages. You can generate a Javascript one for your frontend and a Rust one for your backend.

If you are using Remote Procedure Call (RPC) like gRPC, you can leverage the features of Protobuf and the generators to automatically generate interfaces and code for both your client and server SDKs.
The only thing you have to do next, is implement the methods of your services.

// user.proto

syntax = "proto3";

message LoginRequest {
    string email = 1;
    string password = 2;
}

service Auth {
    rpc Login (LoginRequest) returns (User);
}

2. What are you we to accomplish

Create a PostgreSQL database using Docker and Docker compose
Create Protobuf definitions and use Buf to generate the SDKs
Setup a gRPC server in Rust using Tonic
Create a JWT authentication system with Diesel and Tonic

3. Requirements

You must have a basic understanding of Rust, I will not deep dive into how to write Rust. As I am myself, a Rust beginner. But I encourage you to check the different crates documentation to have a better understanding on how everything works.

You must know how JWT works as I will not explain it in this blog post.

You must have all the required tools installed:

If you are lost somewhere, or you directly want to go straigth to the code, the repository is available here https://github.com/kerwanp/rust-proto-demo.

Get started

1. Create the PostgreSQL database

As we want to persist our users we need a database. We are going to use PostgreSQL.

For simplicity of use, we will use Docker Compose to manage it. Make sure you have Docker installed before.

In your root folder add the following docker-compose.yml file:

# docker-compose.yaml

version: "3"

services:
  db:
    image: postgres
    restart: always
    ports:
      - "5432:5432"
    environment:
      POSTGRES_USER: rustproto
      POSTGRES_PASSWORD: rustproto

You can then run the command docker compose up -d to start your database (-d will start it in background mode).

2. Create the Protobuf definitions

Protobuf definitions allows 2 things:

Define how gRPC server and client communicate together
Generate SDKs using Protobuf generators

We will store them in a folder called proto

Creating the Auth service Protobuf definitions

We will need two methods in our service:

Register: To create a new user. It returns a Token.
Login: To generate a new Token if credentials are valid.

// proto/auth.proto

syntax = "proto3";

package auth;

message LoginRequest {
    string username = 1;
    string password = 2;
}

message RegisterRequest {
    string firstname = 1;
    string lastname = 2;
    string email = 3;
    string password = 4;
}

message Token {
    string access_token = 1;
}

service Auth {
    rpc Login (LoginRequest) returns (Token);
    rpc Register (RegisterRequest) returns (Token);
}

Creating the Greeting service definitions

To test that our authentication works we need a random service that will throw an unauthenticated status if the access token is invalid.

// proto/greeting.proto

syntax = "proto3";

package greeting;

message GreetRequest {
    string message = 1;
}

message GreetResponse {
    string message = 1;
}

service Greeting {
    rpc Greet (GreetRequest) returns (GreetResponse);
}

3. Setup Buf to generate SDKs

To generate SDKs we have different solutions:

Use tonic-build directly from Rust.
Use the Protobuf CLI protoc and the plugin protoc-gen-tonic.
Use Buf CLI to manage the SDKs generation.

The first two solutions would fit perfectly but as we want to overengineer our authentication system and think of the not coming future we will use the Buf CLI.

So make sure to install it.

Create the module configuration

Buf CLI works with workspaces and modules to easily split our protobuf definitions (by APIs, microservices, etc).

Let's make our proto folder our lonely module by creating a buf.yaml file with the following line:

# proto/buf.yaml

version: v1

Create the workspace configuration

Now we need to configure our Buf workspace by creating a buf.work.yaml at the root of our project:

# buf.work.yaml

version: v1
directories:
  - proto # <- We define our module in the workspace

Create the generator configuration

We will use four generators:

protoc-gen-prost The core code generation plugin
protoc-gen-serde Canonical JSON serialization of protobuf types
protoc-gen-tonic gRPC service generation for the Tonic framework
protoc-gen-prost-crate Generates an include file and cargo manifest for turn-key crates

Let's install them all using the following command:

$ cargo install protoc-gen-prost protoc-gen-prost-serde protoc-gen-tonic protoc-gen-prost-crate

Then, create the buf.gen.yaml file to configure the SDKs generation:

# buf.gen.yaml

version: v1
plugins:
  - plugin: prost # Generates the core code
    out: gen/src
    opt:
      - bytes=.
      - compile_well_known_types
      - extern_path=.google.protobuf=::pbjson_types
      - file_descriptor_set
  - plugin: prost-serde # Generates code compatible with JSON serde
    out: gen/src
  - plugin: tonic # Generates the Tonic services
    out: gen/src
    opt:
      - compile_well_known_types
      - extern_path=.google.protobuf=::pbjson_types
  - plugin: prost-crate # Makes the gen folder a crate
    out: gen
    opt:
      - gen_crate=gen/Cargo.toml

We now need to create the Cargo.toml that will be used as a template to generate the new one.

For simplicity we will use the generated one as a template so it can be replaced.

# gen/Cargo.toml

[package]
name = "protos"
version = "0.1.0"
edition = "2021"

[features]
default = ["proto_full"]
# @@protoc_deletion_point(features)
# This section is automatically generated by protoc-gen-prost-crate.
# Changes in this area may be lost on regeneration.
# @@protoc_insertion_point(features)

[dependencies]
bytes = "1.1.0"
prost = "0.12"
pbjson = "0.6"
pbjson-types = "0.6"
serde = "1.0"
tonic = { version = "0.10", features = ["gzip"] }

We can now generate our crate using the command buf generate.

🎉 Our SDK is now ready!

4. Setup the Rust project

Now that we have our SDK we can start creating the core of our server.
First create a Cargo.toml file and add our generate crate to the dependencies aswell as all the requirements for this project.

[package]
name = "proto-auth-demo"
version = "0.1.0"
edition = "2021"

[dependencies]
diesel = { version = "2.1.3", features = ["postgres"] }
dotenvy = "0.15.7"
prost = "0.12.1"
protobuf = "3.3.0"
serde = "1.0.190"
tokio = { version = "1.33.0", features = ["full"] }
tonic = "0.10.2"
protos = { path = "./gen" }
bcrypt = "0.15.0"
jwt = "0.16.0"
hmac = "0.12.1"
sha2 = "0.10.8"

[workspace]
members = [
    "gen"
]

We will need two environment variables:

DATABASE_URL: To authenticate to our PostgreSQL database
APP_KEY: To encrypt user passwords

Store them in a .env file, we will use them later using the dotenvy crate.

# .env

DATABASE_URL=postgres://rustproto:rustproto@localhost/rustproto
APP_KEY="9E3CnfSfsi9BGfX3Dea#tkbs#nDj&6d#6Y&jhNa!"

And we can now create our main.rs file:

// src/main.rs

use dotenvy::dotenv;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    dotenv().ok();
    Ok(())
}

And build our project using cargo build.

5. Use Diesel to manage our users

Setup Diesel

First install the Diesel CLI using the following command:

$ cargo install diesel_cli

You may need to install libpq-dev, libmysqlclient-dev and libsqlite3-dev to install the CLI.

Create the migrations folder:

$ diesel setup

And add the following lines to the top of your main.rs (we will create those modules later):

// src/main.rs

mod models;
mod schema;

Create the migration

Let's first generate the migration to create and drop the users table.

$ diesel migration generate create_users

-- migrations/*-create_users/up.sql
CREATE TABLE users (
    id SERIAL PRIMARY KEY,
    email VARCHAR(255) UNIQUE NOT NULL,
    password VARCHAR(255) NOT NULL,
    firstname VARCHAR(255) NOT NULL,
    lastname VARCHAR(255) NOT NULL
);

-- migrations/*-create_users/down.sql
DROP TABLE users;

And run the following command to run the migrations and generate the schema.rs file:

$ diesel migration run

Create the User model

Create a new file src/models.rs and add our User model to it and implements the methods to create a new User and find one by email.

// src/models.rs

use diesel::{
    ExpressionMethods, Insertable, PgConnection, QueryDsl, Queryable, RunQueryDsl, Selectable,
    SelectableHelper,
};

use crate::schema::users;

#[derive(Queryable, Selectable)]
#[diesel(table_name = users)]
#[diesel(check_for_backend(diesel::pg::Pg))]
pub struct User {
    pub id: i32,
    pub firstname: String,
    pub lastname: String,
    pub email: String,
    pub password: String,
}

#[derive(Insertable)]
#[diesel(table_name = users)]
pub struct NewUser<'a> {
    pub firstname: &'a str,
    pub lastname: &'a str,
    pub email: &'a str,
    pub password: &'a str,
}

impl User {
    pub fn create(
        conn: &mut PgConnection,
        firstname: &str,
        lastname: &str,
        email: &str,
        password: &str,
    ) -> Result<User, diesel::result::Error> {
        let new_user = NewUser {
            firstname,
            lastname,
            email,
            password,
        };

        diesel::insert_into(users::table)
            .values(new_user)
            .returning(User::as_returning())
            .get_result(conn)
    }

    pub fn find_by_email(conn: &mut PgConnection, email: &str) -> Option<User> {
        users::dsl::users
            .select(User::as_select())
            .filter(users::dsl::email.eq(email))
            .first(conn)
            .ok()
    }
}

Create the database connection

We now have to setup the database connection and we will be done with Diesel.

// src/main.rs

use std::env;

use diesel::{PgConnection, Connection};
use dotenvy::dotenv;

mod models;
mod schema;

pub fn connect_db() -> PgConnection {
    let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
    PgConnection::establish(&database_url)
        .unwrap_or_else(|_| panic!("Error connecting to {}", database_url))
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    dotenv().ok();

    let mut database = connect_db();

    Ok(())
}

You can try it by creating a user using the database connection User.create(&mut database, "", "", "", "")

5. Implement the Tokio Services

It is now time to create the business logic for our authentication system.
We will start with the biggest part, the authentication.

Create the skeleton

use std::sync::{Arc, Mutex};

use diesel::PgConnection;
use protos::auth::{auth_server::Auth, LoginRequest, Token, RegisterRequest};
use tonic::{Request, Response, Status};

use crate::models::User;

pub struct Service {
    database: Arc<Mutex<PgConnection>>,
}

impl Service {
    pub fn new(database: PgConnection) -> Self {
        Self {
            database: Arc::new(Mutex::new(database))
        }
    }

    fn generate_token(user: User) -> Token {
        unimplemented!();
    }
}

#[tonic::async_trait]
impl Auth for Service {
    async fn login(&self, request: Request<LoginRequest>) -> Result<Response<Token>, Status> {
        unimplemented!();
    }

    async fn register(&self, request: Request<RegisterRequest>) -> Result<Response<Token>, Status> {
        unimplemented!();
    }
}

Start the Tonic Server and add our unimplemented service

In the main.rs file we are going to create a server, add our service to it and start it.

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    dotenv().ok();

    let database = connect_db();
    let addr = "[::1]:50051".parse()?;

    Server::builder()
        .add_service(AuthServer::new(auth::Service::new(database)))
        .serve(addr)
        .await?;

    Ok(())
}

You can try your server by using the following command to call the auth.Auth/Login function (package.Service/Method).

$ grpcurl -plaintext -import-path ./proto -proto auth.proto '[::1]:50051' auth.Auth/Login

You should see in the console of your server the following error:

thread 'tokio-runtime-worker' panicked at 'not implemented', src/auth.rs:22:9

This error has been thrown by the unimplemented! macro of our Auth Service, it works!

Generate an access token

We are going to use the jwt crate to create a JWT token. It will contains the following claims:

sub: The Subject of the JWT (our user id)
iat: Time at which the JWT was issued
exp: Time after which the JWT expires

Implement the register method

The register method is fairly simple, it directly takes the request message to create the entry in our database after encrypting the password.

// src/auth.rs

async fn register(&self, request: Request<RegisterRequest>) -> Result<Response<Token>, Status> {
    let database = self.database.lock();
    let data = request.into_inner();

    let password = bcrypt::hash(&data.password, 10)
        .map_err(|_| Status::unknown("Error while creating the user"))?;

    let user = NewUser {
        firstname: &data.firstname,
        lastname: &data.lastname,
        email: &data.email,
        password: &password,
    };

    let user = User::create(&mut database.unwrap(), user);

    unimplemented!();
}

You can already try it! It will throw an error because of the unimplemented but your user should be created in the database.

Implement the generate token method

We are going to split our function in two, the part for generating the claims and the one for encoding the token.

We are going to use the APP_KEY defined in our .env file to sign the JWT.

// src/auth.rs
pub struct GenerateTokenError;

pub struct GenerateClaimsError;

fn generate_claims(user: User) -> Result<BTreeMap<&'static str, String>, GenerateClaimsError> {
    let mut claims: BTreeMap<&str, String> = BTreeMap::new();

    claims.insert("sub", user.id.to_string());

    let current_timestamp = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_err(|_| GenerateClaimsError)?
        .as_secs();

    claims.insert("iat", current_timestamp.to_string());
    claims.insert("exp", String::from("3600"));

    Ok(claims)
}

fn generate_token(user: User) -> Result<Token, GenerateTokenError> {
    let app_key: String = env::var("APP_KEY").expect("env APP_KEY is not defined");

    let key: Hmac<Sha256> =
        Hmac::new_from_slice(app_key.as_bytes()).map_err(|_| GenerateTokenError)?;

    let claims = generate_claims(user).map_err(|_| GenerateTokenError)?;

    let access_token = claims.sign_with_key(&key).map_err(|_| GenerateTokenError)?;

    Ok(Token {
        access_token: access_token,
    })
}

And now, let's use this function to return an access token when a user is registered:

// src/auth.Rs
impl Auth for Service {
    [...]
    async fn register(&self, request: Request<RegisterRequest>) -> Result<Response<Token>, Status> {
        let database = self.database.lock();
        let data = request.into_inner();

        let password = bcrypt::hash(&data.password, 10)
            .map_err(|_| Status::unknown("Error while creating the user"))?;

        let user = NewUser {
            firstname: &data.firstname,
            lastname: &data.lastname,
            email: &data.email,
            password: &password,
        };

        let user = User::create(&mut database.unwrap(), user)
            .map_err(|_| Status::already_exists("User already exists in the database"))?;

        let token = generate_token(user).map_err(|_| Status::unknown("Cannot generate a token for the User"))?;

        Ok(Response::new(token))
    }
}

You can create a new user with the following command:

$ grpcurl -plaintext -import-path ./proto -proto auth.proto -d '{"firstname": "John", "lastname": "Doe", "email": "john@doe.com", "password": "rustproto"}' '[::1]:50051' auth.Auth/Register

And if you run it again, you will have the error AlreadyExists. 🎉

Implement the login method

The login method is now really simple, we try to find a user correspond to the email in the message, verify if the password is correct and use the generate_token method to return the Response.

// src/auth.rs
impl Auth for Service {
    [...]
    async fn login(&self, request: Request<LoginRequest>) -> Result<Response<Token>, Status> {
        let data = request.into_inner();

        let database = self.database.lock();

        let user = User::find_by_email(&mut database.unwrap(), &data.email)
            .ok_or(Status::unauthenticated("Invalid email or password"))?;

        match verify(data.password, &user.password) {
            Ok(true) => (),
            Ok(false) | Err(_) => return Err(Status::unauthenticated("Invalid email or password")),
        };

        let reply = generate_token(user)
            .map_err(|_| Status::unauthenticated("Invalid email or password"))?;

        Ok(Response::new(reply))
    }
    [...]
}

And we can now generate a token using the email and password of our previously registered user!

$ grpcurl -plaintext -import-path ./proto -proto auth.proto -d '{"email": "john@doe.com", "password": "rustproto"}' '[::1]:50051' auth.Auth/Login

Implement the verify token method

As this post already starts to be really long, we will make the verification simple, we will only check for the signature and return true if it is valid.

// src/auth.rs
pub struct VerifyTokenError;

pub fn verify_token(token: &str) -> Result<bool, VerifyTokenError> {
    let app_key: String = env::var("APP_KEY").expect("env APP_KEY is not defined");

    let key: Hmac<Sha256> =
        Hmac::new_from_slice(app_key.as_bytes()).map_err(|_| VerifyTokenError)?;

    Ok(token
        .verify_with_key(&key)
        .map(|_: HashMap<String, String>| true)
        .unwrap_or(false))
}

Implement the Greeting service

In a new src/greeting.rs file we are going to implement our Greeting service using the trait provided by our generated SDk.

We will verify that we have a x-authorization token in the metadata and verify its value.

// src/greeting.rs

use protos::greeting::{greeting_server::Greeting, GreetRequest, GreetResponse};
use tonic::{Request, Response, Status};

use crate::auth;

#[derive(Default)]
pub struct Service {}

#[tonic::async_trait]
impl Greeting for Service {
    async fn greet(
        &self,
        request: Request<GreetRequest>,
    ) -> Result<Response<GreetResponse>, Status> {
        let token = request
            .metadata()
            .get("x-authorization")
            .ok_or(Status::unauthenticated("No access token specified"))?
            .to_str()
            .map_err(|_| Status::unauthenticated("No access token specified"))?;

        match auth::verify_token(token) {
            Ok(true) => (),
            Err(_) | Ok(false) => return Err(Status::unauthenticated("Invalid token")),
        }

        let data = request.into_inner();

        Ok(Response::new(GreetResponse {
            message: format!("{} {}", data.message, "Pong!"),
        }))
    }
}

We can now add our service to the Tonic server in main.rs.

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    dotenv().ok();

    let database = connect_db();
    let addr = "[::1]:50051".parse()?;

    Server::builder()
        .add_service(AuthServer::new(auth::Service::new(database)))
        .add_service(GreetingServer::new(greeting::Service::default()))
        .serve(addr)
        .await?;

    Ok(())
}

It's now time to try our authentication system! You can use the following command and replace the access token with one generated with the Login or Register method:

$ grpcurl -plaintext -import-path ./proto -proto greeting.proto -H 'x-authorization: <access_token>' -d '{"message": "Ping!" }' '[::1]:50051' greeting.Greeting/Greet

Try to put a wrong access token, you will get a unauthenticated status!

You made it 🚀

Rust and Protobuf are new to me, but it is exactly the kind of project that helps me having a great learning experience on things that have a steep learning curve.

You can find the repository here https://github.com/kerwanp/rust-proto-demo

In the next blog post we will learn how to consume this API in NextJS, so make sure to stay tuned by following me on [dev.to]

DEV Community