Jos Dirksen

Full-stack developer.

Getting started with Rust - part 2

11 minute read

In the previous article I started experimenting with Rust and build a very simple extension to git, which helps me remember the most recent branches I’ve been working on, or did PRs for. For this second article, I wanted to do look at how you can do network programming in Rust.

(If you just want to skip to the code: https://github.com/josdirksen/dns-proxy-rusts)

For this, and the next article, since I’ve only done the basic so far, we’ll look at the steps needed to create a DNS proxy in Rust. This DNS proxy in the end will have the following features:

  • Proxy all incoming DNS requests to a real DNS server, and respond back to the client with the response from the real server.
  • Keep metrics of how often a specific host was queried.
  • Allow the use of /etc/hosts with wildcards as a source for A-Record queries.
  • And probably some other stuff, which I’ll think of when implementing.

One of reasons I want something simple like this is that at work we use wildcard certificates to connect specific domains to appications. So ui-dev.smartjava.org, would route to a specific instance, while ui-test.smartjava.org would route to another one. In some cases though, for instance when debugging, or testing k8s or kafka configs, it is useful to route these domains to localhost, or another host completely. We can do this by adding all the domains, including the subdomains to your /etc/hosts file, which solves part of the problem, but that only works when working directly on your local machine. If you start containers in docker, they won’t know about these changes. So with a simple DNS proxy, we can have docker instances use the information from the docker host for routing. And of course I’m just interested to see what kind of DNs queries my machine is making.

For this first article we’ll focus on the first of these points, and try to get the basic application up and running. As a more technical / Rust goal I wanted to mostly learn about:

  • Network programming
  • async/await
  • Indirectly about the whole ownership model

Another change I made in regards to the previous article, was that I switched editors a couple of times. Before going into the details on the code I’ve written, a quick note on that.

Visual Studio Code as Rust Editor

For my day job I do a lot of Kotlin, and you don’t really have much choice of editor. You just use Intellij Idea. Which is a great editor, but can sometimes become slow and unresponsive, since it’s got so many features, checks and other smart things running in the background, analyzing your code. In the previous article, I used the Intellij Rust plugin to write the git extension, but for this DNS example I’ve switched to using Visual Studio Code.

Visual studio code

Code has different plugins you can use for editing Rust, and I tried the following two:

I started off with the offical extension, and initially it worked great. I got code completion, could quickly skip through source code, and had nice inline comments. It also provided integration with the Code tasks system to easily run tasks ons cargo. However, when I started adding the async/await stuff, it just stopped working. I still had some code highlighting, and it showed the errors (the squiggly lines), but there was no code completion anymore, and jumping to implementations also pretty much stopped working. Apparently this is a know issue, and is caused by the macros from Tokio, which I used to add async support to my code. At that point I, momentarily, gave up on Rust in Code and switched back to IntelliJ. But, unfortunately, with the Tokio stuff added, IntelliJ was giving the same kind of errors. So back to Visual Studio Code, and then I tried the Rust Analyzer extension. This one was actually working quite ok with the async stuff, so I’ve written the rest of the product with that extension.

( There is also Another rust with RLS extension, on the marketplace, but I didn’t test that one.)

I’ll check back in the near future how the various extensions evolve, but at least I’ve now got a light weight solution to work with in Visual Studio Code, which works quite nice with the tasks in Cargo.

Quick introduction on UDP

For this first version we’re just going to create a DNS Proxy which works with UDP. You can also run DNS queries over TCP, so we’ll add that in the next article. UDP is a connection-less protocol so we need to take that into account when creating our service. The following image from this slidedeck (Introduction into DNS), provides a nice explanation of what we need to do:

UDP explanation

So we need to write the following:

  1. An UDP server which can receive datagrams from a DNS client (we’ll just use dig as the client)
  2. When we receive an UDP request, we use an UDP client to forward the request to an external DNS server.
  3. When that external server responds with a DNS answers, we forward it back to the initial UDP client (dig).

First of, lets look at the libraries used for this proxy

Cargo configuration

We use the following cargo.toml.

[package]
name = "rust-dns"
version = "0.1.0"
authors = ["Jos Dirksen <jos.dirksen@gmail.com>"]
edition = "2018"

[dependencies]
dns-parser = "0.8.0"
simplelog = "^0.7.4"
log = "0.4.8"
tokio = { version = "0.2", features = ["full"] }

We’ve got the following libraries in here:

  • dns-parser: This library can parse incoming byte streams to an UDP packet. We don’t do much with this yet, just use it for some logging to see what is happening.
  • simplelog and log: These two libraries provide logging functionality. The log library is an abstraction on top of the different logging libraries out there, and with simplelog we just get a simple logger, with colors.
  • tokio: Rust has gotten async/await support a couple of months ago, and to use this you need a library that provides the execution runtime. Tokio is the best known one for this.

Before I’ll show my simple code, a quick sidestep on the whole async/await stuff.

Async/await

With async/await we get a lot of syntactic sugar on how to work with futures. So we can write asynchronous code, which can be read as sequential code. Rust has a complete async handbook (which I’ve only skimmed through so far) explaining this concept.

Basically what it allows is to write code like this:

async fn function_1() { ... }
async fn function_2() { ... }

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    function_1().await?
    function_1().await?
    Ok(())
}

