REST APIs & Custom Authentication in RUST— Writing Extractors in Axum

Writing Custom Middleware in Rust : The right way

In this short post, we’ll look at writing some custom logic as a middleware for REST APIs using Rust And Axum Framework. We’ll start by creating a new Rust binary project and add dependencies for Axum Framework. The actual problem statement is given below:

Problem Statement

A common requirement while writing REST Apis is that we may need to extract User information from the Request Context. In a majority of the use cases, it can as simple as extracting the User Authorization Header. It may be a JWT , or it may be a simple user object that is serialized — ( if you already have an auth proxy setup. ).

In either case, we want the following functionality within our APIs:

  1. Parse Headers every time and extract come information.
  2. Throw errors in case the information is missing.
  3. Provide a User Object to the API handlers, so that they can use it.

We’ll look at these use cases in Rust and Axum framework, and write a very simple application that enables us to do everything in a nice, concise way.

Project Setup

Let’s start by setting up a new Rust Project. We’ll create a new Binary.

cargo new axum-custom-auth-demo

Add Dependencies

Add the following dependencies to your project. To write our API, we’ll use rust Axum Framework and tokio as it’s runtime. And we’ll also include JSON and Serde dependencies for Json Serialization and deserialization. s

[dependencies]  
axum = { version = "0.5.16", features=["headers"]}  
tokio = {version = "1.21.2", features=["full"]}  
serde = {version = "1.0.145", features = ["derive"]}  
serde_json = "1.0.85"

Main Function

We’ll create a very simple endpoint that’s a Hello World Get API. But to this API, we would ensure that the request is authenticated, and a user object is provided by the time the request reaches our handler. Makes things really easy. Given below is our base app setup without any authentication logic.

#[tokio::main]  
async fn main() {  
    // build our application with a route  
    let app = Router::_new_()  
        // `GET /` goes to `root`  
        .route("/", get(root));  
  
    // run our app with hyper  
    // `axum::Server` is a re-export of `hyper::Server`  
    let addr = SocketAddr::_from_(([127, 0, 0, 1], 3000));  
    axum::Server::_bind_(&addr)  
        .serve(app.into_make_service())  
        .await  
        .unwrap();  
}// basic handler that responds with a static string  
async fn root() -> &_'static_ str {  
    "Hello, World!"  
}

We define a route called root, that handles all requests matching our root URL localhost:3000

The above program runs fine and outputs Hello World when the root endpoint is called. But what we really want is an Authenticated User within this root handler. Imagine we if multiple endpoints were present, and we want to provide an authenticated user object to every single one of them.

This setup can be easily achieved by Axum. Axum is a really powerful framework and yet very simply designed. To achieve our requirement, we will use an **Extractor in axum**. As the name suggests, an extractor is a way in which we can extract values from our request, or apply certain logic on the incoming request. Think of it like a processing function applied to every incoming request. It can yield an outcome or error out, or not return anything. There can be multiple extractors defined in Axum that run before the request reaches the actual handler. Whether it’s parsing request body bytes as JSON, reading request headers or logging request bodies — every thing is an extraction process. Once the extractors have run, the result of the extractor can be passed to the handler function. Simple, Right ?

Our use case Implementation

Logically, for our use case, we just want to read our request headers, extract a x-user-id header, and create a new user object with this ID. In real world, it may be that you’re using a JWT token to verify a token, extract claims and creating a User Object. But for this post ,we’ll keep it simple. The whole idea of doing this is to achieve two things.

  1. Understand what extractors are by an example.
  2. Provide customization so that we can implement our logic within any APIs we write.

Coming back to our use case, we want to extract x-user-id header within all the requests that our Api receives. If not, then let’s output an unauthorized message to our users.

Implementing an Extractor

Before we go ahead and define a logic for extracting a User, let’s create a new User Object. This user is what will be propagated to all the handlers after authentication.

// the output to our `create_user` handler  
#[derive(Serialize)]  
struct AppUser {  
    username: String,  
}

For our use case, we’ll create a new user object with a username field. Whenever we receive a header with key x-user-id, we will create a new user object and use that within the handler. The only question remaining is how do we really do it.

The way we do this is by Implementing the FromRequest trait provided by Axum on the User Object. This way, whenever we tell axum that we need a User Object, axum will try to extract it. This can be implemented on any struct / method. The output of this trait is a Result Object with any return type, and an optional error that can have the error details.

In our case, we implemented it to return an AppUser type, and an Unauthorized error in case we didn’t find the header. We could silently create the user object and skip the failed error altogether if such a requirement exists. Below is the logic to create our user object from request headers.

The error type is any type that can be converted to a Response Object. In our case, we use a StatusCode error which Axum Provides. If you would like to use JSON errors, please refer to my other article on Rest API creation using Axum , postgres and Rust.

Rest APIs in Rust

How to Write a REST API Using Rust and Axum Framework

With Postgres Connection Pooling

betterprogramming.pub

The beauty of the above method is that the logic is really modular. Now we know that we can use our AppUser to create a new user from HttpRequest by just using this piece, no matter which handler, or which application we are working with. It’s time to update our handler and inject User into our handler object. We update our root handler like this.

Before

// basic handler that responds with a static string  
async fn root() -> &_'static_ str {  
    "Hello, World!"  
}

After

// basic handler that responds with a static string  
async fn root(authenticated_user: AppUser) -> &_'static_ str {  
    "Hello, World!"  
}

By passing an user Object to our handler, axum can infer that we have an extractor in place to get this object. For every HttpRequest, it tries to run our extractor and generate a Object for us. We can freely use this and not have to worry about authentication, creation of objects / errors. Everything is handled by one single trait implementation on AppUser class.

Test Results

Let’s test out our implementation with our fake authentication.

cargo run

And fire a call to  `localhost:3000`

$ **curl -kv localhost:3000**  
*   Trying 127.0.0.1:3000...  
* Connected to localhost (127.0.0.1) port 3000 (#0)  
> GET / HTTP/1.1  
> Host: localhost:3000  
> User-Agent: curl/7.81.0  
> Accept: */*  
>   
* Mark bundle as not supporting multiuse  
< HTTP/1.1 401 Unauthorized  
< content-type: text/plain; charset=utf-8  
< content-length: 54  
< date: Thu, 06 Oct 2022 01:50:52 GM

We now see that we get an Unauthorized Response from our server. Now try passing a header x-user-id value in the request.

curl -kv -H 'x-user-id:shanmukh' localhost:3000  
*   Trying 127.0.0.1:3000...  
* Connected to localhost (127.0.0.1) port 3000 (#0)  
> GET / HTTP/1.1  
> Host: localhost:3000  
> User-Agent: curl/7.81.0  
> Accept: */*  
> x-user-id:shanmukh  
>   
* Mark bundle as not supporting multiuse  
< HTTP/1.1 200 OK  
< content-type: text/plain; charset=utf-8  
< content-length: 13  
< date: Thu, 06 Oct 2022 01:52:33 GMT  
<   
* Connection #0 to host localhost left intactHello, World!

We see that our server now returns Hello World to us. This is exactly what we wanted. And if you debug the program, you’ll notice that your user object is available to you for use within your handler method.

Conclusion

What you saw was a very simple, yet powerful way of parsing requests, and using extractors to play with requests before it reaches to our handlers. It’s a clean way to apply our logic and some middleware without having to pollute actual handlers. Instead it provides ways to injectdependencies of to our handlers with module logic.

By default Axum does provide a good amount of Extractors out of the box. Please check out those before you write your own. This example is a good balance of framework capability and some customization we can do to solve our requirements while writing Rest APIs in Rust.