Build a HTTP server with Rust and tokio - Part 1: serving static files
Tags: rust, http, webdev
Published: 2023-05-20 22:24:12

In this episode, we'll extend our server to serve static files.
We'll also refactor our code to support connection reuse, and we'll implement a graceful shutdown mechanism.

If your didn't follow the previous episode, you can find the code on [GitHub](https://github.com/geoffreycopin/http_server/tree/part_0).

As we will use new dependencies, we'll need to update our `Cargo.toml` file:
```bash
cargo add clap tokio-util futures
```

## Connection reuse
Under HTTP/1.0, a separate TCP connection is established for each request/response pair.
This is inefficient, as it requires a new TCP handshake for each request. HTTP/1.1 introduced connection reuse,
which allows multiple requests to be sent over the same TCP connection. This mechanism is also necessary
to support request pipelining, which we'll see in a future episode.

We'll keep waiting for new requests on the same connection until the client closes it, unless the client sets
the `Connection: close` header. In that case, we'll close the connection after sending the response.
All we have to do to implement this change is to wrap our client handling code in a loop:
```rust
// main.rs

// [...]
info!(?addr, "new connection");

loop {
    let req = match req::parse_request(&mut stream).await {
        Ok(req) => {
            info!(?req, "incoming request");
            req
        }
        Err(e) => {
            error!(?e, "failed to parse request");
            break;
        }
    };

    let close_connection =
        req.headers.get("Connection") == Some(&"close".to_string());

    let resp = resp::Response::from_html(
        resp::Status::NotFound,
        include_str!("../static/404.html"),
    );

    resp.write(&mut stream).await.unwrap();

    if close_connection {
        break;
    }
}
// [...]
```


## Serving static files
Answering every request with the same 'Not found' page was ok for a start, but the time has come to serve
some real content. In this section, we'll implement a handler that serves static files from a directory. 
This directory will be either the current working directory, or a directory specified by the user.

As our CLI is getting more complex, we'll use the [clap](https://crates.io/crates/clap) crate to parse the
command line arguments.

```rust
// args.rs
use std::path::PathBuf;

use clap::Parser;

#[derive(Parser, Debug)]
pub struct Args {
    #[arg(short, long, default_value_t = 8080)]
    pub port: u16,
    #[arg(short, long)]
    pub root: Option<PathBuf>,
}

// main.rs
use clap::Parser;

// [...]

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // [...]
    let args = args::Args::parse();
    let port = args.port;
    let listener = TcpListener::bind(format!("0.0.0.0:{port}")).await.unwrap();
    // [...]
}
```

In order to simplify our handling code, we'll also refactor our `Response` `struct` by 
using a trait object instead of a generic type parameter. This will allow use to use the same
`Response` type for all our handlers, even if they don't return the same type of response.

We'll also add a new `from_file` constructor to our `Response` type, which will allow us to
create a response `struct` from a file on disk. At this stage, we will not implement any
kind of sophisticated content negotiation, so we'll just set the `Content-Type` header to
a mime type based on the file extension.

```rust
// resp.rs
pub struct Response {
    pub status: Status,
    pub headers: HashMap<String, String>,
    pub data: Box<dyn AsyncRead + Unpin + Send>,
}

impl Response {
    // [...]
    pub async fn from_file(path: &Path, file: File) -> anyhow::Result<Response> {
        let headers = hashmap! {
            "Content-Length".to_string() => file.metadata().await?.len().to_string(),
            "Content-Type".to_string() => mime_type(path).to_string(),
        };

        Ok(Response {
            headers,
            status: Status::Ok,
            data: Box::new(file),
        })
    } 
    // [...]
}

fn mime_type(path: &Path) -> &str {
    match path.extension().and_then(|ext| ext.to_str()) {
        Some("html") => "text/html",
        Some("css") => "text/css",
        Some("js") => "text/javascript",
        Some("png") => "image/png",
        Some("jpg") => "image/jpeg",
        Some("gif") => "image/gif",
        _ => "application/octet-stream",
    }
}
```

