command socktes na djjhbfxvr5d ghgd9k7bvtx h

Cargo.lock

@@ -424,6 +424,16 @@ dependencies = [
  "cipher",
 ]
 
+[[package]]
+name = "ctrlc"
+version = "3.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbcf33c2a618cbe41ee43ae6e9f2e48368cd9f9db2896f10167d8d762679f639"
+dependencies = [
+ "nix",
+ "windows-sys 0.45.0",
+]
+
 [[package]]
 name = "curve25519-dalek"
 version = "3.2.0"
@@ -1355,6 +1365,18 @@ dependencies = [
  "tempfile",
 ]
 
+[[package]]
+name = "nix"
+version = "0.26.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bfdda3d196821d6af13126e40375cdf7da646a96114af134d5f417a9a1dc8e1a"
+dependencies = [
+ "bitflags",
+ "cfg-if",
+ "libc",
+ "static_assertions",
+]
+
 [[package]]
 name = "nom"
 version = "7.1.3"
@@ -2299,6 +2321,12 @@ dependencies = [
  "loom",
 ]
 
+[[package]]
+name = "static_assertions"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
+
 [[package]]
 name = "stringprep"
 version = "0.1.2"
@@ -2393,6 +2421,7 @@ name = "tfclient"
 version = "0.1.0"
 dependencies = [
  "clap",
+ "ctrlc",
  "dirs 5.0.0",
  "flate2",
  "hex",

tfclient/Cargo.toml

@@ -18,6 +18,7 @@ url = "2.3.1"
 toml = "0.7.3"
 serde = { version = "1.0.158", features = ["derive"] }
 serde_json = "1.0.94"
+ctrlc = "3.2.5"
 
 [build-dependencies]
 serde = { version = "1.0.157", features = ["derive"] }

tfclient/src/apiworker.rs

@@ -1,11 +1,32 @@
-use std::sync::mpsc::Receiver;
+use std::sync::mpsc::{Receiver, TryRecvError};
+use log::{error, info};
 use crate::config::TFClientConfig;
 use crate::daemon::ThreadMessageSender;
 
 pub enum APIWorkerMessage {
-
+    Shutdown
 }
 
-pub fn apiworker_main(_config: TFClientConfig, _transmitters: ThreadMessageSender, _rx: Receiver<APIWorkerMessage>) {
+pub fn apiworker_main(_config: TFClientConfig, _transmitters: ThreadMessageSender, rx: Receiver<APIWorkerMessage>) {
-
+    loop {
+        match rx.try_recv() {
+            Ok(msg) => {
+                match msg {
+                    APIWorkerMessage::Shutdown => {
+                        info!("recv on command socket: shutdown, stopping");
+                        break;
+                    }
+                }
+            },
+            Err(e) => {
+                match e {
+                    TryRecvError::Empty => {}
+                    TryRecvError::Disconnected => {
+                        error!("apiworker command socket disconnected, shutting down to prevent orphaning");
+                        break;
+                    }
+                }
+            }
+        }
+    }
 }

tfclient/src/config.rs

@@ -10,7 +10,7 @@ fn default_port() -> u16 { DEFAULT_PORT }
 #[derive(Serialize, Deserialize, Clone)]
 pub struct TFClientConfig {
     #[serde(default = "default_port")]
-    listen_port: u16
+    pub listen_port: u16
 }
 
 pub fn create_config(instance: &str) -> Result<(), Box<dyn Error>> {

tfclient/src/daemon.rs

@@ -5,6 +5,7 @@ use log::{error, info};
 
 use crate::apiworker::{apiworker_main, APIWorkerMessage};
 use crate::config::load_config;
+use crate::main;
 use crate::nebulaworker::{nebulaworker_main, NebulaWorkerMessage};
 use crate::socketworker::{socketworker_main, SocketWorkerMessage};
 use crate::util::check_server_url;
@@ -22,7 +23,7 @@ pub fn daemon_main(name: String, server: String) {
         }
     };
 
-    info!("Starting API thread...");
+    info!("Creating transmitter");
 
     let (tx_api, rx_api) = mpsc::channel::<APIWorkerMessage>();
     let (tx_socket, rx_socket) = mpsc::channel::<SocketWorkerMessage>();
@@ -34,6 +35,42 @@ pub fn daemon_main(name: String, server: String) {
         nebula_thread: tx_nebula
     };
 
+    let mainthread_transmitter = transmitter.clone();
+
+    info!("Setting signal trap...");
+
+    match ctrlc::set_handler(move || {
+        info!("Ctrl-C detected. Stopping threads...");
+        match mainthread_transmitter.nebula_thread.send(NebulaWorkerMessage::Shutdown) {
+            Ok(_) => (),
+            Err(e) => {
+                error!("Error sending shutdown message to nebula worker thread: {}", e);
+            }
+        }
+        match mainthread_transmitter.api_thread.send(APIWorkerMessage::Shutdown) {
+            Ok(_) => (),
+            Err(e) => {
+                error!("Error sending shutdown message to api worker thread: {}", e);
+            }
+        }
+        match mainthread_transmitter.socket_thread.send(SocketWorkerMessage::Shutdown) {
+            Ok(_) => (),
+            Err(e) => {
+                error!("Error sending shutdown message to socket worker thread: {}", e);
+            }
+        }
+    }) {
+        Ok(_) => (),
+        Err(e) => {
+            error!("Unable to set sigtrap: {}", e);
+            std::process::exit(1);
+        }
+    }
+
+    info!("Starting API thread...");
+
+
+
     let config_api = config.clone();
     let transmitter_api = transmitter.clone();
     let api_thread = thread::spawn(move || {
@@ -60,6 +97,7 @@ pub fn daemon_main(name: String, server: String) {
             std::process::exit(1);
         }
     }
+    info!("Socket thread exited");
 
     info!("Waiting for API thread to exit...");
     match api_thread.join() {
@@ -69,6 +107,7 @@ pub fn daemon_main(name: String, server: String) {
             std::process::exit(1);
         }
     }
+    info!("API thread exited");
 
     info!("Waiting for Nebula thread to exit...");
     match nebula_thread.join() {
@@ -78,13 +117,14 @@ pub fn daemon_main(name: String, server: String) {
             std::process::exit(1);
         }
     }
+    info!("Nebula thread exited");
 
     info!("All threads exited");
 }
 
 #[derive(Clone)]
 pub struct ThreadMessageSender {
-    socket_thread: Sender<SocketWorkerMessage>,
+    pub socket_thread: Sender<SocketWorkerMessage>,
-    api_thread: Sender<APIWorkerMessage>,
+    pub api_thread: Sender<APIWorkerMessage>,
-    nebula_thread: Sender<NebulaWorkerMessage>
+    pub nebula_thread: Sender<NebulaWorkerMessage>
 }

tfclient/src/nebulaworker.rs

@@ -1,13 +1,34 @@
 // Code to handle the nebula worker
 
-use std::sync::mpsc::Receiver;
+use std::sync::mpsc::{Receiver, TryRecvError};
+use log::{error, info};
 use crate::config::TFClientConfig;
 use crate::daemon::ThreadMessageSender;
 
 pub enum NebulaWorkerMessage {
-
+    Shutdown
 }
 
-pub fn nebulaworker_main(_config: TFClientConfig, _transmitter: ThreadMessageSender, _rx: Receiver<NebulaWorkerMessage>) {
+pub fn nebulaworker_main(_config: TFClientConfig, _transmitter: ThreadMessageSender, rx: Receiver<NebulaWorkerMessage>) {
-    
+    loop {
+        match rx.try_recv() {
+            Ok(msg) => {
+                match msg {
+                    NebulaWorkerMessage::Shutdown => {
+                        info!("recv on command socket: shutdown, stopping");
+                        break;
+                    }
+                }
+            },
+            Err(e) => {
+                match e {
+                    TryRecvError::Empty => {}
+                    TryRecvError::Disconnected => {
+                        error!("nebulaworker command socket disconnected, shutting down to prevent orphaning");
+                        break;
+                    }
+                }
+            }
+        }
+    }
 }

tfclient/src/socketworker.rs

@@ -1,13 +1,259 @@
 // Code to handle the nebula worker
 
-use std::sync::mpsc::Receiver;
+use std::error::Error;
+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 log::{debug, error, info, trace, warn};
+use serde::{Deserialize, Serialize};
+use crate::apiworker::APIWorkerMessage;
 use crate::config::TFClientConfig;
 use crate::daemon::ThreadMessageSender;
+use crate::nebulaworker::NebulaWorkerMessage;
 
 pub enum SocketWorkerMessage {
-
+    Shutdown
 }
 
-pub fn socketworker_main(_config: TFClientConfig, _transmitter: ThreadMessageSender, _rx: Receiver<SocketWorkerMessage>) {
+pub fn socketworker_main(config: TFClientConfig, transmitter: ThreadMessageSender, rx: Receiver<SocketWorkerMessage>) {
+    info!("socketworker_main called, entering realmain");
+    match _main(config, transmitter, rx) {
+        Ok(_) => (),
+        Err(e) => {
+            error!("Error in socket thread: {}", e);
+        }
+    };
+}
 
+fn _main(config: TFClientConfig, 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)?;
+
+    loop {
+        match listener.accept() {
+            Ok(stream) => {
+                let transmitter_clone = transmitter.clone();
+                thread::spawn(|| {
+                    match handle_stream(stream, transmitter_clone) {
+                        Ok(_) => (),
+                        Err(e) => {
+                            error!("Error in client thread: {}", e);
+                        }
+                    }
+                });
+            },
+            Err(e) if e.kind() == io::ErrorKind::WouldBlock => (),
+            Err(e) => { Err(e)?; }
+        }
+
+        match rx.try_recv() {
+            Ok(msg) => {
+                match msg {
+                    SocketWorkerMessage::Shutdown => {
+                        info!("recv on command socket: shutdown, stopping");
+                        break;
+                    }
+                }
+            },
+            Err(e) => {
+                match e {
+                    TryRecvError::Empty => {}
+                    TryRecvError::Disconnected => {
+                        error!("socketworker command socket disconnected, shutting down to prevent orphaning");
+                        break;
+                    }
+                }
+            }
+        }
+    }
+
+    Ok(())
+}
+
+fn handle_stream(stream: (TcpStream, SocketAddr), transmitter: ThreadMessageSender) -> Result<(), io::Error> {
+    info!("Incoming client");
+    match handle_client(stream.0, transmitter) {
+        Ok(()) => (),
+        Err(e) if e.kind() == io::ErrorKind::TimedOut => {
+            warn!("Client timed out, connection aborted");
+        },
+        Err(e) if e.kind() == io::ErrorKind::NotConnected => {
+            warn!("Client connection severed");
+        },
+        Err(e) if e.kind() == io::ErrorKind::BrokenPipe => {
+            warn!("Client connection returned error: broken pipe");
+        },
+        Err(e) if e.kind() == io::ErrorKind::ConnectionAborted => {
+            warn!("Client aborted connection");
+        },
+        Err(e) => {
+            error!("Error in client handler: {}", e);
+            return Err(e);
+        }
+    };
+    Ok(())
+}
+
+fn handle_client(stream: TcpStream, transmitter: ThreadMessageSender) -> 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
+    };
+
+    loop {
+        let mut command = String::new();
+        client.reader.read_line(&mut command)?;
+
+        let command: JsonMessage = serde_json::from_str(&command)?;
+
+        trace!("recv {:?} from {}", command, client.stream.peer_addr()?);
+
+        let should_disconnect;
+
+        match client.state {
+            ClientState::WaitHello => {
+                should_disconnect = waithello_handle(&mut client, &transmitter, command)?;
+            }
+            ClientState::SentHello => {
+                should_disconnect = senthello_handle(&mut client, &transmitter, command)?;
+            }
+        }
+
+        if should_disconnect { break; }
+    }
+
+    // Gracefully close the connection
+    stream.shutdown(Shutdown::Both)?;
+
+    Ok(())
+}
+
+struct Client<'a> {
+    state: ClientState,
+    reader: BufReader<&'a TcpStream>,
+    writer: BufWriter<&'a TcpStream>,
+    stream: &'a TcpStream
+}
+
+fn waithello_handle(client: &mut Client, _transmitter: &ThreadMessageSender, command: JsonMessage) -> Result<bool, io::Error> {
+    trace!("state: WaitHello, handing with waithello_handle");
+    let mut should_disconnect = false;
+
+    match command {
+        JsonMessage::Hello { version } => {
+            if version != JSON_API_VERSION {
+                should_disconnect = true;
+                client.stream.write_all(&ctob(JsonMessage::Goodbye {
+                    reason: DisconnectReason::UnsupportedVersion {
+                        expected: JSON_API_VERSION,
+                        got: version
+                    }
+                }))?;
+            }
+            client.stream.write_all(&ctob(JsonMessage::Hello {
+                version: JSON_API_VERSION
+            }))?;
+            client.state = ClientState::SentHello;
+            trace!("setting state to SentHello");
+        },
+        JsonMessage::Goodbye { reason } => {
+            info!("Client sent disconnect: {:?}", reason);
+            should_disconnect = true;
+        },
+        _ => {
+            debug!("message type unexpected in WaitHello state");
+            should_disconnect = true;
+            client.stream.write_all(&ctob(JsonMessage::Goodbye {
+                reason: DisconnectReason::UnexpectedMessageType,
+            }))?;
+        }
+    }
+
+    Ok(should_disconnect)
+}
+
+fn senthello_handle(client: &mut Client, transmitter: &ThreadMessageSender, command: JsonMessage) -> Result<bool, io::Error> {
+    trace!("state: SentHello, handing with senthello_handle");
+    let mut should_disconnect = false;
+
+    match command {
+        JsonMessage::Goodbye { reason } => {
+            info!("Client sent disconnect: {:?}", reason);
+            should_disconnect = true;
+        },
+
+        JsonMessage::Shutdown {} => {
+            info!("Requested to shutdown by local control socket. Sending shutdown message to threads");
+            match transmitter.nebula_thread.send(NebulaWorkerMessage::Shutdown) {
+                Ok(_) => (),
+                Err(e) => {
+                    error!("Error sending shutdown message to nebula worker thread: {}", e);
+                }
+            }
+            match transmitter.api_thread.send(APIWorkerMessage::Shutdown) {
+                Ok(_) => (),
+                Err(e) => {
+                    error!("Error sending shutdown message to api worker thread: {}", e);
+                }
+            }
+            match transmitter.socket_thread.send(SocketWorkerMessage::Shutdown) {
+                Ok(_) => (),
+                Err(e) => {
+                    error!("Error sending shutdown message to socket worker thread: {}", e);
+                }
+            }
+        }
+
+        _ => {
+            debug!("message type unexpected in SentHello state");
+            should_disconnect = true;
+            client.stream.write_all(&ctob(JsonMessage::Goodbye {
+                reason: DisconnectReason::UnexpectedMessageType,
+            }))?;
+        }
+    }
+
+    Ok(should_disconnect)
+}
+
+fn ctob(command: JsonMessage) -> Vec<u8> {
+    let command_str = serde_json::to_string(&command).unwrap() + "\n";
+    command_str.into_bytes()
+}
+
+enum ClientState {
+    WaitHello,
+    SentHello
+}
+
+pub const JSON_API_VERSION: i32 = 1;
+
+#[derive(Serialize, Deserialize, Debug)]
+#[serde(tag = "method")]
+enum JsonMessage {
+    #[serde(rename = "hello")]
+    Hello {
+        version: i32
+    },
+    #[serde(rename = "goodbye")]
+    Goodbye {
+        reason: DisconnectReason
+    },
+    #[serde(rename = "shutdown")]
+    Shutdown {}
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+enum DisconnectReason {
+    #[serde(rename = "unsupported_version")]
+    UnsupportedVersion { expected: i32, got: i32 },
+    #[serde(rename = "unexpected_message_type")]
+    UnexpectedMessageType,
+    #[serde(rename = "done")]
+    Done
 }