4
\$\begingroup\$

Recently I needed to support someone who wanted to ingest some data from on of our API's but had tooling that struggled with the authentication bit. While I'm sure there are some off-the-shelf products that could have been made do somehow, I decided to take this as a learning opportunity to learn Rust and create a small proxy that would transparently handle the authentication and token refresh.

Took me a while to sift through the myriads of http related crates until I settled on axum. Rouille also looked promising but their proxy implementation can't do TLS (according to the docs). Axum turned out quite nice though, especially the service layer functionality.

I'm looking for general best-practice-Rust feedback, especially a hint on how to improve error handling. I also get the feeling I should be using combinators more but not sure how that would effectively look like.

use axum::{
    body::Body,
    debug_handler,
    extract::{ Extension, Request },
    response::IntoResponse,
    Router,
};
use chrono::{ DateTime, Duration, Local };
use const_format::concatcp;
use flexi_logger::{ Cleanup, Criterion, Duplicate, FileSpec, Logger, Naming };
use log::{ trace, debug, info, warn, error };
use serde::{
    de::{ DeserializeOwned, Error }
    Deserialize,
    Deserializer
};
use serde_aux::prelude::*;
use serde_json::json;
use std::{
    collections::HashMap,
    ops::Deref,
    sync::Arc, 
    sync::RwLock,
};
use tower::ServiceBuilder;
use tower_http::{
    add_extension::AddExtensionLayer,
    trace::TraceLayer,
};

const API_BASE_HOST: &'static str = "downstreamapi.example.com";

const API_BASE_URI: &'static str = concatcp!("https://", API_BASE_HOST);

#[derive(serde_derive::Serialize, serde_derive::Deserialize, Default)]
struct Settings {
    #[serde(deserialize_with = "deserialize_non_empty_string")]
    bind_to: String,
    local_port: u16,
    log_level: String,
    #[serde(deserialize_with = "deserialize_non_empty_string")]
    username: String,
    impersonator: Option<String>,
    #[serde(deserialize_with = "deserialize_non_empty_string")]
    password: String,
}

struct AuthToken {
    token: String,
    expiry: Option<DateTime<Local>>,
}

impl AuthToken {
    fn new(tok: String, expires_in: i64) -> Self {
        Self { token: tok, expiry: Some(Local::now() + Duration::seconds(expires_in)) }
    }
}

type SharedState = Arc<State>;

#[derive(Default)]
struct Authentication {
    access_token: Option<AuthToken>,
    refresh_token: Option<AuthToken>,
}

#[derive(Default)]
struct State {
    auth: RwLock<Authentication>,
    client: reqwest::Client,
}

#[derive(serde_derive::Serialize, serde_derive::Deserialize)]
struct AuthResponse {
    access_token: String,
    #[serde(deserialize_with = "deserialize_number_from_string")]
    expires_in: i64,
    refresh_token: String,
}

#[derive(serde_derive::Serialize, serde_derive::Deserialize)]
struct RefreshResponse {
    access_token: String,
    #[serde(deserialize_with = "deserialize_number_from_string")]
    expires_in: i64,
    refresh_token: String,
    #[serde(deserialize_with = "deserialize_number_from_string")]
    refresh_token_expires_in: i64,
}

fn is_token_valid(tok: &Option<AuthToken>, tok_name: &str) -> bool {
    if tok.as_ref().is_some_and(|t| t.expiry.is_some()) {
        let now = Local::now();
        let expiry = tok.as_ref().unwrap().expiry.unwrap();
        // if the expiry is more than 10 minutes in the future we consider the token valid
        // 10 minutes should be sufficient buffer to account for clock drifts and finishing
        // an active request
        if now + Duration::minutes(10) < expiry {
            return true;
        }
    }
    info!("{} token not present or expired", tok_name);
    false
}