These preparatory steps allow us to implement our static file handler in a few lines of code:
```rust
// handler.rs
use std::{env::current_dir, io, path::PathBuf};

use crate::{
    req::Request,
    resp::{Response, Status},
};

#[derive(Debug, Clone)]
pub struct StaticFileHandler {
    root: PathBuf,
}

impl StaticFileHandler {
    pub fn in_current_dir() -> io::Result<StaticFileHandler> {
        current_dir().map(StaticFileHandler::with_root)
    }

    pub fn with_root(root: PathBuf) -> StaticFileHandler {
        StaticFileHandler { root }
    }

    pub async fn handle(&self, request: Request) -> anyhow::Result<Response> {
        let path = self.root.join(request.path.strip_prefix('/').unwrap());

        if !path.is_file() {
            return Ok(Response::from_html(
                Status::NotFound,
                include_str!("../static/404.html"),
            ));
        }

        let file = tokio::fs::File::open(&path).await?;
        Response::from_file(&path, file).await
    }
}


// main.rs

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // [...]
    let handler = args
        .root
        .map(handler::StaticFileHandler::with_root)
        .unwrap_or_else(|| {
            handler::StaticFileHandler::in_current_dir().expect("failed to get current dir")
        });
    // [...]
    match handler.handle(req).await {
        Ok(resp) => {
            resp.write(stream).await.unwrap();
        }
        Err(e) => {
            error!(?e, "failed to handle request");
            return Ok(false);
        }
    };
    // [...]
}

```

After copying these files to the `static` directory, we can now serve them with our server:
```bash
cargo run -- --root static
```

Pressing `Ctrl+C` will stop the server, but in a rather abrupt way: the server will not log any
message regarding the shutdown, and will not wait for the pending requests to be processed before
closing the connection. This is not a big deal for a toy server, but in a real world application,
we would want to handle this more gracefully.

## Graceful shutdown
The most straightforward way to stop our server would be to collect the client handling tasks
into a `JoinSet` and abort them when we receive a `SIGINT` signal. However, by doing so, we would
have no way to wait for the pending requests to be processed before exiting the program.

Instead, we'll use a `CancellationToken` to signal that the server should:
* stop accepting new connection, 
* stop processing requests from the current connections

We'll use the two methods provided by the `CancellationToken`: 
* `cancel`: request cancellation from the main thread when the user presses `Ctrl+C`
* `cancelled`: return a future that resolves when the cancellation is requested. Used in conjunction
  with `select!`, this will allow us to stop the client handling tasks when the cancellation is
  requested.

Our `main` function now looks like this:
```rust
// [...]
async fn main() -> anyhow::Result<()> {
    // [...]
    let cancel_token = CancellationToken::new();

    tokio::spawn({
        let cancel_token = cancel_token.clone();
        async move {
            if let Ok(()) = signal::ctrl_c().await {
                info!("received Ctrl-C, shutting down");
                cancel_token.cancel();
            }
        }
    });
    // [...]

    let mut tasks = Vec::new();

    loop {
        let cancel_token = cancel_token.clone();

        tokio::select! {
            Ok((stream, addr)) = listener.accept() => {
                let handler = handler.clone();
                let client_task = tokio::spawn(async move {
                    if let Err(e) = handle_client(cancel_token, stream, addr, &handler).await {
                        error!(?e, "failed to handle client");
                    }
                });
                tasks.push(client_task);
            },
            _ = cancel_token.cancelled() => {
                info!("stop listening");
                break;
            }
        }
    }

    futures::future::join_all(tasks).await;

    Ok(())
}

async fn handle_client(
    cancel_token: CancellationToken,
    stream: TcpStream,
    addr: SocketAddr,
    handler: &handler::StaticFileHandler,
) -> anyhow::Result<()> {
    let mut stream = BufStream::new(stream);

    info!(?addr, "new connection");

    loop {
        tokio::select! {
            req = req::parse_request(&mut stream) => {
                match req {
                    Ok(req) => {
                        info!(?req, "incoming request");
                        let close_conn = handle_req(req, &handler, &mut stream).await?;
                        if close_conn {
                            break;
                        }
                    }
                    Err(e) => {
                        error!(?e, "failed to parse request");
                        break;
                    }
                }
            }
            _ = cancel_token.cancelled() => {
                info!(?addr, "closing connection");
                break;
            }
        }
    }

    Ok(())
}

async fn handle_req<S: AsyncWrite + Unpin>(
    req: req::Request,
    handler: &handler::StaticFileHandler,
    stream: &mut S,
) -> anyhow::Result<bool> {
    let close_connection = req.headers.get("Connection") == Some(&"close".to_string());

    match handler.handle(req).await {
        Ok(resp) => {
            resp.write(stream).await.unwrap();
        }
        Err(e) => {
            error!(?e, "failed to handle request");
            return Ok(false);
        }
    };

    Ok(close_connection)
}
```

We can now stop the server by pressing `Ctrl+C`, and the server will wait for the pending requests
to be processed before exiting.

The full source code for this part is available [here](https://github.com/geoffreycopin/http_server).





Looking for a Rust dev? [Let’s get in touch!](mailto:copin.geoffrey@gmail.com)