tcp-over-http/src/tcp_relay_client/mod.rs

extern crate core;

use std::error::Error;
use std::io::ErrorKind;
use std::sync::Arc;

use futures::future::join_all;
use reqwest::{Certificate, Identity};

use crate::base::err_utils::{encpasulate_error, new_err};
use crate::base::RemoteConfig;
use crate::tcp_relay_client::client_config::ClientConfig;
use crate::tcp_relay_client::relay_client::relay_client;

pub mod client_config;
mod relay_client;

/// Get remote server config i.e. get the list of forwarded ports
async fn get_server_config(conf: &ClientConfig) -> Result<RemoteConfig, Box<dyn Error>> {
    let url = format!("{}/config", conf.relay_url);
    log::info!("Retrieving configuration on {}", url);

    let mut client = reqwest::Client::builder();

    // Specify root certificate, if any was specified in the command line
    if let Some(cert) = conf.get_root_certificate() {
        client = client.add_root_certificate(Certificate::from_pem(&cert)?);
    }

    // Specify client certificate, if any
    if let Some(kp) = conf.get_merged_client_keypair() {
        let identity = Identity::from_pem(&kp)
            .map_err(|e| encpasulate_error(e, "Failed to load certificates for reqwest!"))?;
        client = client.identity(identity).use_rustls_tls();
    }

    let client = client.build().expect("Failed to build reqwest client");

    let req = client
        .get(url)
        .header("Authorization", format!("Bearer {}", conf.get_auth_token()))
        .send()
        .await?;
    if req.status().as_u16() != 200 {
        Err(std::io::Error::new(
            ErrorKind::Other,
            format!(
                "Could not retrieve configuration! (got status {})",
                req.status()
            ),
        ))?;
    }

    Ok(req.json::<RemoteConfig>().await?)
}

/// Core logic of the application
pub async fn run_app(mut args: ClientConfig) -> std::io::Result<()> {
    args.load_certificates()?;
    let args = Arc::new(args);

    // Check arguments coherence
    if args.tls_cert.is_some() != args.tls_key.is_some() {
        return Err(new_err(
            "If you specify one of TLS certificate / key, you must then specify the other!",
        ));
    }

    if args.get_client_keypair().is_some() {
        log::info!("Using client-side authentication");
    }

    // Get server relay configuration (fetch the list of port to forward)
    let remote_conf = match get_server_config(&args).await {
        Ok(c) => c,
        Err(e) => {
            Err(std::io::Error::new(
                ErrorKind::Other,
                format!("Failed to fetch relay configuration from server! {}", e),
            ))?;
            unreachable!();
        }
    };

    // Start to listen port
    let mut handles = vec![];
    for port in remote_conf {
        let listen_address = format!("{}:{}", args.listen_address, port.port);

        let h = tokio::spawn(relay_client(
            format!(
                "{}/ws?id={}&token={}",
                args.relay_url,
                port.id,
                urlencoding::encode(args.get_auth_token())
            )
            .replace("http", "ws"),
            listen_address,
            args.clone(),
        ));
        handles.push(h);
    }

    join_all(handles).await;

    Ok(())
}