352 lines
11 KiB
Rust
352 lines
11 KiB
Rust
//! Client control socket implementation.
|
|
use std::io::prelude::*;
|
|
use std::io::BufReader;
|
|
use std::io::LineWriter;
|
|
use std::os::unix::net::UnixListener;
|
|
use std::os::unix::net::UnixStream;
|
|
use std::path::PathBuf;
|
|
use std::{io, net};
|
|
|
|
use radicle::node::Handle;
|
|
|
|
use crate::identity::Id;
|
|
use crate::node;
|
|
use crate::node::FetchLookup;
|
|
use crate::runtime;
|
|
|
|
#[derive(thiserror::Error, Debug)]
|
|
pub enum Error {
|
|
#[error("failed to bind control socket listener: {0}")]
|
|
Bind(io::Error),
|
|
#[error("invalid socket path specified: {0}")]
|
|
InvalidPath(PathBuf),
|
|
}
|
|
|
|
/// Listen for commands on the control socket, and process them.
|
|
pub fn listen<H: Handle<Error = runtime::HandleError>>(
|
|
listener: UnixListener,
|
|
mut handle: H,
|
|
) -> Result<(), Error> {
|
|
log::debug!(target: "control", "Control thread listening on socket..");
|
|
|
|
for incoming in listener.incoming() {
|
|
match incoming {
|
|
Ok(mut stream) => {
|
|
log::debug!(target: "control", "Accepted new client on control socket..");
|
|
|
|
if let Err(e) = drain(&stream, &mut handle) {
|
|
if let DrainError::Shutdown = e {
|
|
log::debug!(target: "control", "Shutdown requested..");
|
|
// Channel might already be disconnected if shutdown
|
|
// came from somewhere else. Ignore errors.
|
|
handle.shutdown().ok();
|
|
break;
|
|
}
|
|
writeln!(stream, "error: {}", e).ok();
|
|
|
|
stream.flush().ok();
|
|
stream.shutdown(net::Shutdown::Both).ok();
|
|
}
|
|
}
|
|
Err(e) => log::error!(target: "control", "Failed to accept incoming connection: {}", e),
|
|
}
|
|
}
|
|
log::debug!(target: "control", "Exiting control loop..");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[derive(thiserror::Error, Debug)]
|
|
enum DrainError {
|
|
#[error("invalid command argument `{0}`, {1}")]
|
|
InvalidCommandArg(String, Box<dyn std::error::Error>),
|
|
#[error("unknown command `{0}`")]
|
|
UnknownCommand(String),
|
|
#[error("runtime error: {0}")]
|
|
Runtime(#[from] runtime::HandleError),
|
|
#[error("i/o error: {0}")]
|
|
Io(#[from] io::Error),
|
|
#[error("shutdown requested")]
|
|
Shutdown,
|
|
}
|
|
|
|
fn drain<H: Handle<Error = runtime::HandleError>>(
|
|
stream: &UnixStream,
|
|
handle: &mut H,
|
|
) -> Result<(), DrainError> {
|
|
let mut reader = BufReader::new(stream);
|
|
let mut writer = LineWriter::new(stream);
|
|
let mut line = String::new();
|
|
|
|
reader.read_line(&mut line)?;
|
|
|
|
let cmd = line.trim_end();
|
|
|
|
log::debug!(target: "control", "Received `{cmd}` on control socket");
|
|
|
|
// TODO: refactor to include helper
|
|
match cmd.split_once(' ') {
|
|
Some(("fetch", arg)) => match arg.parse() {
|
|
Ok(id) => {
|
|
fetch(id, LineWriter::new(stream), handle)?;
|
|
}
|
|
Err(err) => {
|
|
return Err(DrainError::InvalidCommandArg(arg.to_owned(), Box::new(err)));
|
|
}
|
|
},
|
|
Some(("track-repo", arg)) => match arg.parse() {
|
|
Ok(id) => match handle.track_repo(id) {
|
|
Ok(updated) => {
|
|
if updated {
|
|
writeln!(writer, "{}", node::RESPONSE_OK)?;
|
|
} else {
|
|
writeln!(writer, "{}", node::RESPONSE_NOOP)?;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
return Err(DrainError::Runtime(e));
|
|
}
|
|
},
|
|
Err(err) => {
|
|
return Err(DrainError::InvalidCommandArg(arg.to_owned(), Box::new(err)));
|
|
}
|
|
},
|
|
Some(("untrack-repo", arg)) => match arg.parse() {
|
|
Ok(id) => match handle.untrack_repo(id) {
|
|
Ok(updated) => {
|
|
if updated {
|
|
writeln!(writer, "{}", node::RESPONSE_OK)?;
|
|
} else {
|
|
writeln!(writer, "{}", node::RESPONSE_NOOP)?;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
return Err(DrainError::Runtime(e));
|
|
}
|
|
},
|
|
Err(err) => {
|
|
return Err(DrainError::InvalidCommandArg(arg.to_owned(), Box::new(err)));
|
|
}
|
|
},
|
|
Some(("track-node", args)) => {
|
|
let (peer, alias) = if let Some((peer, alias)) = args.split_once(' ') {
|
|
(peer, Some(alias.to_owned()))
|
|
} else {
|
|
(args, None)
|
|
};
|
|
match peer.parse() {
|
|
Ok(id) => match handle.track_node(id, alias) {
|
|
Ok(updated) => {
|
|
if updated {
|
|
writeln!(writer, "{}", node::RESPONSE_OK)?;
|
|
} else {
|
|
writeln!(writer, "{}", node::RESPONSE_NOOP)?;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
return Err(DrainError::Runtime(e));
|
|
}
|
|
},
|
|
Err(err) => {
|
|
return Err(DrainError::InvalidCommandArg(
|
|
args.to_owned(),
|
|
Box::new(err),
|
|
));
|
|
}
|
|
}
|
|
}
|
|
Some(("untrack-node", arg)) => match arg.parse() {
|
|
Ok(id) => match handle.untrack_node(id) {
|
|
Ok(updated) => {
|
|
if updated {
|
|
writeln!(writer, "{}", node::RESPONSE_OK)?;
|
|
} else {
|
|
writeln!(writer, "{}", node::RESPONSE_NOOP)?;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
return Err(DrainError::Runtime(e));
|
|
}
|
|
},
|
|
Err(err) => {
|
|
return Err(DrainError::InvalidCommandArg(arg.to_owned(), Box::new(err)));
|
|
}
|
|
},
|
|
Some(("announce-refs", arg)) => match arg.parse() {
|
|
Ok(id) => {
|
|
if let Err(e) = handle.announce_refs(id) {
|
|
return Err(DrainError::Runtime(e));
|
|
}
|
|
writeln!(writer, "{}", node::RESPONSE_OK)?;
|
|
}
|
|
Err(err) => {
|
|
return Err(DrainError::InvalidCommandArg(arg.to_owned(), Box::new(err)));
|
|
}
|
|
},
|
|
Some((cmd, _)) => return Err(DrainError::UnknownCommand(cmd.to_owned())),
|
|
|
|
// Commands with no arguments.
|
|
None => match cmd {
|
|
"status" => {
|
|
println!("RECEIVED 'status'");
|
|
writeln!(writer, "{}", node::RESPONSE_OK).ok();
|
|
}
|
|
"routing" => match handle.routing() {
|
|
Ok(c) => {
|
|
for (id, seed) in c.iter() {
|
|
writeln!(writer, "{id} {seed}",)?;
|
|
}
|
|
}
|
|
Err(e) => return Err(DrainError::Runtime(e)),
|
|
},
|
|
"inventory" => match handle.inventory() {
|
|
Ok(c) => {
|
|
for id in c.iter() {
|
|
writeln!(writer, "{id}")?;
|
|
}
|
|
}
|
|
Err(e) => return Err(DrainError::Runtime(e)),
|
|
},
|
|
"shutdown" => {
|
|
return Err(DrainError::Shutdown);
|
|
}
|
|
_ => {
|
|
return Err(DrainError::UnknownCommand(line));
|
|
}
|
|
},
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn fetch<W: Write, H: Handle<Error = runtime::HandleError>>(
|
|
id: Id,
|
|
mut writer: W,
|
|
handle: &mut H,
|
|
) -> Result<(), DrainError> {
|
|
match handle.fetch(id) {
|
|
Err(e) => {
|
|
return Err(DrainError::Runtime(e));
|
|
}
|
|
Ok(FetchLookup::Found { seeds, results }) => {
|
|
let seeds = Vec::from(seeds);
|
|
|
|
writeln!(
|
|
writer,
|
|
"ok: found {} seeds for {} ({:?})", // TODO: Better output
|
|
seeds.len(),
|
|
&id,
|
|
&seeds,
|
|
)?;
|
|
|
|
for result in results
|
|
.iter()
|
|
.take(results.capacity().unwrap_or(seeds.len()))
|
|
{
|
|
match result.result {
|
|
Ok(updated) => {
|
|
writeln!(writer, "ok: {} fetched from {}", &id, result.remote)?;
|
|
|
|
for update in updated {
|
|
writeln!(writer, "{}", update)?;
|
|
}
|
|
}
|
|
Err(err) => {
|
|
writeln!(
|
|
writer,
|
|
"error: {} failed to fetch from {}: {}",
|
|
&id, result.remote, err
|
|
)?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Ok(FetchLookup::NotFound) => {
|
|
writeln!(writer, "error: {} was not found", &id)?;
|
|
}
|
|
Ok(FetchLookup::NotTracking) => {
|
|
writeln!(writer, "error: {} is not tracked", &id)?;
|
|
}
|
|
Ok(FetchLookup::Error(err)) => {
|
|
writeln!(writer, "error: {}", err)?;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::io::prelude::*;
|
|
use std::os::unix::net::UnixStream;
|
|
use std::thread;
|
|
|
|
use super::*;
|
|
use crate::identity::Id;
|
|
use crate::node::Handle;
|
|
use crate::node::{Node, NodeId};
|
|
use crate::test;
|
|
|
|
#[test]
|
|
fn test_control_socket() {
|
|
let tmp = tempfile::tempdir().unwrap();
|
|
let handle = test::handle::Handle::default();
|
|
let socket = tmp.path().join("alice.sock");
|
|
let projs = test::arbitrary::set::<Id>(1..3);
|
|
let listener = UnixListener::bind(&socket).unwrap();
|
|
|
|
thread::spawn({
|
|
let handle = handle.clone();
|
|
|
|
move || listen(listener, handle)
|
|
});
|
|
|
|
for proj in &projs {
|
|
let mut buf = [0; 2];
|
|
let mut stream = loop {
|
|
if let Ok(stream) = UnixStream::connect(&socket) {
|
|
break stream;
|
|
}
|
|
};
|
|
writeln!(&stream, "announce-refs {proj}").unwrap();
|
|
stream.read_exact(&mut buf).unwrap();
|
|
assert_eq!(&buf, &[b'o', b'k']);
|
|
}
|
|
|
|
for proj in &projs {
|
|
assert!(handle.updates.lock().unwrap().contains(proj));
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_track_untrack() {
|
|
let tmp = tempfile::tempdir().unwrap();
|
|
let socket = tmp.path().join("node.sock");
|
|
let proj = test::arbitrary::gen::<Id>(1);
|
|
let peer = test::arbitrary::gen::<NodeId>(1);
|
|
let listener = UnixListener::bind(&socket).unwrap();
|
|
let mut handle = Node::new(&socket);
|
|
|
|
thread::spawn({
|
|
let handle = crate::test::handle::Handle::default();
|
|
|
|
move || crate::control::listen(listener, handle)
|
|
});
|
|
|
|
// Wait for node to be online.
|
|
while !handle.is_running() {}
|
|
|
|
assert!(handle.track_repo(proj).unwrap());
|
|
assert!(!handle.track_repo(proj).unwrap());
|
|
assert!(handle.untrack_repo(proj).unwrap());
|
|
assert!(!handle.untrack_repo(proj).unwrap());
|
|
|
|
assert!(handle
|
|
.track_node(peer, Some(String::from("alice")))
|
|
.unwrap());
|
|
assert!(!handle
|
|
.track_node(peer, Some(String::from("alice")))
|
|
.unwrap());
|
|
assert!(handle.untrack_node(peer).unwrap());
|
|
assert!(!handle.untrack_node(peer).unwrap());
|
|
}
|
|
}
|