This commit is contained in:
2025-05-23 15:28:30 +04:00
parent 92b9c356b8
commit 0e545e56de
144 changed files with 294 additions and 1907 deletions

View File

@@ -0,0 +1,117 @@
# This file is automatically @generated by Cargo.
# It is not intended for manual editing.
version = 4
[[package]]
name = "itoa"
version = "1.0.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c"
[[package]]
name = "memchr"
version = "2.7.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
[[package]]
name = "proc-macro2"
version = "1.0.94"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a31971752e70b8b2686d7e46ec17fb38dad4051d94024c88df49b667caea9c84"
dependencies = [
"unicode-ident",
]
[[package]]
name = "quote"
version = "1.0.40"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d"
dependencies = [
"proc-macro2",
]
[[package]]
name = "rustclients"
version = "0.1.0"
dependencies = [
"serde",
"serde_json",
"thiserror",
]
[[package]]
name = "ryu"
version = "1.0.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f"
[[package]]
name = "serde"
version = "1.0.219"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5f0e2c6ed6606019b4e29e69dbaba95b11854410e5347d525002456dbbb786b6"
dependencies = [
"serde_derive",
]
[[package]]
name = "serde_derive"
version = "1.0.219"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5b0276cf7f2c73365f7157c8123c21cd9a50fbbd844757af28ca1f5925fc2a00"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "serde_json"
version = "1.0.140"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "20068b6e96dc6c9bd23e01df8827e6c7e1f2fddd43c21810382803c136b99373"
dependencies = [
"itoa",
"memchr",
"ryu",
"serde",
]
[[package]]
name = "syn"
version = "2.0.100"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b09a44accad81e1ba1cd74a32461ba89dee89095ba17b32f5d03683b1b1fc2a0"
dependencies = [
"proc-macro2",
"quote",
"unicode-ident",
]
[[package]]
name = "thiserror"
version = "1.0.69"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52"
dependencies = [
"thiserror-impl",
]
[[package]]
name = "thiserror-impl"
version = "1.0.69"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "unicode-ident"
version = "1.0.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512"

View File

@@ -0,0 +1,13 @@
[package]
name = "rustclients"
version = "0.1.0"
edition = "2021"
[dependencies]
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
thiserror = "1.0"
[[example]]
name = "fakehandler_example"
path = "examples/fakehandler_example.rs"

View File

@@ -0,0 +1,111 @@
use std::time::Duration;
use std::thread;
use std::io::{Read, Write};
use std::os::unix::net::{UnixListener, UnixStream};
use std::fs;
// Import directly from the lib.rs
use rustclients::FakeHandlerClient;
use rustclients::Result;
// Simple mock server that handles Unix socket connections
fn start_mock_server(socket_path: &str) -> std::thread::JoinHandle<()> {
let socket_path = socket_path.to_string();
thread::spawn(move || {
// Remove the socket file if it exists
let _ = fs::remove_file(&socket_path);
// Create a Unix socket listener
let listener = match UnixListener::bind(&socket_path) {
Ok(listener) => listener,
Err(e) => {
println!("Failed to bind to socket {}: {}", socket_path, e);
return;
}
};
println!("Mock server listening on {}", socket_path);
// Accept connections and handle them
for stream in listener.incoming() {
match stream {
Ok(mut stream) => {
println!("Mock server: Accepted new connection");
// Read from the stream
let mut buffer = [0; 1024];
match stream.read(&mut buffer) {
Ok(n) => {
let request = String::from_utf8_lossy(&buffer[0..n]);
println!("Mock server received: {}", request);
// Send a welcome message first
let welcome = "Welcome to the mock server\n";
let _ = stream.write_all(welcome.as_bytes());
// Send a response
let response = "OK: Command processed\n> ";
let _ = stream.write_all(response.as_bytes());
},
Err(e) => println!("Mock server error reading from stream: {}", e),
}
},
Err(e) => println!("Mock server error accepting connection: {}", e),
}
}
})
}
fn main() -> Result<()> {
// Define the socket path
let socket_path = "/tmp/heroagent/test.sock";
// Start the mock server
println!("Starting mock server...");
let server_handle = start_mock_server(socket_path);
// Give the server time to start
thread::sleep(Duration::from_millis(500));
// Initialize the client
let client = FakeHandlerClient::new(socket_path)
.with_timeout(Duration::from_secs(5));
println!("\n--- Test 1: Making first request ---");
// This should open a new connection
match client.return_success(Some("Test 1")) {
Ok(response) => println!("Response: {}", response),
Err(e) => println!("Error: {}", e),
}
// Wait a moment
thread::sleep(Duration::from_millis(500));
println!("\n--- Test 2: Making second request ---");
// This should open another new connection
match client.return_success(Some("Test 2")) {
Ok(response) => println!("Response: {}", response),
Err(e) => println!("Error: {}", e),
}
// Wait a moment
thread::sleep(Duration::from_millis(500));
println!("\n--- Test 3: Making third request ---");
// This should open yet another new connection
match client.return_success(Some("Test 3")) {
Ok(response) => println!("Response: {}", response),
Err(e) => println!("Error: {}", e),
}
println!("\nTest completed. Check the debug output to verify that a new connection was opened for each request.");
// Clean up
let _ = fs::remove_file(socket_path);
// Wait for the server to finish (in a real application, you might want to signal it to stop)
println!("Waiting for server to finish...");
// In a real application, we would join the server thread here
Ok(())
}

