From b04588ee6cf2325c00125c7959e89ba2e43591f1 Mon Sep 17 00:00:00 2001 From: Jeremy Andrews Date: Mon, 15 May 2023 14:44:39 +0200 Subject: [PATCH] only expect hash for worker messages --- src/controller.rs | 59 ++++++++++++++++++++++---------------------- src/gaggle/worker.rs | 2 +- tests/controller.rs | 8 +++--- 3 files changed, 35 insertions(+), 34 deletions(-) diff --git a/src/controller.rs b/src/controller.rs index 7fc0a5d7..5da77225 100644 --- a/src/controller.rs +++ b/src/controller.rs @@ -1621,38 +1621,39 @@ impl ControllerState { if let Ok(command_string) = self.get_command_string(buf).await { // Extract the command and value in a generic way. if let Ok(request_message) = self.get_match(command_string.trim()).await { - let hash = if let Some(ControllerValue::Text(hash)) = - request_message.value.as_ref() - { - // Clone the value. - hash.to_string() - } else { - unreachable!("Hash must exist, enforced by regex"); - }; // Workers using Telnet socket to connect to the Manager. if request_message.command == ControllerCommand::WorkerConnect { info!("Worker instance connecting ..."); - if request_message.command == ControllerCommand::WorkerConnect { - // A Worker is trying to connect, send the connection to the Parent. - if self - .channel_tx - .try_send(ControllerRequest { - response_channel: None, - client_id: self.thread_id, - request: ControllerRequestMessage { - command: ControllerCommand::WorkerConnect, - value: Some(ControllerValue::Socket( - WorkerConnection { hash, socket }, - )), - }, - }) - .is_err() - { - warn!("failed to send Worker socket to parent"); - }; - break; - // ELSE? - } + // Workers must include hash when connecting. + let hash = if let Some(ControllerValue::Text(hash)) = + request_message.value.as_ref() + { + // Clone the value. + hash.to_string() + } else { + // @TODO: Don't panic, instead cancel the connection ... + panic!("Hash must exist, enforced by regex"); + }; + + // A Worker is trying to connect, send the connection to the Parent. + if self + .channel_tx + .try_send(ControllerRequest { + response_channel: None, + client_id: self.thread_id, + request: ControllerRequestMessage { + command: ControllerCommand::WorkerConnect, + value: Some(ControllerValue::Socket( + WorkerConnection { hash, socket }, + )), + }, + }) + .is_err() + { + warn!("failed to send Worker socket to parent"); + }; + break; + // ELSE? } // Act on the commmand received. diff --git a/src/gaggle/worker.rs b/src/gaggle/worker.rs index dc533fe0..6d62635e 100644 --- a/src/gaggle/worker.rs +++ b/src/gaggle/worker.rs @@ -120,7 +120,7 @@ impl GooseConfiguration { message: "manager_host", }, ]) - .unwrap_or_else(|| "".to_string()); + .unwrap_or_default(); // Set `manager_port` on Worker. self.manager_port = self diff --git a/tests/controller.rs b/tests/controller.rs index e7801346..5a254a12 100644 --- a/tests/controller.rs +++ b/tests/controller.rs @@ -206,8 +206,8 @@ async fn run_standalone_test(test_type: TestType) { if let Some(stream) = test_state.telnet_stream.as_mut() { let _ = match stream.read(&mut test_state.buf).await { Ok(data) => data, - Err(_) => { - panic!("ERROR: server disconnected!"); + Err(e) => { + panic!("ERROR: server disconnected: {}", e); } }; response = str::from_utf8(&test_state.buf).unwrap(); @@ -233,7 +233,7 @@ async fn run_standalone_test(test_type: TestType) { unreachable!(); } - //println!("{:?}: {}", test_state.command, response); + println!("{:?}: {}", test_state.command, response); match test_state.command { ControllerCommand::Exit => { match test_state.step { @@ -751,7 +751,7 @@ async fn update_state(test_state: Option, test_type: &TestType) -> Te } else { // Connect to telnet controller. let telnet_stream = match test_type { - TestType::Telnet => Some(TcpStream::connect("127.0.0.1:5116").await.unwrap()), + TestType::Telnet => Some(TcpStream::connect("127.0.0.1:5115").await.unwrap()), _ => None, };