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 log::{debug, info};
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;
fn default_port() -> u16 { DEFAULT_PORT }
@ -13,6 +13,11 @@ pub struct TFClientConfig {
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>> {
info!("Creating config directory...");
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)?;
info!("Loaded config successfully");
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...");
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...");

View File

@ -10,4 +10,12 @@ pub fn get_config_dir(instance: &str) -> Option<PathBuf> {
pub fn get_config_file(instance: &str) -> Option<PathBuf> {
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::net::{IpAddr, Shutdown, SocketAddr, TcpListener, TcpStream};
use std::sync::mpsc::{Receiver, TryRecvError};
use clap::builder::Str;
use log::{debug, error, info, trace, warn};
use serde::{Deserialize, Serialize};
use crate::apiworker::APIWorkerMessage;
use crate::config::TFClientConfig;
use crate::config::{load_cdata, TFClientConfig};
use crate::daemon::ThreadMessageSender;
use crate::nebulaworker::NebulaWorkerMessage;
@ -16,9 +17,9 @@ pub enum SocketWorkerMessage {
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");
match _main(config, transmitter, rx) {
match _main(config, instance, transmitter, rx) {
Ok(_) => (),
Err(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))?;
listener.set_nonblocking(true)?;
@ -34,8 +35,10 @@ fn _main(config: TFClientConfig, transmitter: ThreadMessageSender, rx: Receiver<
match listener.accept() {
Ok(stream) => {
let transmitter_clone = transmitter.clone();
let config_clone = config.clone();
let instance_clone = instance.clone();
thread::spawn(|| {
match handle_stream(stream, transmitter_clone) {
match handle_stream(stream, transmitter_clone, config_clone, instance_clone) {
Ok(_) => (),
Err(e) => {
error!("Error in client thread: {}", e);
@ -71,9 +74,9 @@ fn _main(config: TFClientConfig, transmitter: ThreadMessageSender, rx: Receiver<
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");
match handle_client(stream.0, transmitter) {
match handle_client(stream.0, transmitter, config, instance) {
Ok(()) => (),
Err(e) if e.kind() == io::ErrorKind::TimedOut => {
warn!("Client timed out, connection aborted");
@ -95,14 +98,16 @@ fn handle_stream(stream: (TcpStream, SocketAddr), transmitter: ThreadMessageSend
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()?);
let mut client = Client {
state: ClientState::WaitHello,
reader: BufReader::new(&stream),
writer: BufWriter::new(&stream),
stream: &stream
stream: &stream,
config,
instance,
};
loop {
@ -137,7 +142,9 @@ struct Client<'a> {
state: ClientState,
reader: BufReader<&'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> {
@ -207,6 +214,20 @@ fn senthello_handle(client: &mut Client, transmitter: &ThreadMessageSender, comm
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
},
#[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)]
#[serde(tag = "type")]
enum DisconnectReason {
#[serde(rename = "unsupported_version")]
UnsupportedVersion { expected: i32, got: i32 },