Finally solve deadlock problem

This commit is contained in:
Michael Mikovsky
2025-06-08 17:16:08 -06:00
parent fda5e9ea02
commit e2764edfae
13 changed files with 310 additions and 212 deletions
-1
View File
@@ -2,7 +2,6 @@ use std::{
error::Error,
net::SocketAddr,
sync::{Arc, Mutex},
thread,
};
use serde::{Deserialize, Serialize};
+2
View File
@@ -2,3 +2,5 @@ mod packets;
pub use packets::C2Packet;
pub use packets::ErrorPacket;
pub use packets::Parameter;
pub use packets::Parameters;
+21 -1
View File
@@ -1,4 +1,7 @@
use std::fmt;
use std::{
collections::HashMap,
fmt::{self, Display},
};
use serde::{Deserialize, Serialize};
use serde_json::Result;
@@ -16,6 +19,15 @@ pub enum C2Packet {
SetCampaign(CampaignConfig),
AckSetCampaign,
GetParameter(String),
AckGetParameter(String, Option<Parameter>),
ParameterUpate(String, Parameter),
SetParameter(String, Parameter),
AckSetParameter(bool),
SetAllParameters(Parameters),
Error(ErrorPacket),
Sysinfo { hostname: String },
@@ -32,6 +44,14 @@ impl fmt::Debug for CampaignConfig {
}
}
pub type Parameters = HashMap<String, Parameter>;
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum Parameter {
Test1,
CurrentTab(i32),
}
impl C2Packet {
pub fn encode(&self) -> Result<String> {
serde_json::to_string(self)
+17
View File
@@ -10,6 +10,17 @@ pub trait Connection: Send + Sync {
fn write(&mut self, data: &str) -> Result<(), Self::Error>;
}
pub trait AsyncConnection<C>
where
C: Connection,
{
type Error: std::fmt::Debug;
fn as_async<T: Serialize + DeserializeOwned + Send + 'static>(
connection: C,
) -> (Sender<T>, Receiver<T>);
}
pub trait ServerTrait<C: Connection> {
type Error: std::fmt::Debug;
@@ -95,6 +106,12 @@ use std::net::SocketAddr;
use std::sync::Arc;
use std::thread;
use crossbeam_channel::Receiver;
use crossbeam_channel::Sender;
use serde::Serialize;
use serde::de::DeserializeOwned;
pub use tcp::TCPClient;
pub use tcp::TCPConnection;
pub use tcp::TCPServer;
use crate::connection;
+67 -1
View File
@@ -1,9 +1,14 @@
use std::{
error::Error,
io::{self, BufRead, BufReader, Write},
net::{SocketAddr, TcpListener, TcpStream},
thread,
};
use crate::networkers::{ClientTrait, Connection, ServerTrait};
use crossbeam_channel::{Receiver, Sender};
use serde::{Serialize, de::DeserializeOwned};
use crate::networkers::{AsyncConnection, ClientTrait, Connection, ServerTrait};
pub struct TCPConnection {
stream: TcpStream,
@@ -48,6 +53,67 @@ impl Connection for TCPConnection {
}
}
impl AsyncConnection<TCPConnection> for TCPConnection {
type Error = io::Error;
fn as_async<T: Serialize + DeserializeOwned + Send + 'static>(
connection: TCPConnection,
) -> (Sender<T>, Receiver<T>) {
let (send_tx, send_rx) = crossbeam_channel::unbounded::<T>();
let (recv_tx, recv_rx) = crossbeam_channel::unbounded::<T>();
let tx_clone = send_tx.clone();
thread::spawn(move || {
let mut reader = connection.reader;
let mut read = || -> Result<String, Self::Error> {
let mut line = String::new();
let n = reader.read_line(&mut line)?;
Ok(line.trim_end().to_string())
};
loop {
if let Ok(data) = read() {
if data.is_empty() {
break;
}
info!("Got {}", data);
if let Ok(decoded) = serde_json::from_str::<T>(&data) {
if let Err(e) = tx_clone.send(decoded) {
error!("Got error: {}", e);
}
}
}
}
});
let rx_clone = recv_rx.clone();
thread::spawn(move || {
let mut stream = connection.stream;
let mut write = |data: String| -> Result<(), Self::Error> {
writeln!(stream, "{}", data)?;
stream.flush()?;
Ok(())
};
loop {
if let Ok(data) = rx_clone.recv() {
if let Ok(encoded) = serde_json::to_string(&data) {
info!("Write {}", encoded);
if let Err(e) = write(encoded) {
error!("Got error: {}", e);
}
}
}
}
});
(recv_tx, send_rx)
}
}
pub struct TCPServer {
listener: TcpListener,
}