kochab/src/handling.rs

//! Types for handling requests
//!
//! The main type is the [`Handler`], which wraps a more specific type of handler and
//! manages delegating responses to it.
//!
//! For most purposes, you should never have to manually create any of these structs
//! yourself, though it may be useful to look at the implementations of [`From`] on
//! [`Handler`], as these are the things that can be used as handlers for routes.
use anyhow::Result;

use std::{
    pin::Pin,
    future::Future,
    task::Poll,
    panic::{catch_unwind, AssertUnwindSafe},
};

use crate::types::{Response, Request};

/// A struct representing something capable of handling a request.
///
/// In the future, this may have multiple varieties, but at the minute, it just wraps an
/// [`Fn`](std::ops::Fn).
///
/// The most useful part of the documentation for this is the implementations of [`From`]
/// on it, as this is what can be passed to
/// [`Builder::add_route`](crate::Builder::add_route) in order to create a new route.
/// Each implementation has bespoke docs that describe how the type is used, and what
/// response is produced.
pub enum Handler {
    FnHandler(HandlerInner),
}

/// Since we can't store train objects, we need to wrap fn handlers in a box
type HandlerInner = Box<dyn Fn(Request) -> HandlerResponse + Send + Sync>;
/// Same with dyn Futures
type HandlerResponse = Pin<Box<dyn Future<Output = Result<Response>> + Send>>;

impl Handler {
    /// Handle an incoming request
    ///
    /// This delegates to the request to the appropriate method of handling it, whether
    /// that's fetching a file or directory listing, cloning a static response, or handing
    /// the request to a wrapped handler function.
    ///
    /// Any unexpected errors that occur will be printed to the log and potentially
    /// reported to the user, depending on the handler type.
    pub async fn handle(&self, request: Request) -> Response {
        match self {
            Self::FnHandler(inner) => {
                let fut_handle = (inner)(request);
                let fut_handle = AssertUnwindSafe(fut_handle);

                HandlerCatchUnwind::new(fut_handle).await
                    .unwrap_or_else(|err| {
                        error!("Handler failed: {:?}", err);
                        Response::server_error("").unwrap()
                    })
            },
        }
    }
}

impl<H, R> From<H> for Handler
where
    H: 'static + Fn(Request) -> R + Send + Sync,
    R: 'static + Future<Output = Result<Response>> + Send,
{
    /// Wrap an [`Fn`] in a [`Handler`] struct
    ///
    /// This automatically boxes both the [`Fn`] and the [`Fn`]'s response.
    ///
    /// Any requests passed to the handler will be directly handed down to the handler,
    /// with the request as the first argument.  The response provided will be sent to the
    /// requester.  If the handler panics or returns an [`Err`], this will be logged, and
    /// the requester will be sent a [`SERVER_ERROR`](Response::server_error()).
    fn from(handler: H) -> Self {
        Self::FnHandler(
            Box::new(move|req| Box::pin((handler)(req)) as HandlerResponse)
        )
    }
}

/// A utility for catching unwinds on Futures.
///
/// This is adapted from the futures-rs CatchUnwind, in an effort to reduce the large
/// amount of dependencies tied into the feature that provides this simple struct.
#[must_use = "futures do nothing unless polled"]
struct HandlerCatchUnwind {
    future: AssertUnwindSafe<HandlerResponse>,
}

impl HandlerCatchUnwind {
    fn new(future: AssertUnwindSafe<HandlerResponse>) -> Self {
        Self { future }
    }
}

impl Future for HandlerCatchUnwind {
    type Output = Result<Response>;

    fn poll(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context
    ) -> Poll<Self::Output> {
        match catch_unwind(AssertUnwindSafe(|| self.future.as_mut().poll(cx))) {
            Ok(res) => res,
            Err(e) => {
                error!("Handler panic! {:?}", e);
                Poll::Ready(Response::server_error(""))
            }
        }
    }
}