View File

@@ -0,0 +1,100 @@
use std::time::Duration;
// Import directly from the lib.rs
use rustclients::FakeHandlerClient;
use rustclients::Result;
fn main() -> Result<()> {
// Create a new fake handler client
// Replace with the actual socket path used in your environment
let socket_path = "/tmp/heroagent/fakehandler.sock";
// Initialize the client with a timeout
let client = FakeHandlerClient::new(socket_path)
.with_timeout(Duration::from_secs(5));
println!("Connecting to fake handler at {}", socket_path);
// Connect to the server
match client.connect() {
Ok(_) => println!("Successfully connected to fake handler"),
Err(e) => {
eprintln!("Failed to connect: {}", e);
eprintln!("Make sure the fake handler server is running and the socket path is correct");
return Err(e);
}
}
// Test various commands
// 1. Get help information
println!("\n--- Help Information ---");
match client.help() {
Ok(help) => println!("{}", help),
Err(e) => eprintln!("Error getting help: {}", e),
}
// 2. Return success message
println!("\n--- Success Message ---");
match client.return_success(Some("Custom success message")) {
Ok(response) => println!("Success response: {}", response),
Err(e) => eprintln!("Error getting success: {}", e),
}
// 3. Return JSON response
println!("\n--- JSON Response ---");
match client.return_json(Some("JSON message"), Some("success"), Some(200)) {
Ok(response) => println!("JSON response: {:?}", response),
Err(e) => eprintln!("Error getting JSON: {}", e),
}
// 4. Return error message (this will return a ClientError)
println!("\n--- Error Message ---");
match client.return_error(Some("Custom error message")) {
Ok(response) => println!("Error response (unexpected success): {}", response),
Err(e) => eprintln!("Expected error received: {}", e),
}
// 5. Return empty response
println!("\n--- Empty Response ---");
match client.return_empty() {
Ok(response) => println!("Empty response (length: {})", response.len()),
Err(e) => eprintln!("Error getting empty response: {}", e),
}
// 6. Return large response
println!("\n--- Large Response ---");
match client.return_large(Some(10)) {
Ok(response) => {
let lines: Vec<&str> = response.lines().collect();
println!("Large response (first 3 lines of {} total):", lines.len());
for i in 0..std::cmp::min(3, lines.len()) {
println!(" {}", lines[i]);
}
println!(" ...");
},
Err(e) => eprintln!("Error getting large response: {}", e),
}
// 7. Return invalid JSON (will cause a JSON parsing error)
println!("\n--- Invalid JSON ---");
match client.return_invalid_json() {
Ok(response) => println!("Invalid JSON response (unexpected success): {:?}", response),
Err(e) => eprintln!("Expected JSON error received: {}", e),
}
// 8. Return malformed error
println!("\n--- Malformed Error ---");
match client.return_malformed_error() {
Ok(response) => println!("Malformed error response: {}", response),
Err(e) => eprintln!("Error with malformed error: {}", e),
}
// Close the connection
println!("\nClosing connection");
client.close()?;
println!("Example completed successfully");
Ok(())
}

View File

