rust warp wrap_fn no-cache filter example

warp_wrap_fn_no_cache_example.rs

/*
[dependencies]
log = "0.4"
env_logger = "0.7"

futures = "0.3"
tokio = { version = "0.2", features = ["full"] }

warp = "0.2"

serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
*/

use env_logger::{Builder, Env};
use std::env;
use std::net::SocketAddr;
use std::convert::Infallible;

use serde::Serialize;

use warp::{Filter, Rejection, Reply, reply};
use warp::http::{StatusCode, HeaderValue};

#[tokio::main]
async fn main() {
    env::set_var("RUST_LOG", "warp=info, warn");
    Builder::from_env(Env::default().default_filter_or("info")).init();

    let addr = SocketAddr::from(([0, 0, 0, 0], 8080));

    warp::serve(router()).run(addr).await;
}

async fn ok() -> Result<impl Reply, Rejection> {
    Ok(warp::reply::json(&"OK".to_string()))
}

fn ok_route(
) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone + Send + Sync + 'static {
    warp::path!("ok")
        .and(warp::path::end())
        .and_then(ok)
}

#[derive(Serialize)]
struct ErrorMessage {
    code: u16,
    message: String,
}

async fn rejection_handler(r: Rejection) -> Result<impl Reply, Infallible> {
    let code;
    let message;

    if r.is_not_found() {
        code = StatusCode::NOT_FOUND;
        message = "NOT_FOUND";
    } else if let Some(_) =
    r.find::<warp::filters::body::BodyDeserializeError>()
    {
        code = StatusCode::BAD_REQUEST;
        message = "BAD_REQUEST";
    } else if let Some(_) = r.find::<warp::reject::MethodNotAllowed>() {
        code = StatusCode::METHOD_NOT_ALLOWED;
        message = "METHOD_NOT_ALLOWED";
    } else {
        code = StatusCode::INTERNAL_SERVER_ERROR;
        message = "INTERNAL_SERVER_ERROR";
    }

    let json = warp::reply::json(&ErrorMessage {
        code: code.as_u16(),
        message: message.to_string()
    });

    Ok(warp::reply::with_status(json, code))
}

fn router() -> impl warp::Filter<
    Extract = impl warp::Reply,
    Error = std::convert::Infallible,
> + Clone + Send + Sync + 'static {
    ok_route().recover(rejection_handler)
        .with(warp::wrap_fn(no_cache_filter))
}

fn disable_cache(reply: impl Reply) -> impl Reply {
    reply::with_header(
        reply,
        "cache_control",
        HeaderValue::from_static("no-cache"),
    )
}

fn no_cache_filter<F, T>(
    filter: F,
) -> impl Filter<Extract = impl Reply, Error = std::convert::Infallible> + Clone + Send + Sync + 'static
    where
        F: Filter<Extract = (T,), Error = std::convert::Infallible> + Clone + Send + Sync + 'static,
        F::Extract: warp::Reply,
        T: warp::Reply,
{
    warp::any()
        .and(filter)
        .map(|r| disable_cache(r))
}