157 lines
3.7 KiB
Rust
157 lines
3.7 KiB
Rust
use rand::{thread_rng, RngCore};
|
|
use std::env;
|
|
use std::fs::remove_file;
|
|
use std::io::{BufRead, BufReader, Write};
|
|
use std::os::unix::net::{UnixListener, UnixStream};
|
|
use std::path::PathBuf;
|
|
use std::process::Command;
|
|
use std::thread;
|
|
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
|
|
|
const ENV_VAR: &str = "IPC_SOCK_PATH";
|
|
|
|
struct SocketServer {
|
|
socket_path: PathBuf,
|
|
listener: UnixListener,
|
|
}
|
|
|
|
impl SocketServer {
|
|
fn new() -> Self {
|
|
let socket_path = create_socket();
|
|
|
|
SocketServer {
|
|
socket_path: socket_path.clone(),
|
|
listener: UnixListener::bind(&socket_path).unwrap(),
|
|
}
|
|
}
|
|
|
|
fn run(&self) {
|
|
println!("Listening on {}", self.socket_path.display());
|
|
|
|
for stream in self.listener.incoming() {
|
|
println!("New connection");
|
|
|
|
match stream {
|
|
Ok(stream) => {
|
|
thread::spawn(|| handle_client(stream));
|
|
}
|
|
Err(err) => {
|
|
println!("Connection failed: {:?}", err);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for SocketServer {
|
|
fn drop(&mut self) {
|
|
remove_file(&self.socket_path).unwrap();
|
|
}
|
|
}
|
|
|
|
fn create_socket() -> PathBuf {
|
|
let mut temp_dir = env::temp_dir();
|
|
let now = SystemTime::now();
|
|
|
|
let seed = thread_rng().next_u32();
|
|
|
|
let secs = now.duration_since(UNIX_EPOCH).unwrap().as_secs();
|
|
temp_dir.push(&format!("rust-ipc-{secs}-{seed}.sock"));
|
|
|
|
temp_dir
|
|
}
|
|
|
|
fn handle_client(mut stream: UnixStream) {
|
|
const TIMEOUT: Duration = Duration::from_secs(3);
|
|
|
|
let reader = BufReader::new(stream.try_clone().unwrap());
|
|
|
|
stream.set_read_timeout(Some(TIMEOUT)).unwrap();
|
|
|
|
let mut lines = reader.lines();
|
|
|
|
match lines.next() {
|
|
Some(Err(e)) if e.kind() == std::io::ErrorKind::WouldBlock => {
|
|
println!("Timeout");
|
|
return;
|
|
}
|
|
Some(Err(e)) => {
|
|
println!("Error reading {}", e);
|
|
return;
|
|
}
|
|
Some(Ok(msg)) => {
|
|
if msg != "handshake" {
|
|
println!("Got incorrect handshake: {:?}", msg);
|
|
return;
|
|
}
|
|
println!("Got correct handshake");
|
|
}
|
|
None => {
|
|
return;
|
|
}
|
|
}
|
|
|
|
// Restore blocking socket for increased efficiency
|
|
stream.set_read_timeout(None).unwrap();
|
|
|
|
for line in lines {
|
|
dbg!(line.unwrap());
|
|
|
|
// Panics if the socket dies mid-loop
|
|
writeln!(stream, "REPLY").unwrap();
|
|
}
|
|
|
|
println!("Closed - killing thread");
|
|
}
|
|
|
|
fn client(path: String) {
|
|
let mut stream = UnixStream::connect(path).unwrap();
|
|
|
|
let reader = BufReader::new(stream.try_clone().unwrap());
|
|
|
|
thread::sleep(Duration::from_secs(2));
|
|
|
|
writeln!(stream, "handshake").unwrap();
|
|
writeln!(stream, "handshake1").unwrap();
|
|
writeln!(stream, "handshake2").unwrap();
|
|
|
|
thread::sleep(Duration::from_secs(2));
|
|
|
|
writeln!(stream, "Hello world").unwrap();
|
|
|
|
let resp = reader.lines().next();
|
|
|
|
dbg!(resp);
|
|
|
|
thread::sleep(Duration::from_secs(2));
|
|
}
|
|
|
|
fn monitor_child(socket_path: PathBuf) {
|
|
loop {
|
|
println!("Starting child process");
|
|
let mut child_proc = Command::new(env::current_exe().unwrap())
|
|
.env(ENV_VAR, &socket_path)
|
|
.spawn()
|
|
.unwrap();
|
|
|
|
child_proc.wait().unwrap();
|
|
|
|
println!("Child process died");
|
|
|
|
thread::sleep(Duration::from_secs(5));
|
|
}
|
|
}
|
|
|
|
fn main() {
|
|
if let Ok(path) = env::var(ENV_VAR) {
|
|
client(path);
|
|
} else {
|
|
let server = SocketServer::new();
|
|
|
|
let socket_path = server.socket_path.clone();
|
|
thread::spawn(|| monitor_child(socket_path));
|
|
|
|
server.run();
|
|
}
|
|
}
|