Integrating OAuth into your Rust App

More and more developers these days are turning to modern, compiled, and strongly-typed languages like Rust and Go to develop web services. For Rust, there are great frameworks like Actix, Axum, Yew, and Dioxus which can be used to make reliable, robust, and fast web applications. These tools are outstanding, but when it comes to deploying your app, they will only get you so far.

For years, ngrok has been a powerful tool to quickly get local web applications on the internet with all of today’s necessities: TLS, identity and access management, load balancing, and much more. In the past, this has required running a separate process to connect a local port to the ngrok platform. But today, we’re going to take this one step further by building the ngrok tunnel straight into our app.

A few weeks ago, I stumbled upon Shub Argha’s blog post about building a chat app and using ngrok to integrate OAuth for authenticating and identifying users. I had used ngrok in the past but had no idea it could be used to handle authentication. This saves a considerable amount of work on our side so we can just focus on building cool apps. So why not try rewriting this app in Rust to check out ngrok’s new Rust Crate and also learn how websockets work along the way?

Getting Started with axum

We’re going to build upon Shub’s example but with a Rusty twist. We will use axum, tokio, and a few other handy Crates to build a websocket-powered chat app which uses OAuth to support authentication for multiple users.

Create an app with cargo:

cargo new rust-chat

And add the following dependencies:

• axum: The web framework
• tokio: Async runtime
• askama and askama_axum: Templating (similar to Jinja2)
• anyhow: Easy error-handling
• log and env_logger: You guessed it… logging
• futures: For handling incoming and outgoing websocket messages simultaneously

App State

The first thing we need for our app is some State to keep track of things.

pub struct AppState {
    chat_history: Mutex<Vec<ChatMsg>>,
    pub tx: broadcast::Sender<String>,
}

chat_history is a Vec of ChatMsg (we’ll define that in a minute). But we wrap this in a Mutex because we’re going to need to write to this Vec elsewhere. We need this so when new clients connect, they can see the previous chat history.

tx is the sending half of an MPMC (multi-producer, multi-consumer) channel which will be used to send messages to the websocket.

And the messages that will be sent are:

#[derive(Debug, Clone)]
pub struct ChatMsg {
    pub username: String,
    pub text: String,
}

A simple struct with both a username and a text field.

Finally, we’re going to want to implement a few methods on our state struct:

impl AppState {
    pub fn new() -> Self {
        let chat_history = Mutex::new(Vec::new());
        let (tx, _rx) = broadcast::channel(100);
        Self { chat_history, tx }
    }

    pub fn get_history(&self) -> Vec {
        self.chat_history.lock().unwrap().clone()
    }

    pub fn push_history(&self, msg: ChatMsg) {
        let mut history = self.chat_history.lock().unwrap();
        history.push(msg);
        history.reverse();
        history.truncate(100);
        history.reverse();
    }
}

First, a new method for creating a new AppState and a getter and setter for handling adding items to the chat history. We want to limit the size of our chat history so we use the .truncate method sandwiched between two .reverse calls so we’re keeping the most recent messages. This isn’t pretty but it works.

Main

Our entry point will look similar to most other axum apps.

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    env_logger::init_from_env(env_logger::Env::new().default_filter_or("info"));

    let app_state = Arc::new(AppState::new());

    let app = Router::new()
        .route("/", get(routes::index))
        .route("/ws", get(routes::websocket))
        .with_state(app_state);

    let addr = SocketAddr::from(([127, 0, 0, 1], 8000));
    log::info!("App started on URL: {:?}", addr.to_string());

    axum::Server::bind(&addr)
        .serve(app.into_make_service_with_connect_info::())
        .await?;

    Ok(())
}

Our app has only two endpoints: / for serving up the static HTML content and /ws for handling the websocket connections.

Index Route

The index route does really only one thing: renders HTML. We’re going to shamelessly rip off the HTML template (here) from Shub and make some minor tweaks.

This template will work pretty well with our template crate askama but we will be passing messages in just a bit differently. So change this:

{% for message in chat_messages %}
<div class="message">
  	<span class="name"><b>{{ message[0] }}</b>:</span>
  	<span class="text">{{ message[1] }}</span>
</div>
{% endfor %}

To this:

{% for message in chat_messages %}
<div class="message">
  <span class="name"><b>{{ message.username }}</b>:</span>
  <span class="text">{{ message.text }}</span>
</div>
{% endfor %}

