piuszgdpiugdszfoigfxliudxfg fixup thing and stuff and consistentifiy thhhhhhh

This commit is contained in:
c0repwn3r 2023-03-23 13:50:21 -04:00
parent 76dc3fc477
commit 990758c27e
Signed by: core
GPG Key ID: FDBF740DADDCEECF
4 changed files with 97 additions and 13 deletions

View File

@ -2,7 +2,7 @@ use std::error::Error;
use std::fs; use std::fs;
use log::{debug, info}; use log::{debug, info};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::dirs::{get_config_dir, get_config_file}; use crate::dirs::{get_cdata_dir, get_cdata_file, get_config_dir, get_config_file, get_data_dir};
pub const DEFAULT_PORT: u16 = 8157; pub const DEFAULT_PORT: u16 = 8157;
fn default_port() -> u16 { DEFAULT_PORT } fn default_port() -> u16 { DEFAULT_PORT }
@ -13,6 +13,11 @@ pub struct TFClientConfig {
pub listen_port: u16 pub listen_port: u16
} }
#[derive(Serialize, Deserialize, Clone)]
pub struct TFClientData {
pub host_id: Option<String>
}
pub fn create_config(instance: &str) -> Result<(), Box<dyn Error>> { pub fn create_config(instance: &str) -> Result<(), Box<dyn Error>> {
info!("Creating config directory..."); info!("Creating config directory...");
fs::create_dir_all(get_config_dir(instance).ok_or("Unable to load config dir")?)?; fs::create_dir_all(get_config_dir(instance).ok_or("Unable to load config dir")?)?;
@ -39,4 +44,46 @@ pub fn load_config(instance: &str) -> Result<TFClientConfig, Box<dyn Error>> {
let config: TFClientConfig = toml::from_str(&config_str)?; let config: TFClientConfig = toml::from_str(&config_str)?;
info!("Loaded config successfully"); info!("Loaded config successfully");
Ok(config) Ok(config)
}
pub fn create_cdata(instance: &str) -> Result<(), Box<dyn Error>> {
info!("Creating data directory...");
fs::create_dir_all(get_cdata_dir(instance).ok_or("Unable to load data dir")?)?;
info!("Copying default data file to config directory...");
let config = TFClientData { host_id: None };
let config_str = toml::to_string(&config)?;
fs::write(get_cdata_file(instance).ok_or("Unable to load data dir")?, config_str)?;
Ok(())
}
pub fn load_cdata(instance: &str) -> Result<TFClientData, Box<dyn Error>> {
info!("Loading cdata...");
let config_file = get_cdata_file(instance).ok_or("Unable to load cdata dir")?;
if !config_file.exists() {
create_cdata(instance)?;
}
debug!("opening {}", config_file.as_path().display());
let config_str = fs::read_to_string(config_file)?;
debug!("parsing cdata file");
let config: TFClientData = toml::from_str(&config_str)?;
info!("Loaded cdata successfully");
Ok(config)
}
pub fn save_cdata(instance: &str, data: TFClientData) -> Result<(), Box<dyn Error>> {
info!("Saving cdata...");
let config_file = get_cdata_file(instance).ok_or("Unable to load cdata dir")?;
if !config_file.exists() {
create_cdata(instance)?;
}
debug!("serializing cdata file");
let config: String = toml::to_string(&data)?;
debug!("writing to {}", config_file.as_path().display());
fs::write(config_file, config)?;
info!("Saved cdata successfully");
Ok(())
} }

View File

@ -86,7 +86,7 @@ pub fn daemon_main(name: String, server: String) {
info!("Starting socket worker thread..."); info!("Starting socket worker thread...");
let socket_thread = thread::spawn(move || { let socket_thread = thread::spawn(move || {
socketworker_main(config, transmitter, rx_socket); socketworker_main(config, name.clone(), transmitter, rx_socket);
}); });
info!("Waiting for socket thread to exit..."); info!("Waiting for socket thread to exit...");

View File

@ -10,4 +10,12 @@ pub fn get_config_dir(instance: &str) -> Option<PathBuf> {
pub fn get_config_file(instance: &str) -> Option<PathBuf> { pub fn get_config_file(instance: &str) -> Option<PathBuf> {
get_config_dir(instance).map(|f| f.join("tfclient.toml")) get_config_dir(instance).map(|f| f.join("tfclient.toml"))
}
pub fn get_cdata_dir(instance: &str) -> Option<PathBuf> {
dirs::config_dir().map(|f| f.join("tfclient_data/").join(format!("{}/", instance)))
}
pub fn get_cdata_file(instance: &str) -> Option<PathBuf> {
get_cdata_dir(instance).map(|f| f.join("tfclient.toml"))
} }

View File

@ -5,10 +5,11 @@ use std::{io, thread};
use std::io::{BufRead, BufReader, BufWriter, Read, Write}; use std::io::{BufRead, BufReader, BufWriter, Read, Write};
use std::net::{IpAddr, Shutdown, SocketAddr, TcpListener, TcpStream}; use std::net::{IpAddr, Shutdown, SocketAddr, TcpListener, TcpStream};
use std::sync::mpsc::{Receiver, TryRecvError}; use std::sync::mpsc::{Receiver, TryRecvError};
use clap::builder::Str;
use log::{debug, error, info, trace, warn}; use log::{debug, error, info, trace, warn};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::apiworker::APIWorkerMessage; use crate::apiworker::APIWorkerMessage;
use crate::config::TFClientConfig; use crate::config::{load_cdata, TFClientConfig};
use crate::daemon::ThreadMessageSender; use crate::daemon::ThreadMessageSender;
use crate::nebulaworker::NebulaWorkerMessage; use crate::nebulaworker::NebulaWorkerMessage;
@ -16,9 +17,9 @@ pub enum SocketWorkerMessage {
Shutdown Shutdown
} }
pub fn socketworker_main(config: TFClientConfig, transmitter: ThreadMessageSender, rx: Receiver<SocketWorkerMessage>) { pub fn socketworker_main(config: TFClientConfig, instance: String, transmitter: ThreadMessageSender, rx: Receiver<SocketWorkerMessage>) {
info!("socketworker_main called, entering realmain"); info!("socketworker_main called, entering realmain");
match _main(config, transmitter, rx) { match _main(config, instance, transmitter, rx) {
Ok(_) => (), Ok(_) => (),
Err(e) => { Err(e) => {
error!("Error in socket thread: {}", e); error!("Error in socket thread: {}", e);
@ -26,7 +27,7 @@ pub fn socketworker_main(config: TFClientConfig, transmitter: ThreadMessageSende
}; };
} }
fn _main(config: TFClientConfig, transmitter: ThreadMessageSender, rx: Receiver<SocketWorkerMessage>) -> Result<(), Box<dyn Error>> { fn _main(config: TFClientConfig, instance: String, transmitter: ThreadMessageSender, rx: Receiver<SocketWorkerMessage>) -> Result<(), Box<dyn Error>> {
let listener = TcpListener::bind(SocketAddr::new(IpAddr::from([127, 0, 0, 1]), config.listen_port))?; let listener = TcpListener::bind(SocketAddr::new(IpAddr::from([127, 0, 0, 1]), config.listen_port))?;
listener.set_nonblocking(true)?; listener.set_nonblocking(true)?;
@ -34,8 +35,10 @@ fn _main(config: TFClientConfig, transmitter: ThreadMessageSender, rx: Receiver<
match listener.accept() { match listener.accept() {
Ok(stream) => { Ok(stream) => {
let transmitter_clone = transmitter.clone(); let transmitter_clone = transmitter.clone();
let config_clone = config.clone();
let instance_clone = instance.clone();
thread::spawn(|| { thread::spawn(|| {
match handle_stream(stream, transmitter_clone) { match handle_stream(stream, transmitter_clone, config_clone, instance_clone) {
Ok(_) => (), Ok(_) => (),
Err(e) => { Err(e) => {
error!("Error in client thread: {}", e); error!("Error in client thread: {}", e);
@ -71,9 +74,9 @@ fn _main(config: TFClientConfig, transmitter: ThreadMessageSender, rx: Receiver<
Ok(()) Ok(())
} }
fn handle_stream(stream: (TcpStream, SocketAddr), transmitter: ThreadMessageSender) -> Result<(), io::Error> { fn handle_stream(stream: (TcpStream, SocketAddr), transmitter: ThreadMessageSender, config: TFClientConfig, instance: String) -> Result<(), io::Error> {
info!("Incoming client"); info!("Incoming client");
match handle_client(stream.0, transmitter) { match handle_client(stream.0, transmitter, config, instance) {
Ok(()) => (), Ok(()) => (),
Err(e) if e.kind() == io::ErrorKind::TimedOut => { Err(e) if e.kind() == io::ErrorKind::TimedOut => {
warn!("Client timed out, connection aborted"); warn!("Client timed out, connection aborted");
@ -95,14 +98,16 @@ fn handle_stream(stream: (TcpStream, SocketAddr), transmitter: ThreadMessageSend
Ok(()) Ok(())
} }
fn handle_client(stream: TcpStream, transmitter: ThreadMessageSender) -> Result<(), io::Error> { fn handle_client(stream: TcpStream, transmitter: ThreadMessageSender, config: TFClientConfig, instance: String) -> Result<(), io::Error> {
info!("Handling connection from {}", stream.peer_addr()?); info!("Handling connection from {}", stream.peer_addr()?);
let mut client = Client { let mut client = Client {
state: ClientState::WaitHello, state: ClientState::WaitHello,
reader: BufReader::new(&stream), reader: BufReader::new(&stream),
writer: BufWriter::new(&stream), writer: BufWriter::new(&stream),
stream: &stream stream: &stream,
config,
instance,
}; };
loop { loop {
@ -137,7 +142,9 @@ struct Client<'a> {
state: ClientState, state: ClientState,
reader: BufReader<&'a TcpStream>, reader: BufReader<&'a TcpStream>,
writer: BufWriter<&'a TcpStream>, writer: BufWriter<&'a TcpStream>,
stream: &'a TcpStream stream: &'a TcpStream,
config: TFClientConfig,
instance: String
} }
fn waithello_handle(client: &mut Client, _transmitter: &ThreadMessageSender, command: JsonMessage) -> Result<bool, io::Error> { fn waithello_handle(client: &mut Client, _transmitter: &ThreadMessageSender, command: JsonMessage) -> Result<bool, io::Error> {
@ -207,6 +214,20 @@ fn senthello_handle(client: &mut Client, transmitter: &ThreadMessageSender, comm
error!("Error sending shutdown message to socket worker thread: {}", e); error!("Error sending shutdown message to socket worker thread: {}", e);
} }
} }
},
JsonMessage::GetHostID {} => {
let data = match load_cdata(&client.instance) {
Ok(d) => d,
Err(e) => {
error!("Error loading cdata: {}", e);
panic!("{}", e); // TODO: Find a better way of handling this
}
};
client.stream.write_all(&ctob(JsonMessage::HostID {
has_id: data.host_id.is_some(),
id: data.host_id
}))?;
} }
_ => { _ => {
@ -245,10 +266,18 @@ enum JsonMessage {
reason: DisconnectReason reason: DisconnectReason
}, },
#[serde(rename = "shutdown")] #[serde(rename = "shutdown")]
Shutdown {} Shutdown {},
#[serde(rename = "get_host_id")]
GetHostID {},
#[serde(rename = "host_id")]
HostID {
has_id: bool,
id: Option<String>
}
} }
#[derive(Serialize, Deserialize, Debug)] #[derive(Serialize, Deserialize, Debug)]
#[serde(tag = "type")]
enum DisconnectReason { enum DisconnectReason {
#[serde(rename = "unsupported_version")] #[serde(rename = "unsupported_version")]
UnsupportedVersion { expected: i32, got: i32 }, UnsupportedVersion { expected: i32, got: i32 },