added reconnect

This commit is contained in:
sad 2024-03-26 15:54:35 -06:00
parent 5d5ed9da8e
commit 6bf609c162
WARNING! Although there is a key with this ID in the database it does not verify this commit! This commit is SUSPICIOUS.
GPG Key ID: 28D3A882F3E6AD02
2 changed files with 46 additions and 25 deletions

View File

@ -10,7 +10,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
use colored::*; use colored::*;
use tokio_socks::tcp::Socks5Stream; use tokio_socks::tcp::Socks5Stream;
#[derive(Deserialize)] #[derive(Deserialize, Clone)]
struct Config { struct Config {
server: String, server: String,
port: u16, port: u16,
@ -22,6 +22,9 @@ struct Config {
sasl_password: Option<String>, sasl_password: Option<String>,
capabilities: Option<Vec<String>>, capabilities: Option<Vec<String>>,
reconnect_delay: u64,
reconnect_attempts: u64,
// Proxy // Proxy
use_proxy: bool, use_proxy: bool,
proxy_type: Option<String>, proxy_type: Option<String>,
@ -45,43 +48,62 @@ use mods::ascii::handle_ascii_command;
#[tokio::main(flavor = "multi_thread", worker_threads = 12)] #[tokio::main(flavor = "multi_thread", worker_threads = 12)]
async fn main() -> Result<(), Box<dyn std::error::Error>> { async fn main() -> Result<(), Box<dyn std::error::Error>> {
tokio::spawn(async move { //tokio::spawn(async move {
println!("Loading Config..."); println!("Loading Config...");
let config = loaded_config().expect("Error parsing config.toml"); let config = loaded_config().expect("Error parsing config.toml");
println!("Config loaded!"); println!("Config loaded!");
let mut reconnect_attempts = 0;
let server = format!("{}:{}", config.server, config.port); while reconnect_attempts < config.reconnect_attempts {
let configc = config.clone();
if config.use_ssl { let server = format!("{}:{}", configc.server, configc.port);
if config.use_proxy { let connection_result = tokio::spawn(async move {
let tcp_stream = proxy_exec(&config).await; let config = configc.clone();
match tcp_stream { if config.use_ssl {
Ok(tcp_stream) => { if config.use_proxy {
let tls_stream = tls_exec(&config, tcp_stream).await; let tcp_stream = proxy_exec(&config).await;
match tls_stream { match tcp_stream {
Ok(tls_stream) => { Ok(tcp_stream) => {
if let Err(e) = handler(tls_stream, config).await { let tls_stream = tls_exec(&config, tcp_stream).await;
println!("Error handling TLS connection: {}", e); match tls_stream {
Ok(tls_stream) => {
if let Err(e) = handler(tls_stream, config).await {
println!("Error handling TLS connection: {}", e);
}
},
Err(e) => {
println!("Error establishing TLS connection: {}", e);
}
} }
}, },
Err(e) => { Err(e) => {
println!("Error establishing TLS connection: {}", e); println!("Error connecting to proxy: {}", e);
} }
} }
}, } else {
Err(e) => { let tcp_stream = TcpStream::connect(server).await.expect("Error connecting to server");
println!("Error connecting to proxy: {}", e); let tls_stream = tls_exec(&config, tcp_stream).await.expect("Error establishing TLS connection");
handler(tls_stream, config).await.unwrap();
} }
} else {
println!("Non-SSL connection not implemented.");
}
Ok::<(), Box<dyn std::error::Error + Send>>(())
}).await.unwrap();
match connection_result {
Ok(_) => {
println!("Connection established successfully!");
reconnect_attempts = 0;
},
Err(e) => {
println!("Error handling connection: {}", e);
reconnect_attempts += 1;
tokio::time::sleep(tokio::time::Duration::from_secs(config.reconnect_delay)).await;
} }
} else {
let tcp_stream = TcpStream::connect(server).await.expect("Error connecting to server");
let tls_stream = tls_exec(&config, tcp_stream).await.expect("Error establishing TLS connection");
handler(tls_stream, config).await.unwrap();
} }
} else {
println!("Non-SSL connection not implemented.");
} }
}).await.unwrap(); println!("Reconnect attempts exceeded. Exiting...");
Ok(()) Ok(())
} }

View File

@ -27,7 +27,6 @@ async fn send_ansi_art<W: AsyncWriteExt + Unpin>(writer: &mut W, file_path: &str
let reader = BufReader::new(file); let reader = BufReader::new(file);
let mut lines = reader.lines(); let mut lines = reader.lines();
while let Some(line) = lines.next_line().await? { while let Some(line) = lines.next_line().await? {
if line.len() > CHUNK_SIZE { if line.len() > CHUNK_SIZE {