diff --git a/src/coordinator.rs b/src/coordinator.rs index ac8d4fb..e761a24 100644 --- a/src/coordinator.rs +++ b/src/coordinator.rs @@ -1,21 +1,23 @@ -use std::{path::Path, sync::Arc, time::Duration}; use std::io::Write; +use std::{path::Path, sync::Arc, time::Duration}; use http_body_util::Full; use serde::{Deserialize, Serialize}; +use tor_cell::relaycell::msg::Connected; use tor_config::CfgPath; -use tor_hsservice::{config::OnionServiceConfigBuilder, HsNickname, RendRequest, StreamRequest, RunningOnionService}; +use tor_hsservice::{ + config::OnionServiceConfigBuilder, HsNickname, RendRequest, RunningOnionService, StreamRequest, +}; use tor_proto::stream::IncomingStreamRequest; -use tor_cell::relaycell::msg::Connected; use veq::veq::ConnectionInfo; use crate::protocol::ConnectionManager; use futures::Stream; use futures::StreamExt; -use hyper_util::rt::TokioIo; -use hyper::{body::Bytes, server::conn::http1, Method, Response, StatusCode}; use hyper::service::service_fn; +use hyper::{body::Bytes, server::conn::http1, Method, Response, StatusCode}; +use hyper_util::rt::TokioIo; use arti_client::{TorClient, TorClientConfig}; @@ -28,21 +30,37 @@ pub struct AugmentedInfo { pub display_name: String, } -pub async fn create_tor_client(config_dir: &Path, nickname: String) -> (TorClient, Arc, impl Stream) { +pub async fn create_tor_client( + config_dir: &Path, + nickname: String, +) -> ( + TorClient, + Arc, + impl Stream, +) { let tor_cache_dir = config_dir.join("tor-cache"); let tor_state_dir = config_dir.join("tor-state"); let mut client_config_builder = TorClientConfig::builder(); - client_config_builder.storage().cache_dir(CfgPath::new_literal(tor_cache_dir)).state_dir(CfgPath::new_literal(tor_state_dir)); - client_config_builder.stream_timeouts().connect_timeout(Duration::from_secs(10)); - let client_config = client_config_builder.build().expect("Failed to set up tor client config."); + client_config_builder + .storage() + .cache_dir(CfgPath::new_literal(tor_cache_dir)) + .state_dir(CfgPath::new_literal(tor_state_dir)); + client_config_builder + .stream_timeouts() + .connect_timeout(Duration::from_secs(10)); + let client_config = client_config_builder + .build() + .expect("Failed to set up tor client config."); let client_handle = TorClient::create_bootstrapped(client_config.clone()); log::info!("Bootstrapping Tor client..."); // Including print to stdout since TUI doesn't appear until Tor is started. print!("Bootstrapping Tor client..."); std::io::stdout().flush().expect("Failed to flush stdout."); - let mut client = client_handle.await.expect("Failed to bootstrap tor client."); + let mut client = client_handle + .await + .expect("Failed to bootstrap tor client."); log::info!("Bootstrapping Tor client done."); println!("done."); @@ -55,14 +73,22 @@ pub async fn create_tor_client(config_dir: &Path, nickname: String) -> (TorClien client.set_stream_prefs(stream_prefs); let hs_nickname = HsNickname::new(nickname).expect("Failed to create tor onion nickname."); - let onion_service_config = OnionServiceConfigBuilder::default().nickname(hs_nickname).build().expect("Failed to build tor onion service config."); - let (onion_service, request_stream) = client.launch_onion_service(onion_service_config).expect("Failed to launch tor onion service"); + let onion_service_config = OnionServiceConfigBuilder::default() + .nickname(hs_nickname) + .build() + .expect("Failed to build tor onion service config."); + let (onion_service, request_stream) = client + .launch_onion_service(onion_service_config) + .expect("Failed to launch tor onion service"); (client, onion_service, request_stream) } - -pub async fn forward_onion_connections(request_stream: impl Stream + std::marker::Unpin, connection_manager: Arc, display_name: String) { +pub async fn forward_onion_connections( + request_stream: impl Stream + std::marker::Unpin, + connection_manager: Arc, + display_name: String, +) { log::info!("New onion service started."); let stream_requests = tor_hsservice::handle_rend_requests(request_stream); tokio::pin!(stream_requests); @@ -78,13 +104,24 @@ pub async fn forward_onion_connections(request_stream: impl Stream, display_name: String + stream_request: StreamRequest, + connection_manager: Arc, + display_name: String, ) { match stream_request.request() { &IncomingStreamRequest::Begin(ref begin) if begin.port() == COORDINATOR_PORT => { - let onion_service_stream = stream_request.accept(Connected::new_empty()).await.unwrap(); - let io = TokioIo::new(onion_service_stream); - http1::Builder::new().serve_connection(io, service_fn(|request| serve(request, connection_manager.clone(), display_name.clone()))).await.unwrap(); + if let Ok(onion_service_stream) = stream_request.accept(Connected::new_empty()).await { + let io = TokioIo::new(onion_service_stream); + if let Ok(_) = http1::Builder::new() + .serve_connection( + io, + service_fn(|request| { + serve(request, connection_manager.clone(), display_name.clone()) + }), + ) + .await + {} + } } _ => { log::debug!("Received request to onion service on wrong port."); @@ -93,7 +130,11 @@ async fn handle_stream_request( } } -async fn serve(request: hyper::Request, connection_manager: Arc, display_name: String) -> Result>, http::Error> { +async fn serve( + request: hyper::Request, + connection_manager: Arc, + display_name: String, +) -> Result>, http::Error> { let path = request.uri().path(); log::debug!("Path: {path}"); @@ -105,10 +146,10 @@ async fn serve(request: hyper::Request, connection_manage match (request.method(), first) { (&Method::GET, "hello") => { let name = parts.next().ok_or(()).unwrap(); - Ok(hyper::Response::new(Full::new(Bytes::from( - format!("Hello, {name}!") - )))) - }, + Ok(hyper::Response::new(Full::new(Bytes::from(format!( + "Hello, {name}!" + ))))) + } (&Method::GET, "info") => { log::debug!("Path info."); let info = AugmentedInfo { @@ -116,10 +157,8 @@ async fn serve(request: hyper::Request, connection_manage display_name, }; let json = serde_json::to_string(&info).unwrap(); - Ok(hyper::Response::new(Full::new(Bytes::from( - json - )))) - }, + Ok(hyper::Response::new(Full::new(Bytes::from(json)))) + } _ => { // TODO: probably something better to respond with then empty bytes. log::error!("Unexpected path."); @@ -128,4 +167,4 @@ async fn serve(request: hyper::Request, connection_manage Ok(not_found) } } -} \ No newline at end of file +}