I hadn’t seen then question mark thingy ? at the end of a function call beforehand, but it appears to be some syntactic sugar for dealing with Result types. It allows us to just chain together functions that return a Result without having to and_then or pattern match the results. Basically a minimal version of Haskell’s do-notation or Scala’s for-comprehensions.

Now let’s look at the code

The code

The code in itself isn’t that large, so I’ll first show the complete project so far with all its comments, and below that explain some of the stuff I was struggling with, or found interesting.

use dns_parser::rdata::a::Record;
use dns_parser::{Builder, Error as DNSError, Packet, RData, ResponseCode};
use dns_parser::{QueryClass, QueryType};
use io::Result as ioResult;
use log::*;
use simplelog::{Config, LevelFilter, TermLogger, TerminalMode};
use std::error::Error;
use std::net::SocketAddr;
use std::str;
use tokio::net::UdpSocket;
use tokio::prelude::*;

///
/// Simple dns proxy. For the first version, we'll just listen to
/// an UDP socket, parse an incoming DNS query and write the output
/// to the console.
///
/// Next steps:
///   Support TCP protocol
///   Keep track of hosts (and subhosts queried)
///   Allow resolving through local host file
///   Add error handling
///
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    init_logging();

    info!("Starting server, setting up listener for port 12345");

    // chain the await calls using '?'
    let listen_socket = create_udp_socket_receiver("0.0.0.0:12345").await?;
    let sender_socket = create_udp_socket_sender().await?;

    start_listening_udp(listen_socket, sender_socket).await?;
    Ok(())
}

///
/// Create an async UDP socket.
///
async fn create_udp_socket_receiver(host: &str) -> ioResult<UdpSocket> {
    debug!("initializing listener udp socket on {}", host);
    let socket = UdpSocket::bind(&host).await?;
    return Ok(socket);
}

///
/// Create the sender to forward the UDP request
///
async fn create_udp_socket_sender() -> ioResult<UdpSocket> {
    let local_address = "0.0.0.0:0";
    let socket = UdpSocket::bind(local_address).await?;
    let socket_address: SocketAddr = "8.8.8.8:53"
        .parse::<SocketAddr>()
        .expect("Invalid forwarding address specified");
    socket.connect(&socket_address).await?;
    debug!("initializing listener udp socket on {}", local_address);
    return Ok(socket);
}

///
/// Asynchronously run the handling of incoming messages. Whenever something comes in on the socket_rcv_from
/// we try and convert it to a DNS query, and log it to the output. UDP is a connectionless protocol, so we
/// can use this socket to send and receive messages to our client and other servers.
///
async fn start_listening_udp(
    mut listen_socket: UdpSocket,
    mut sender_socket: UdpSocket,
) -> ioResult<()> {
    // 1. Wait for a request from a DNS client.
    // 2. Then forward the request to a remote dns server
    // 3. The response from the remote DNS server is then send back to the initial client.
    loop {
        let (request, peer) = receive_request(&mut listen_socket).await?;
        let forward_response = forward_request(&mut sender_socket, &request[..]).await?;
        listen_socket.send_to(&forward_response[..], &peer).await?;
    }
}

///
/// Forward a request to the provided UDP socket, and wait for an answer.
///
async fn forward_request(sender_socket: &mut UdpSocket, request: &[u8]) -> ioResult<Vec<u8>> {
    let mut buf = [0; 4096];
    info!("Forwarding to target DNS");
    sender_socket.send(request).await?;
    let (amt, _) = sender_socket.recv_from(&mut buf).await?;
    let filled_buf = &mut buf[..amt];
    // let answer_received = parse_incoming_stream(filled_buf).expect("Something went wrong");

    let v = Vec::from(filled_buf);
    return Ok(v);
}