async fn post_json<T>(
    cfg: &Settings,
    map: HashMap<&str, String>,
    client: &reqwest::Client,
    path: &str
) -> Result<T, anyhow::Error>
    where T: DeserializeOwned
{
    let mut request = client
        .post(format!("{API_BASE_URI}/api/v2.4/{path}"))
        .header("Content-Type", "application/json")
        .body(json!(map).to_string());
    let mut suffix = "no impersonation".to_owned();
    if cfg.impersonator.is_some() {
        request = request.header("x-impersonated-user", cfg.username.to_string());
        suffix = format!("impersonating {}", cfg.username.to_string()).to_owned();
    }
    debug!("POST to {}, {}", path, suffix);

    let res = request.send().await?;
    
    match res.error_for_status_ref() {
        Ok(_) => Ok(res.json::<T>().await?),
        Err(e) => Err(e.into()),
    }
}

async fn get_access_token(auth: &RwLock<Authentication>, client: reqwest::Client) -> Result<String, anyhow::Error> {
    if !is_token_valid(&auth.read().unwrap().access_token, "auth") {
        let cfg: Settings = confy::load_path("config.toml")?;

        // if we don't have a valid access token we need to get a new one
        // either use the refresh_token to exchange (if it's still valid) or
        // re-login with credentials
        if is_token_valid(&auth.read().unwrap().refresh_token, "refresh") {
            info!("Attempting to obtain new access token via refresh");
            let mut map: HashMap<&str, String> = HashMap::new();
            map.insert("refreshToken", auth.read().unwrap().refresh_token.as_ref().unwrap().token.to_string());
            let resp: RefreshResponse = post_json(&cfg, map, &client, "security/refresh").await?;
            auth.write().unwrap().access_token = Some(AuthToken::new(resp.access_token, resp.expires_in));
            auth.write().unwrap().refresh_token = Some(AuthToken::new(resp.refresh_token, resp.refresh_token_expires_in));
            info!("Obtained new access token via refresh");
        } else {
            info!("Attempting to obtain new access token via login ({})", cfg.username);
            // ok, so we have no access and no refresh token which means we need to login with our credentials
            let mut map: HashMap<&str, String> = HashMap::new();
            map.insert("username", cfg.impersonator.as_ref().unwrap_or(&cfg.username).to_string());
            map.insert("password", cfg.password.to_string());
            let resp: AuthResponse = post_json(&cfg, map, &client, "security/auth").await?;

            auth.write().unwrap().access_token = Some(AuthToken::new(resp.access_token, resp.expires_in));
            // The auth response doesn't explicitly return the expiry of the refresh token and you are not
            // supposed to decode the refresh token either, so for now we can only assume the default 14 day validity
            auth.write().unwrap().refresh_token = Some(AuthToken::new(resp.refresh_token, 14 * 24 * 3600));
            info!("Obtained new access token via login");
        }
    }

    // at this point we should have a valid token or have already returned an error
    return Ok(auth.read().unwrap().access_token.as_ref().unwrap().token.to_string());
}

fn translate_method(axum_method: &axum::http::Method) -> Result<reqwest::Method, anyhow::Error> {
    match axum_method {
        &axum::http::Method::GET => Ok(reqwest::Method::GET),
        &axum::http::Method::POST => Ok(reqwest::Method::POST),
        _ => anyhow::bail!("Unsupported request method {}", axum_method),
    }
}

