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

extern crate core;

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

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

use base::RemoteConfig;

use crate::client_config::ClientConfig;
use crate::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).expect("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 {
        log::error!(
            "Could not retrieve configuration! (got status {})",
            req.status()
        );
        std::process::exit(2);
    }

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

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

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

    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 = get_server_config(&args).await?;

    // 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(())
}