Which will represent each of the ChatMsgs we pass in to the template from our AppState.chat_history field. Put the modified template file in templates/index.html in the base directory of your app.

Finally, we render our HTML template and serve it up from our index method.

#[derive(Template)]
#[template(path = "./index.html")]
pub struct IndexTemplate {
    pub chat_messages: Vec<ChatMsg>,
}

pub async fn index(State(state): State<Arc<AppState>>) -> impl IntoResponse {
    let chat_messages = state.get_history();
    IndexTemplate { chat_messages }
}

Websocket Route

The websocket route handles the incoming connection and starts the chat service:

pub async fn websocket(
    ws: WebSocketUpgrade,
    State(state): State<Arc<AppState>>,
) -> impl IntoResponse {
    let username = "None".to_string();
    ws.on_upgrade(|socket| chat_service(socket, state, username))
}

For now, we’re going to set a username of “None”. We’ll fix this later.

The chat_service function is a bit more complicated and we don’t need to go into the details of it here. But the source code I’m using can be found here.

Run It

Now, if we run the app with cargo run, we will get a simple chat app that has minimal functionality running on port 8000. If we wanted to get this app online, we’d have a significant amount of work ahead of us still. We’d need to configure a reverse-proxy and open a port on our router. We’d also want to set up TLS encryption which would require a domain name and a valid x509 certificate. And that’s the easy part. The next step would be to configure some kind of authentication mechanism and there are a myriad of options here. Most of them will require us to add a substantial amount of application logic. Monitoring and load balancing would be nice to have. Where to even start? And worse yet, none of this is related to chat.

Perhaps there’s an easier way…

Ngrok It

With ngrok’s Rust Crate, we can get everything listed above in only a few lines of code.

First we add the ngrok crate to our project. I like to use cargo-edit:

cargo add ngrok

Then we configure the tunnel in our main function:

let tunnel = ngrok::Session::builder()
        .authtoken_from_env()
        .connect()
        .await?
        .http_endpoint()
        .oauth(OauthOptions::new("google"))
        .listen()
        .await?;

This grabs our ngrok API token from the environment ($NGROK_AUTHTOKEN), establishes an http tunnel, and protects it with Google’s identity and access management service.

Cool. Now what do we do with this? Well previously, we were serving our app on port 8000. We’re going to skip that part entirely and stream our app directly into the tunnel.

Change this:

let addr = SocketAddr::from(([127, 0, 0, 1], 8000));
    log::info!("App started on URL: {:?}", addr.to_string());

    axum::Server::bind(&addr)
        .serve(app.into_make_service_with_connect_info::<SocketAddr>())
        .await?;

To this:

log::info!("Tunnel started on URL: {:?}", tunnel.url());

    axum::Server::builder(tunnel)
        .serve(app.into_make_service_with_connect_info::<SocketAddr>())
        .await?;

We got rid of the SocketAddr and changed axum::Server::bind() to axum::Server::builder() and passed our tunnel directly into this. Now when we run our app, we get a public https URL from which our app is being served. Online in one(ish) lines!

Getting Usernames

The final piece of our app is integrating usernames for messages as they come in. Since we’re using Single Sign On (SSO), we already have usernames from the Identity Provider (Google in our case). ngrok adds these usernames as a header ngrok-auth-user-name which we can feed into our chat_service function. Modify the websocket route function we made earlier:

pub async fn websocket(
    ws: WebSocketUpgrade,
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
) -> impl IntoResponse {
    let username = match headers.get("ngrok-auth-user-name") {
        Some(header) => header.to_str().unwrap_or("None").to_string(),
        None => "None".to_string(),
    };
    ws.on_upgrade(|socket| chat_service(socket, state, username))
}

Now, user messages will display whatever name they have configured with Google. Pretty clean!

Conclusion

We saw the common approach to building a web app and getting it on the internet, which requires running a local service followed by a few hours of configuration.

Then we saw how, with just a few lines of added code, we could have a public and secure web service up and running without having to fool around with ports and proxies.

Integrating ngrok directly into the code makes serving web applications effortless and painless.

This was a pretty cool project for me because I’ve never built a chat app and never built anything which used websockets. Figuring out how to translate the high-level Python abstractions for websockets into slightly lower-level Rust code was a bit of a challenge as well. But at the end of the day, I got to learn something new, integrate a powerful web service (for free!), and build something fun.

‍

‍

‍

‍