From 2dca737392a208e094b6f054f39100b8d31ed80d Mon Sep 17 00:00:00 2001 From: bnewbold Date: Wed, 12 Oct 2016 00:22:29 -0700 Subject: partially working shell command --- Cargo.lock | 37 ++++++++++++ Cargo.toml | 1 + src/bin/einhyrningsinsctl.rs | 136 +++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 77 ++++++++++++++++++++++-- 4 files changed, 247 insertions(+), 4 deletions(-) create mode 100644 src/bin/einhyrningsinsctl.rs diff --git a/Cargo.lock b/Cargo.lock index d784638..4ed9eae 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10,6 +10,7 @@ dependencies = [ "json 0.10.2 (registry+https://github.com/rust-lang/crates.io-index)", "log 0.3.6 (registry+https://github.com/rust-lang/crates.io-index)", "nix 0.7.0 (registry+https://github.com/rust-lang/crates.io-index)", + "rustyline 1.0.0 (registry+https://github.com/rust-lang/crates.io-index)", "slog 1.0.1 (registry+https://github.com/rust-lang/crates.io-index)", "slog-envlogger 0.5.0 (registry+https://github.com/rust-lang/crates.io-index)", "time 0.1.35 (registry+https://github.com/rust-lang/crates.io-index)", @@ -87,6 +88,11 @@ name = "crossbeam" version = "0.2.10" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "encode_unicode" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "env_logger" version = "0.3.5" @@ -168,6 +174,15 @@ dependencies = [ "libc 0.2.16 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "nix" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "bitflags 0.4.0 (registry+https://github.com/rust-lang/crates.io-index)", + "libc 0.2.16 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "nix" version = "0.7.0" @@ -246,6 +261,19 @@ dependencies = [ "semver 0.1.20 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "rustyline" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "encode_unicode 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", + "kernel32-sys 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)", + "libc 0.2.16 (registry+https://github.com/rust-lang/crates.io-index)", + "nix 0.5.1 (registry+https://github.com/rust-lang/crates.io-index)", + "unicode-width 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", + "winapi 0.2.8 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "semver" version = "0.1.20" @@ -348,6 +376,11 @@ name = "unicode-normalization" version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "unicode-width" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "url" version = "1.2.1" @@ -388,6 +421,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" "checksum chan-signal 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)" = "afbba6202dc1d10ff08c3b04e00e4d2d6cf5effee56cd9fee92928be6692379a" "checksum chrono 0.2.25 (registry+https://github.com/rust-lang/crates.io-index)" = "9213f7cd7c27e95c2b57c49f0e69b1ea65b27138da84a170133fd21b07659c00" "checksum crossbeam 0.2.10 (registry+https://github.com/rust-lang/crates.io-index)" = "0c5ea215664ca264da8a9d9c3be80d2eaf30923c259d03e870388eb927508f97" +"checksum encode_unicode 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)" = "28d65f1f5841ef7c6792861294b72beda34c664deb8be27970f36c306b7da1ce" "checksum env_logger 0.3.5 (registry+https://github.com/rust-lang/crates.io-index)" = "15abd780e45b3ea4f76b4e9a26ff4843258dd8a3eed2775a0e7368c2e7936c2f" "checksum getopts 0.2.14 (registry+https://github.com/rust-lang/crates.io-index)" = "d9047cfbd08a437050b363d35ef160452c5fe8ea5187ae0a624708c91581d685" "checksum idna 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "1053236e00ce4f668aeca4a769a09b3bf5a682d802abd6f3cb39374f6b162c11" @@ -400,6 +434,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" "checksum log 0.3.6 (registry+https://github.com/rust-lang/crates.io-index)" = "ab83497bf8bf4ed2a74259c1c802351fcd67a65baa86394b6ba73c36f4838054" "checksum matches 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)" = "bcc3ad8109fa4b522f9b0cd81440422781f564aaf8c195de6b9d6642177ad0dd" "checksum memchr 0.1.11 (registry+https://github.com/rust-lang/crates.io-index)" = "d8b629fb514376c675b98c1421e80b151d3817ac42d7c667717d282761418d20" +"checksum nix 0.5.1 (registry+https://github.com/rust-lang/crates.io-index)" = "bfb3ddedaa14746434a02041940495bf11325c22f6d36125d3bdd56090d50a79" "checksum nix 0.7.0 (registry+https://github.com/rust-lang/crates.io-index)" = "a0d95c5fa8b641c10ad0b8887454ebaafa3c92b5cd5350f8fc693adafd178e7b" "checksum num 0.1.36 (registry+https://github.com/rust-lang/crates.io-index)" = "bde7c03b09e7c6a301ee81f6ddf66d7a28ec305699e3d3b056d2fc56470e3120" "checksum num-integer 0.1.32 (registry+https://github.com/rust-lang/crates.io-index)" = "fb24d9bfb3f222010df27995441ded1e954f8f69cd35021f6bef02ca9552fb92" @@ -409,6 +444,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" "checksum regex 0.1.77 (registry+https://github.com/rust-lang/crates.io-index)" = "64b03446c466d35b42f2a8b203c8e03ed8b91c0f17b56e1f84f7210a257aa665" "checksum regex-syntax 0.3.7 (registry+https://github.com/rust-lang/crates.io-index)" = "48f0573bcee95a48da786f8823465b5f2a1fae288a55407aca991e5b3e0eae11" "checksum rustc_version 0.1.7 (registry+https://github.com/rust-lang/crates.io-index)" = "c5f5376ea5e30ce23c03eb77cbe4962b988deead10910c372b226388b594c084" +"checksum rustyline 1.0.0 (registry+https://github.com/rust-lang/crates.io-index)" = "00b06ac9c8e8e3e83b33d175d39a9f7b6c2c930c82990593719c8e48788ae2d9" "checksum semver 0.1.20 (registry+https://github.com/rust-lang/crates.io-index)" = "d4f410fedcf71af0345d7607d246e7ad15faaadd49d240ee3b24e5dc21a820ac" "checksum slog 1.0.1 (registry+https://github.com/rust-lang/crates.io-index)" = "d816659be2527f0f85437f31ebd3bea98a2553f83c41a6404abae9f530c9ab62" "checksum slog-envlogger 0.5.0 (registry+https://github.com/rust-lang/crates.io-index)" = "dfea715bb310c33c8f90e659bce5b95e39851348b9a7e2a77495a069662def78" @@ -421,6 +457,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" "checksum timer 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)" = "a9522a9ec40055e2f9e514e38d2415a496e81dbfc1ece15d98d2fe55c44946b3" "checksum unicode-bidi 0.2.3 (registry+https://github.com/rust-lang/crates.io-index)" = "c1f7ceb96afdfeedee42bade65a0d585a6a0106f681b6749c8ff4daa8df30b3f" "checksum unicode-normalization 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)" = "26643a2f83bac55f1976fb716c10234485f9202dcd65cfbdf9da49867b271172" +"checksum unicode-width 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)" = "2d6722facc10989f63ee0e20a83cd4e1714a9ae11529403ac7e0afd069abc39e" "checksum url 1.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "8527c62d9869a08325c38272b3f85668df22a65890c61a639d233dc0ed0b23a2" "checksum utf8-ranges 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)" = "a1ca13c08c41c9c3e04224ed9ff80461d97e121589ff27c753a16cb10830ae0f" "checksum void 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)" = "6a02e4885ed3bc0f2de90ea6dd45ebcbb66dacffe03547fadbb0eeae2770887d" diff --git a/Cargo.toml b/Cargo.toml index 2d69c85..795ff18 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ slog = "1.0" slog-envlogger = "0.5" json = "*" url = "1.2" +rustyline = "1.0" [replace] "chan-signal:0.1.6" = { git = 'https://github.com/bnewbold/chan-signal' } diff --git a/src/bin/einhyrningsinsctl.rs b/src/bin/einhyrningsinsctl.rs new file mode 100644 index 0000000..c057b9e --- /dev/null +++ b/src/bin/einhyrningsinsctl.rs @@ -0,0 +1,136 @@ +/* + * einhyrningsinsctl: controller/shell for einhyrningsins + * Copyright (C) 2016 Bryan Newbold + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU General Public License as published by + * the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU General Public License for more details. + * + * You should have received a copy of the GNU General Public License + * along with this program. If not, see . + */ + +#[macro_use] +extern crate json; + +extern crate getopts; +extern crate log; +extern crate env_logger; +extern crate nix; +extern crate timer; +extern crate time; +extern crate chan_signal; +extern crate url; +extern crate rustyline; + + +use std::io::prelude::*; +use std::io::{BufReader, BufWriter}; +use std::env; +use std::path::Path; +use std::process::exit; +use std::os::unix::net::UnixStream; +use getopts::Options; + +use rustyline::error::ReadlineError; +use rustyline::Editor; + + +// This is the main event loop +fn shell(ctrl_stream: UnixStream) { + + let mut reader = BufReader::new(&ctrl_stream); + let mut writer = BufWriter::new(&ctrl_stream); + + // `()` can be used when no completer is required + let mut rl = Editor::<()>::new(); + + loop { + let readline = rl.readline("> "); + match readline { + Ok(line) => { + rl.add_history_entry(&line); + if line.len() == 0 { continue; }; + let mut chunks = line.split(' '); + let cmd = chunks.nth(0).unwrap(); + let args = chunks.collect(); + send_msg(&mut reader, &mut writer, cmd, args).unwrap(); + }, + Err(ReadlineError::Interrupted) | Err(ReadlineError::Eof) => { + println!("Quitting..."); + break + }, + Err(err) => { + println!("Shell Error: {:?}", err); + break + } + } + } +} + +fn send_msg(reader: &mut BufRead, writer: &mut Write, cmd: &str, args: Vec<&str>) -> Result { + let mut buffer = String::new(); + let mut arg_list = json::JsonValue::new_array(); + for a in args { + arg_list.push(a).unwrap(); + } + let req = object!{ + "command" => cmd, + "args" => arg_list + }; + //println!("Sending: {}", req.dump()); + writer.write_all(req.dump().as_bytes()).unwrap(); + writer.write_all("\n".as_bytes()).unwrap(); + writer.flush().unwrap(); + + reader.read_line(&mut buffer).unwrap(); + //println!("Got: {}", buffer); + let reply = json::parse(&buffer).unwrap(); + println!("{}", reply.as_str().unwrap()); + Ok(reply.as_str().unwrap().to_string()) +} + +fn print_usage(opts: Options) { + let brief = "usage:\teinhyrningsinsctl [options] program"; + println!(""); + print!("{}", opts.usage(&brief)); +} + +fn main() { + + let args: Vec = env::args().collect(); + + let mut opts = Options::new(); + opts.optflag("h", "help", "print this help menu"); + + let matches = match opts.parse(&args[1..]) { + Ok(m) => { m } + Err(f) => { println!("{}", f.to_string()); print_usage(opts); exit(-1); } + }; + + if matches.opt_present("h") { + print_usage(opts); + return; + } + + // Bind to Control Socket + let ctrl_path = Path::new("/tmp/einhorn.sock"); + // XXX: handle this more gracefully (per-process) + if !ctrl_path.exists() { + println!("Couldn't find control socket: {:?}", ctrl_path); + exit(-1); + } + println!("Connecting to control socket: {:?}", ctrl_path); + let ctrl_stream = UnixStream::connect(ctrl_path).unwrap(); + + send_msg(&mut BufReader::new(&ctrl_stream), &mut BufWriter::new(&ctrl_stream), "ehlo", vec![]).unwrap(); + + shell(ctrl_stream); + exit(0); +} diff --git a/src/main.rs b/src/main.rs index 0141a5f..725cdbd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -18,6 +18,7 @@ #[macro_use] extern crate chan; +extern crate json; extern crate getopts; extern crate log; @@ -47,7 +48,6 @@ use time::Duration; use std::collections::HashMap; use getopts::Options; -use url::percent_encoding; use chan_signal::Signal; use chan::{Sender, Receiver}; use std::os::unix::io::{RawFd, IntoRawFd}; @@ -167,6 +167,7 @@ enum TimerAction { enum CtrlAction { Increment, Decrement, + ManualAck(u32), SigAll(Signal), ShutdownAll, UpgradeAll, @@ -278,6 +279,14 @@ fn shepard(mut cfg: EinConfig, signal_rx: Receiver) { CtrlAction::Status => { req.tx.send(format!("UNIMPLEMENTED")); }, + CtrlAction::ManualAck(pid) => { + if let Some(o) = brood.get_mut(&pid) { + if o.is_active() { + o.state = OffspringState::Healthy; + } + } + req.tx.send(format!("Acknowledged!")); + }, } }, signal_rx.recv() -> sig => match sig.expect("Error with signal handler") { @@ -508,14 +517,74 @@ fn ctrl_socket_handle(stream: UnixStream, ctrl_req_tx: Sender) { let reader = BufReader::new(&stream); let mut writer = BufWriter::new(&stream); for rawline in reader.lines() { + let rawline = rawline.unwrap(); - let line = percent_encoding::percent_decode(rawline.as_bytes()).decode_utf8().unwrap(); - println!("Decoded message: {}", line); + println!("Got line: {}", rawline); + if rawline.len() == 0 { + continue; + } + + // Parse message + let req_action = if let Ok(msg) = json::parse(&rawline) { + match msg["command"].as_str() { + Some("worker:ack") => { + CtrlAction::ManualAck(msg["pid"].as_u32().unwrap()) + }, + Some("signal") => { + CtrlAction::SigAll(match msg["args"][0].as_str() { + Some("SIGHUP") | Some("HUP") => Signal::HUP, + Some("SIGINT") | Some("INT") => Signal::INT, + Some("SIGTERM") | Some("TERM") => Signal::TERM, + Some("SIGKILL") | Some("KILL") => Signal::KILL, + Some("SIGUSR1") | Some("USR1") => Signal::KILL, + Some("SIGUSR2") | Some("USR2") => Signal::USR2, + Some("SIGSTOP") | Some("STOP") => Signal::STOP, + Some("SIGCONT") | Some("CONT") => Signal::CONT, + Some(_) | None => { + writer.write_all("\"Missing or unhandled 'signal'\"\n".as_bytes()).unwrap(); + writer.flush().unwrap(); + continue; + }, + }) + }, + Some("inc") => CtrlAction::Increment, + Some("dec") => CtrlAction::Decrement, + Some("status") => CtrlAction::Status, + Some("die") => CtrlAction::ShutdownAll, + Some("upgrade") => CtrlAction::UpgradeAll, + Some("ehlo") => { + writer.write_all("\"Hi there!\"\n\r".as_bytes()).unwrap(); + writer.flush().unwrap(); + continue; + }, + Some("help") => { + writer.write_all("\"Command Listing: \"\n".as_bytes()).unwrap(); // TODO + writer.flush().unwrap(); + continue; + }, + Some(_) | None => { + writer.write_all("\"Missing or unhandled 'command'\"\n".as_bytes()).unwrap(); + writer.flush().unwrap(); + continue; + }, + } + } else { + writer.write_all("\"Expected valid JSON!\"\n".as_bytes()).unwrap(); + writer.flush().unwrap(); + continue; + }; + + // Send request let (tx, rx): (Sender, Receiver) = chan::async(); - let req = CtrlRequest{ action: CtrlAction::Status, tx: tx }; + let req = CtrlRequest{ action: req_action, tx: tx }; ctrl_req_tx.send(req); + + // Send reply let resp = rx.recv().unwrap(); + writer.write_all("\"".as_bytes()).unwrap(); writer.write_all(resp.as_bytes()).unwrap(); + writer.write_all("\"\n".as_bytes()).unwrap(); + writer.flush().unwrap(); } stream.shutdown(std::net::Shutdown::Both).unwrap(); } -- cgit v1.2.3