new-home-proxy 0.1.2

extern crate actix_rt;

use std::sync::mpsc::{channel, Sender};
use std::sync::Arc;
use std::thread;
use std::thread::JoinHandle;

use actix::clock::Duration;
use actix_http::http::header::{
    ACCESS_CONTROL_ALLOW_HEADERS, ACCESS_CONTROL_ALLOW_METHODS, ACCESS_CONTROL_ALLOW_ORIGIN,
    LOCATION,
};
use actix_http::http::HeaderValue;
use actix_web::dev::{Service, ServiceResponse};
use actix_web::web::{route, Data};
use actix_web::{App, HttpResponse, HttpServer};

use new_home_proxy::proxy_error::ProxyResult;
use new_home_proxy::server::client_manager::ClientManager;
use new_home_proxy::server::config::Config;
use new_home_proxy::server::handler::{handle_connect, handle_proxy};

#[actix_rt::main]
async fn main() {
    if let Err(error) = run().await {
        println!("{}", error);
    }
}

async fn run_server(manager: Arc<ClientManager>, bind: String) -> ProxyResult<()> {
    let server = HttpServer::new(move || {
        App::new()
            .app_data(Data::new(Arc::clone(&manager)))
            .route("/connect/{client_id}", route().to(handle_connect))
            .route(
                "/proxy/{client_id}/{request_uri:.*}",
                route().to(handle_proxy),
            )
            .route(
                "{_:.*}",
                route().to(|| {
                    let mut response = HttpResponse::PermanentRedirect();
                    response.header(LOCATION, "https://gitlab.com/y_software/new-home-proxy");

                    response
                }),
            )
            .wrap_fn(|req, srv| {
                let origin = match req.headers().get("Origin") {
                    Some(origin) => String::from(origin.to_str().unwrap_or_default()),
                    _ => String::new(),
                };
                let fut = srv.call(req);

                async move {
                    let mut res: ServiceResponse = fut.await.unwrap();

                    res.headers_mut().insert(
                        ACCESS_CONTROL_ALLOW_HEADERS,
                        HeaderValue::from_str("Authorization, Content-Type").unwrap(),
                    );
                    res.headers_mut().insert(
                        ACCESS_CONTROL_ALLOW_METHODS,
                        HeaderValue::from_str("*").unwrap(),
                    );
                    res.headers_mut().insert(
                        ACCESS_CONTROL_ALLOW_ORIGIN,
                        HeaderValue::from_str(origin.as_str()).unwrap(),
                    );

                    Ok(res)
                }
            })
    })
    .bind(bind)?
    .run();

    println!("Server running");

    Ok(server.await?)
}

fn run_ping(manager: Arc<ClientManager>) -> (Sender<bool>, JoinHandle<()>) {
    let (thread_send, thread_recv) = channel::<bool>();

    let ping_thread = thread::spawn(move || {
        println!("Ping thread started");

        loop {
            manager.ping_all_clients();

            if thread_recv
                .recv_timeout(Duration::from_millis(1000))
                .unwrap_or(false)
            {
                break;
            }
        }

        println!("Ping stopped");
    });

    (thread_send, ping_thread)
}

async fn run() -> ProxyResult<()> {
    let manager = Arc::new(ClientManager::new());
    let path = std::env::current_dir()?;
    let config = Config::initialize(path.join("server.yaml"))?;
    let (thread_send, ping_thread) = run_ping(manager.clone());
    let bind_address = std::env::var("ADDRESS").unwrap_or(config.bind_address.clone());
    let bind_port = std::env::var("PORT").unwrap_or(config.bind_port.to_string());

    run_server(manager, format!("{}:{}", bind_address, bind_port)).await?;

    thread_send.send(true).unwrap_or_default();
    ping_thread.join().unwrap_or_default();

    Ok(())
}