@@ -0,0 +1,134 @@
use serde::{Deserialize, Serialize};
use std::time::Duration;
use crate::{Client, Result, ClientError};
/// Response from the fake handler
#[derive(Debug, Serialize, Deserialize, Default)]
pub struct FakeResponse {
#[serde(default)]
pub message: String,
#[serde(default)]
pub status: String,
#[serde(default)]
pub code: i32,
}
/// Client for the fake handler
pub struct FakeHandlerClient {
client: Client,
}
impl FakeHandlerClient {
/// Create a new fake handler client
pub fn new(socket_path: &str) -> Self {
Self {
client: Client::new(socket_path),
}
}
/// Set the connection timeout
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.client = self.client.with_timeout(timeout);
self
}
/// Connect to the server
pub fn connect(&self) -> Result<()> {
self.client.connect()
}
/// Close the connection
pub fn close(&self) -> Result<()> {
self.client.close()
}
/// Return a success message
pub fn return_success(&self, message: Option<&str>) -> Result<String> {
let mut script = "!!fake.return_success".to_string();
if let Some(msg) = message {
script.push_str(&format!(" message:'{}'", msg));
}
self.client.send_command(&script)
}
/// Return an error message
pub fn return_error(&self, message: Option<&str>) -> Result<String> {
let mut script = "!!fake.return_error".to_string();
if let Some(msg) = message {
script.push_str(&format!(" message:'{}'", msg));
}
// This will return a ClientError::ServerError with the error message
self.client.send_command(&script)
}
/// Return a JSON response
pub fn return_json(&self, message: Option<&str>, status: Option<&str>, code: Option<i32>) -> Result<FakeResponse> {
let mut script = "!!fake.return_json".to_string();
if let Some(msg) = message {
script.push_str(&format!(" message:'{}'", msg));
}
if let Some(status_val) = status {
script.push_str(&format!(" status:'{}'", status_val));
}
if let Some(code_val) = code {
script.push_str(&format!(" code:{}", code_val));
}
let response = self.client.send_command(&script)?;
// Parse the JSON response
match serde_json::from_str::<FakeResponse>(&response) {
Ok(result) => Ok(result),
Err(e) => Err(ClientError::JsonError(e)),
}
}
/// Return an invalid JSON response
pub fn return_invalid_json(&self) -> Result<FakeResponse> {
let script = "!!fake.return_invalid_json";
let response = self.client.send_command(&script)?;
// This should fail with a JSON parsing error
match serde_json::from_str::<FakeResponse>(&response) {
Ok(result) => Ok(result),
Err(e) => Err(ClientError::JsonError(e)),
}
}
/// Return an empty response
pub fn return_empty(&self) -> Result<String> {
let script = "!!fake.return_empty";
self.client.send_command(&script)
}
/// Return a large response
pub fn return_large(&self, size: Option<i32>) -> Result<String> {
let mut script = "!!fake.return_large".to_string();
if let Some(size_val) = size {
script.push_str(&format!(" size:{}", size_val));
}
self.client.send_command(&script)
}
/// Return a malformed error message
pub fn return_malformed_error(&self) -> Result<String> {
let script = "!!fake.return_malformed_error";
self.client.send_command(&script)
}
/// Get help information
pub fn help(&self) -> Result<String> {
let script = "!!fake.help";
self.client.send_command(&script)
}
}

View File