async fn reverse_proxy(req: axum::extract::Request<Body>, headers: http::HeaderMap, state: SharedState) -> Result<axum::response::Response, AppError> {
    let token = get_access_token(&state.deref().auth, state.deref().client.clone()).await?;
    let method = translate_method(req.method())?;
    let destination = format!("{API_BASE_URI}{}", req.uri());

    debug!("Forwarding {} request to {}", method, destination);
    trace!("original request: {:?}", req);

    let (_, req_body) = req.into_parts();

    let mut r = state.client
                .request(method, reqwest::Url::parse( &destination)?)
                .bearer_auth(token)
                .body(axum::body::to_bytes(req_body, 1 * 1024 * 1024).await?)
                .version(reqwest::Version::HTTP_2);

    // copy all headers, except:
    // user agent - modify to include ourself,
    // host - replace with downstream target
    // connection - need to remove if present since it's not valid for HTTP/2
    for (name, value) in headers.iter() {
        let val: String = value.to_str().unwrap().to_string();
        if let Some(v) = match name {
            &http::header::USER_AGENT => Some(format!("{} (via {}/{})", val, env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION")).to_string()),
            &http::header::HOST => Some(concatcp!(API_BASE_HOST, ":443").to_string()),
            &http::header::CONNECTION => None,
            _ => Some(val)
        } {
            r = r.header(name.to_string(), v);
        }
    }

    let outgoing_request = r.build()?;
    trace!("forwarded request: {:?}", outgoing_request);

    Ok(state.client.execute(outgoing_request).await?.text().await?.into_response())
}

#[debug_handler]
async fn handler(Extension(state): Extension<SharedState>, headers: http::HeaderMap, req: Request<Body>) -> impl axum::response::IntoResponse {
    debug!("Handling {} request to {}", req.method(), req.uri());
    match reverse_proxy(req, headers, state).await {
        Ok(r) => r,
        Err(e) => e.into_response(),
    }
}

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    let cfg: Settings = confy::load_path("config.toml")?;

     // we're going to log everything to files under logs/ and duplicate info and higher to console
    Logger::try_with_str(cfg.log_level)
                .unwrap()
                .log_to_file(FileSpec::default().directory("logs"))
                .rotate(
                    Criterion::Size(10 * 1024 * 1024),
                    Naming::Timestamps,
                    Cleanup::KeepLogFiles(7),
                 )
                .duplicate_to_stderr(Duplicate::Info) 
                .start()
                .unwrap();

    let app = Router::new()
        .fallback(handler)
        .layer(
            ServiceBuilder::new()
                .layer(AddExtensionLayer::new(SharedState::default()))
                .layer(TraceLayer::new_for_http())
                .into_inner(),
        );

    let listener = tokio::net::TcpListener::bind(format!("{}:{}", cfg.bind_to, cfg.local_port)).await.unwrap();
    info!("Listening on {}", listener.local_addr().unwrap());
    Ok(axum::serve(listener, app.into_make_service()).await?)
}

// Make our own error that wraps `anyhow::Error`.
struct AppError(anyhow::Error);

// Tell axum how to convert `AppError` into a response.
impl axum::response::IntoResponse for AppError {
    fn into_response(self) -> axum::response::Response {
        if self.0.is::<reqwest::Error>() {
            warn!("Downstream API returned an HTTP error");
            let code = http::StatusCode::from_u16(self.0.downcast_ref::<reqwest::Error>().unwrap().status().unwrap().as_u16()).unwrap();
            return (
                code,
                format!("HTTP request failed (code {code}): {}", self.0),
            ).into_response();
        }
        error!("{:?}", self.0.source());
        (
            http::StatusCode::INTERNAL_SERVER_ERROR,
            format!("Something went wrong: {}", self.0),
        ).into_response()
    }
}

// This enables using `?` on functions that return `Result<_, anyhow::Error>` to turn them into
// `Result<_, AppError>`. That way you don't need to do that manually.
impl<E> From<E> for AppError
where
    E: Into<anyhow::Error>,
{
    fn from(err: E) -> Self {
        Self(err.into())
    }
}

fn deserialize_non_empty_string<'de, D>(deserializer: D) -> Result<String, D::Error>
where
    D: Deserializer<'de>,
{
    let s: String = Deserialize::deserialize(deserializer)?;
    if s.is_empty() {
        Err(D::Error::custom("empty string not allowed"))
    } else {
        Ok(s)
    }
}
\$\endgroup\$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Browse other questions tagged or ask your own question.