///
/// Receive a request on the reference to the socket. We're not the owner, but we need a mutable
/// reference. The result is a Vec, and the response address. Both are copies, which can be safely
/// modified.
///
async fn receive_request(from_socket: &mut UdpSocket) -> ioResult<(Vec<u8>, SocketAddr)> {
    let mut buf = [0; 4096];

    let (amt, peer) = from_socket.recv_from(&mut buf).await?;
    let filled_buf = &mut buf[..amt];
    // info!("Received length {}", amt);
    // let packet_received = parse_incoming_stream(filled_buf).expect("Something went wrong");
    // info!("Received package {:?}", packet_received);

    let v = Vec::from(filled_buf);
    return Ok((v, peer));
}

///
/// Parse the packet and return it
///
// fn parse_incoming_stream(incoming: &[u8]) -> Result<Packet, DNSError> {
//     let pkt = Packet::parse(incoming)?;
//     return Ok(pkt);
// }

///
/// Initializes the logging library. We just simply log to console
///
fn init_logging() {
    TermLogger::init(LevelFilter::Debug, Config::default(), TerminalMode::Mixed).unwrap();
}

Let’s start with the main function. In our case this creates a client UDP socket, and a server UDP socket.

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    init_logging();

    info!("Starting server, setting up listener for port 12345");

    let listen_socket = create_udp_socket_receiver("0.0.0.0:12345").await?;
    let sender_socket = create_udp_socket_sender().await?;

    start_listening_udp(listen_socket, sender_socket).await?;
    Ok(())
}

I’ll skip the implementations of the two create functions, which internally use the Tokio provided UDP network sockets. We use the Tokio provided ones, since they already work with async and await, so we don’t have to create and implement our own Future implementation on top of setting up the connections. When both of the sockets have been created, we pass them into the start_listening_udp function, which will continuously wait for new datagrams. Tokio provides an enhanced main function, which we can define as async itself, so we don’t have to do any blocking here.

In the function above you can also see how await together with the ? function really allows for easily readable asynchronous code. The next interesting piece of code is the start_listening_udp function (in the future I want something similar for the TCP functionality):

async fn start_listening_udp(
    mut listen_socket: UdpSocket,
    mut sender_socket: UdpSocket,
) -> ioResult<()> {
    loop {
        let (request, peer) = receive_request(&mut listen_socket).await?;
        let forward_response = forward_request(&mut sender_socket, &request[..]).await?;
        listen_socket.send_to(&forward_response[..], &peer).await?;
    }
}

This function is a never-ending loop on top of three async functions. First we wait to receive a request on the socket our server is listening on. Once we get that request (e.g a message from dig), we call the forward_request function. This function will forward the request to a real DNS server, wait for a response, and return with the response. Finally we send the response back to our original client (dig) and wait for new requests to come in. The nice thing about this code, is that it is very easy to reason about. The whole await/async stuff is nicely hidden from view, and we can sequence calls by usinig the question mark. I was impressed that a low-level language like Rust had such advanced constructs.

Let’s look at one of the functions, since the others are pretty much the same:

async fn forward_request(sender_socket: &mut UdpSocket, request: &[u8]) -> ioResult<Vec<u8>> {
    let mut buf = [0; 4096];
    info!("Forwarding to target DNS");
    sender_socket.send(request).await?;
    let (amt, _) = sender_socket.recv_from(&mut buf).await?;
    let filled_buf = &mut buf[..amt];
    let v = Vec::from(filled_buf);
    return Ok(v);
}

Here we have the same couple of await? steps. This time to make a call, and wait for a response. Easy right?

So finally when we run this, we get the following when we make a call:

Final output

Just like the previous project, it’s been really a fun process of doing this.

Conclusions

So what was easy and what wasn’t?

  • I’m still struggling with ownership, and when Rust is cleaning up variables used in functions. I’m slowly starting to understand it, but it does cause some strange and error messages when developing.
  • On those error messages. Whiile they usually are really helpful, sometimes it still takes a lot of time to actually solve the problem. This isn’t so much a Rust issue, as it is an issue for me that I’m still learning about specific concepts.
  • And it is always good to finish with a positivbe :) The whole async/await feels nice, works great and combined with the ? thingy, results in compact and readable code.

Updated:

You may also enjoy

Working with nix-shell

11 minute read

Managing runtime dependencies using Nix We’re setting up a new project at work, and wanted to apply some lessons learned from previous projects. One of the ...

Exploring Cats Effect

13 minute read

In two previous articles we explored how you can use ZIO to easily create programs in a type-safe manner: Exploring ZIO - Part I Exploring ZIO - Part I...

Exploring ZIO - Part I

10 minute read

I’ve done a lot of Kotlin for the last two years, and have mainly followed Scala and done some pet projects. Since a lot of interesting stuff has been going ...