@@ -0,0 +1,242 @@
use std::io::{Read, Write};
use std::os::unix::net::UnixStream;
use std::time::Duration;
use thiserror::Error;
use std::fmt;
use std::error::Error as StdError;
mod processmanager;
mod fakehandler;
pub use processmanager::ProcessManagerClient;
pub use fakehandler::FakeHandlerClient;
/// Standard error response from the telnet server
#[derive(Debug, Clone)]
pub struct ServerError {
pub message: String,
pub raw_response: String,
}
impl fmt::Display for ServerError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.message)
}
}
impl StdError for ServerError {}
/// Error type for the client
#[derive(Error, Debug)]
pub enum ClientError {
#[error("IO error: {0}")]
IoError(#[from] std::io::Error),
#[error("JSON parsing error: {0}")]
JsonError(#[from] serde_json::Error),
#[error("Connection error: {0}")]
ConnectionError(String),
#[error("Command error: {0}")]
CommandError(String),
#[error("Server error: {0}")]
ServerError(String),
}
pub type Result<T> = std::result::Result<T, ClientError>;
/// A client for connecting to a Unix socket server with improved error handling
pub struct Client {
socket_path: String,
timeout: Duration,
secret: Option<String>,
}
impl Client {
/// Create a new Unix socket client
pub fn new(socket_path: &str) -> Self {
Self {
socket_path: socket_path.to_string(),
timeout: Duration::from_secs(10),
secret: None,
}
}
/// Set the connection timeout
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.timeout = timeout;
self
}
/// Set the authentication secret
pub fn with_secret(mut self, secret: &str) -> Self {
self.secret = Some(secret.to_string());
self
}
/// Connect to the Unix socket and return the stream
fn connect_socket(&self) -> Result<UnixStream> {
println!("DEBUG: Opening new connection to {}", self.socket_path);
// Connect to the socket
let stream = UnixStream::connect(&self.socket_path)
.map_err(|e| ClientError::ConnectionError(format!("Failed to connect to socket {}: {}", self.socket_path, e)))?;
// Set read timeout
stream.set_read_timeout(Some(self.timeout))?;
stream.set_write_timeout(Some(self.timeout))?;
// Read welcome message
let mut buffer = [0; 4096];
match stream.try_clone()?.read(&mut buffer) {
Ok(n) => {
let welcome = String::from_utf8_lossy(&buffer[0..n]);
if !welcome.contains("Welcome") {
return Err(ClientError::ConnectionError("Invalid welcome message".to_string()));
}
},
Err(e) => {
return Err(ClientError::IoError(e));
}
}
// Authenticate if a secret is provided
if let Some(secret) = &self.secret {
self.authenticate_stream(&stream, secret)?;
}
Ok(stream)
}
/// Authenticate with the server using the provided stream
fn authenticate_stream(&self, stream: &UnixStream, secret: &str) -> Result<()> {
let mut stream_clone = stream.try_clone()?;
let auth_command = format!("auth {}\n\n", secret);
// Send the auth command
stream_clone.write_all(auth_command.as_bytes())
.map_err(|e| ClientError::CommandError(format!("Failed to send auth command: {}", e)))?;
stream_clone.flush()
.map_err(|e| ClientError::CommandError(format!("Failed to flush auth command: {}", e)))?;
// Add a small delay to ensure the server has time to process the command
std::thread::sleep(Duration::from_millis(100));
// Read the response
let mut buffer = [0; 4096];
let n = stream_clone.read(&mut buffer)
.map_err(|e| ClientError::CommandError(format!("Failed to read auth response: {}", e)))?;
if n == 0 {
return Err(ClientError::ConnectionError("Connection closed by server during authentication".to_string()));
}
let response = String::from_utf8_lossy(&buffer[0..n]).to_string();
// Check for authentication success
if response.contains("Authentication successful") || response.contains("authenticated") {
Ok(())
} else {
Err(ClientError::ServerError(format!("Authentication failed: {}", response)))
}
}
/// Send a command to the server and get the response
pub fn send_command(&self, command: &str) -> Result<String> {
// Connect to the socket for this command
let mut stream = self.connect_socket()?;
// Ensure command ends with double newlines to execute it
let command = if command.ends_with("\n\n") {
command.to_string()
} else if command.ends_with('\n') {
format!("{}\n", command)
} else {
format!("{}\n\n", command)
};
// Send the command
stream.write_all(command.as_bytes())
.map_err(|e| ClientError::CommandError(format!("Failed to send command: {}", e)))?;
stream.flush()
.map_err(|e| ClientError::CommandError(format!("Failed to flush command: {}", e)))?;
// Add a small delay to ensure the server has time to process the command
std::thread::sleep(Duration::from_millis(100));
// Read the response
let mut buffer = [0; 8192]; // Use a larger buffer for large responses
let n = stream.read(&mut buffer)
.map_err(|e| ClientError::CommandError(format!("Failed to read response: {}", e)))?;
if n == 0 {
return Err(ClientError::ConnectionError("Connection closed by server".to_string()));
}
let response = String::from_utf8_lossy(&buffer[0..n]).to_string();
// Remove the prompt if present
let response = response.trim_end_matches("> ").trim().to_string();
// Check for standard error format
if response.starts_with("Error:") {
return Err(ClientError::ServerError(response));
}
// Close the connection by dropping the stream
println!("DEBUG: Closing connection to {}", self.socket_path);
drop(stream);
Ok(response)
}
/// Send a command and parse the JSON response
pub fn send_command_json<T: serde::de::DeserializeOwned>(&self, command: &str) -> Result<T> {
let response = self.send_command(command)?;
// If the response is empty, return an error
if response.trim().is_empty() {
return Err(ClientError::CommandError("Empty response".to_string()));
}
// Handle "action not supported" errors specially
if response.contains("action not supported") {
return Err(ClientError::ServerError(response));
}
// Try to parse the JSON response
match serde_json::from_str::<T>(&response) {
Ok(result) => Ok(result),
Err(e) => {
// If parsing fails, check if it's an error message
if response.starts_with("Error:") || response.contains("error") || response.contains("failed") {
Err(ClientError::ServerError(response))
} else {
Err(ClientError::JsonError(e))
}
},
}
}
/// For backward compatibility
pub fn connect(&self) -> Result<()> {
// Just verify we can connect
let stream = self.connect_socket()?;
drop(stream);
Ok(())
}
/// For backward compatibility
pub fn close(&self) -> Result<()> {
// No-op since we don't maintain a persistent connection
Ok(())
}
/// Authenticate with the server - kept for backward compatibility
pub fn authenticate(&self, secret: &str) -> Result<()> {
// Create a temporary connection to authenticate
let stream = self.connect_socket()?;
self.authenticate_stream(&stream, secret)
}
}

View File

@@ -0,0 +1,164 @@
use serde::{Deserialize, Serialize};
use std::time::Duration;
use crate::{Client, Result};
/// Information about a process
#[derive(Debug, Serialize, Deserialize, Default)]
pub struct ProcessInfo {
#[serde(default)]
pub name: String,
#[serde(default)]
pub command: String,
#[serde(default)]
pub status: String,
#[serde(default)]
pub pid: i32,
#[serde(default)]
pub start_time: String,
#[serde(default)]
pub uptime: String,
#[serde(default)]
pub cpu: String,
#[serde(default)]
pub memory: String,
#[serde(default)]
pub cron: Option<String>,
#[serde(default)]
pub job_id: Option<String>,
}
/// Client for the process manager
pub struct ProcessManagerClient {
client: Client,
}
impl ProcessManagerClient {
/// Create a new process manager client
pub fn new(socket_path: &str) -> Self {
Self {
client: Client::new(socket_path),
}
}
/// Set the connection timeout
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.client = self.client.with_timeout(timeout);
self
}
/// Set the authentication secret
pub fn with_secret(mut self, secret: &str) -> Self {
self.client = self.client.with_secret(secret);
self
}
/// Connect to the server
pub fn connect(&self) -> Result<()> {
self.client.connect()
}
/// Close the connection
pub fn close(&self) -> Result<()> {
self.client.close()
}
/// Start a new process
pub fn start(&self, name: &str, command: &str, log_enabled: bool, deadline: Option<i32>, cron: Option<&str>, job_id: Option<&str>) -> Result<String> {
let mut script = format!("!!process.start name:'{}' command:'{}' log:{}", name, command, log_enabled);
if let Some(deadline_val) = deadline {
script.push_str(&format!(" deadline:{}", deadline_val));
}
if let Some(cron_val) = cron {
script.push_str(&format!(" cron:'{}'", cron_val));
}
if let Some(job_id_val) = job_id {
script.push_str(&format!(" job_id:'{}'", job_id_val));
}
self.client.send_command(&script)
}
/// Stop a running process
pub fn stop(&self, name: &str) -> Result<String> {
let script = format!("!!process.stop name:'{}'", name);
self.client.send_command(&script)
}
/// Restart a process
pub fn restart(&self, name: &str) -> Result<String> {
let script = format!("!!process.restart name:'{}'", name);
self.client.send_command(&script)
}
/// Delete a process
pub fn delete(&self, name: &str) -> Result<String> {
let script = format!("!!process.delete name:'{}'", name);
self.client.send_command(&script)
}
/// List all processes
pub fn list(&self) -> Result<Vec<ProcessInfo>> {
let script = "!!process.list format:'json'";
let response = self.client.send_command(&script)?;
// Handle empty responses
if response.trim().is_empty() {
return Ok(Vec::new());
}
// Try to parse the response as JSON
match serde_json::from_str::<Vec<ProcessInfo>>(&response) {
Ok(processes) => Ok(processes),
Err(_) => {
// If parsing as a list fails, try parsing as a single ProcessInfo
match serde_json::from_str::<ProcessInfo>(&response) {
Ok(process) => Ok(vec![process]),
Err(_) => {
// If both parsing attempts fail, check if it's a "No processes found" message
if response.contains("No processes found") {
Ok(Vec::new())
} else {
// Otherwise, try to send it as JSON
self.client.send_command_json(&script)
}
}
}
}
}
}
/// Get the status of a specific process
pub fn status(&self, name: &str) -> Result<ProcessInfo> {
let script = format!("!!process.status name:'{}' format:'json'", name);
// Use the send_command_json method which handles JSON parsing with better error handling
self.client.send_command_json(&script)
}
/// Get the logs of a specific process
pub fn logs(&self, name: &str, lines: Option<i32>) -> Result<String> {
let mut script = format!("!!process.logs name:'{}'", name);
if let Some(lines_val) = lines {
script.push_str(&format!(" lines:{}", lines_val));
}
self.client.send_command(&script)
}
/// Set the logs path for the process manager
pub fn set_logs_path(&self, path: &str) -> Result<String> {
let script = format!("!!process.set_logs_path path:'{}'", path);
self.client.send_command(&script)
}
/// Get help information for the process manager
pub fn help(&self) -> Result<String> {
let script = "!!process.help";
self.client.send_command(&script)
}
}