diff --git a/jgrab-client/src/main.rs b/jgrab-client/src/main.rs index bb84676..69b56c4 100644 --- a/jgrab-client/src/main.rs +++ b/jgrab-client/src/main.rs @@ -1,20 +1,22 @@ -use dirs::home_dir; +extern crate dirs; +extern crate wait_timeout; + use std::env; -use std::fs::{create_dir_all, read_to_string, File}; -use std::io::{stdin, stdout, Cursor, Error, Read, Result, Stdin, Write}; +use std::fs::{create_dir_all, File, read_to_string}; +use std::io::{Cursor, Error, Read, Result, stdin, Stdin, stdout, Write}; use std::iter::Iterator; use std::net::{Shutdown, TcpStream}; use std::option::Option; use std::path::{Path, PathBuf}; -use std::process::{exit, Child, Command}; +use std::process::{Child, Command, exit}; use std::str; use std::thread::sleep; use std::time::Duration; + +use dirs::home_dir; use wait_timeout::ChildExt; -use Input::*; -extern crate dirs; -extern crate wait_timeout; +use Input::*; const MAX_RETRIES: usize = 5; @@ -72,6 +74,7 @@ fn main() { let jar_path = jar_path(&jgrab_home); let input: Input; + let mut ignore_read_error = false; if args.is_empty() { // no args, pipe stdin @@ -87,18 +90,14 @@ fn main() { return; } "--start" | "-t" => { - send_message_retrying( - TextIn(Cursor::new("-e null".to_string())), - &token_path, - &jar_path, - ); - return; + input = TextIn(Cursor::new("-e null".to_string())) } "--stop" | "-s" => { if connect().is_err() { log("daemon is not running"); return; } + ignore_read_error = true; input = TextIn(Cursor::new("--stop".to_string())) } "--version" | "-v" => { @@ -127,7 +126,7 @@ fn main() { } } - send_message_retrying(input, &token_path, &jar_path); + send_message_retrying(input, &token_path, &jar_path, ignore_read_error); } fn file_input(file_name: &String) -> Input { @@ -157,8 +156,12 @@ fn jar_path(home: &Path) -> PathBuf { path } -fn send_message_retrying(mut reader: R, token_path: &Path, jar_path: &Path) { - if send_message(&mut reader, false, token_path).is_some() { +fn send_message_retrying( + mut reader: R, + token_path: &Path, + jar_path: &Path, + ignore_read_error: bool) { + if send_message(&mut reader, false, token_path, ignore_read_error).is_some() { // failed to connect, try to start the daemon, then retry let mut retries = MAX_RETRIES; @@ -166,7 +169,7 @@ fn send_message_retrying(mut reader: R, token_path: &Path, jar_path: &P check_status(&mut child); while retries > 0 { - if let Some(err) = send_message(&mut reader, true, token_path) { + if let Some(err) = send_message(&mut reader, true, token_path, ignore_read_error) { check_status(&mut child); log(&format!("unable to connect to JGrab daemon: {}", err)); @@ -190,7 +193,11 @@ fn connect() -> Result { TcpStream::connect("127.0.0.1:5002") } -fn send_message(reader: &mut R, is_retry: bool, token_path: &Path) -> Option { +fn send_message( + reader: &mut R, + is_retry: bool, + token_path: &Path, + ignore_read_error: bool) -> Option { match connect() { Ok(mut stream) => { if is_retry { @@ -199,8 +206,8 @@ fn send_message(reader: &mut R, is_retry: bool, token_path: &Path) -> O match read_to_string(token_path) { Ok(token) => { - stream.write_all(token.as_bytes()).unwrap(); - let _ = stream.write(&[b'\n']).unwrap(); + stream.write_all(token.as_bytes()).expect("socket write error (token)"); + stream.write_all(&[b'\n']).expect("socket write error (newline)"); } Err(err) => return Some(err), }; @@ -213,33 +220,38 @@ fn send_message(reader: &mut R, is_retry: bool, token_path: &Path) -> O if n == 0 { break; } else { - stream.write_all(&socket_message[0..n]).unwrap(); + stream.write_all(&socket_message[0..n]) + .expect("socket write error (message)"); } } Err(err) => error(&err.to_string()), } } - stream.shutdown(Shutdown::Write).unwrap(); + stream.shutdown(Shutdown::Write).expect("shutdown socket write"); let mut client_buffer = socket_message; + let stdout = stdout(); + let mut lock = stdout.lock(); + loop { match stream.read(&mut client_buffer) { Ok(n) => { if n == 0 { break; } else { - stdout().write_all(&client_buffer[0..n]).unwrap(); + lock.write_all(&client_buffer[0..n]).expect("stdout write error"); } } + Err(_) if ignore_read_error => break, Err(err) => error(&err.to_string()), } } - Option::None + None } - Err(err) => Option::Some(err), + Err(err) => Some(err), } } @@ -313,9 +325,8 @@ fn show_daemon_version(token_path: &Path) { &mut TextIn(Cursor::new("--version".to_string())), false, token_path, - ) - .is_some() - { + false, + ).is_some() { println!("(Run the JGrab daemon to see its version)"); } } diff --git a/jgrab-client/test.sh b/jgrab-client/test.sh index 392353c..9fd1c0e 100755 --- a/jgrab-client/test.sh +++ b/jgrab-client/test.sh @@ -3,12 +3,19 @@ set -e DIR=$(dirname "$0") +F=$(basename "$0") export JGRAB_HOME="$DIR/.jgrab" mkdir "$JGRAB_HOME" cp "$DIR"/../jgrab-runner/build/libs/jgrab.jar "$JGRAB_HOME"/jgrab.jar ls -al "$JGRAB_HOME" jgrab="$DIR"/target/debug/jgrab-client +trap 'catch $? $LINENO' ERR +catch() { + echo "Error code: $1 on $F:$2" + rm -r "$JGRAB_HOME" || true +} + # start the daemon $jgrab -t @@ -24,6 +31,4 @@ fi # cleanup $jgrab -s - -# on Windows, this may fail rm -r "$JGRAB_HOME" || true diff --git a/jgrab-runner/src/main/java/com/athaydes/jgrab/daemon/CodeRunRequest.java b/jgrab-runner/src/main/java/com/athaydes/jgrab/daemon/CodeRunRequest.java new file mode 100644 index 0000000..8c6e4ab --- /dev/null +++ b/jgrab-runner/src/main/java/com/athaydes/jgrab/daemon/CodeRunRequest.java @@ -0,0 +1,21 @@ +package com.athaydes.jgrab.daemon; + +import com.athaydes.jgrab.code.JavaCode; + +interface Request { + +} + +enum StatelessRequest implements Request { + DO_NOTHING, DIE +} + +final class CodeRunRequest implements Request { + public final JavaCode code; + public final String[] args; + + public CodeRunRequest( JavaCode code, String[] args ) { + this.code = code; + this.args = args; + } +} diff --git a/jgrab-runner/src/main/java/com/athaydes/jgrab/daemon/JGrabDaemon.java b/jgrab-runner/src/main/java/com/athaydes/jgrab/daemon/JGrabDaemon.java index 943ef69..cc14d46 100644 --- a/jgrab-runner/src/main/java/com/athaydes/jgrab/daemon/JGrabDaemon.java +++ b/jgrab-runner/src/main/java/com/athaydes/jgrab/daemon/JGrabDaemon.java @@ -73,36 +73,58 @@ public static void start( RunArgs runArgs ) { System.exit( 1 ); return; } - - while ( true ) { + var running = true; + while ( running ) { try ( Socket clientSocket = serverSocket.accept(); final PrintStream out = new PrintStream( clientSocket.getOutputStream(), true ); BufferedReader in = new BufferedReader( new InputStreamReader( clientSocket.getInputStream() ) ) ) { + // the first line sent must be the current token if ( token.equals( in.readLine() ) ) { - handleClient( in, out, runArgs ); + running = handleClient( in, out, runArgs ); } else { logger.info( "Rejecting client as it did not present the current token" ); out.println( "=== JGrab authorization error ===" ); - clientSocket.close(); + running = false; } } catch ( IOException e ) { logger.warn( "Problem handling client message", e ); } } + logger.info( "Stopped JGrab daemon" ); + try { + // sleep a bit to allow the client to receive the response without errors + Thread.sleep( 100 ); + serverSocket.close(); + } catch ( IOException e ) { + logger.debug( "Exception closing server", e ); + } catch ( InterruptedException e ) { + logger.trace( "Interrupted while closing server" ); + } }, "jgrab-daemon" ).start(); } - private static void handleClient( + /** + * Handle the client synchronously. + * + * @param in client input + * @param out client output + * @param runArgs args + * @return true to continue, false to stop accepting connections + * @throws IOException on IO errors + */ + private static boolean handleClient( BufferedReader in, PrintStream out, RunArgs runArgs ) throws IOException { var request = parseRequest( in, out ); - if ( request == null ) return; + if ( request == StatelessRequest.DO_NOTHING ) return true; + if ( request == StatelessRequest.DIE ) return false; + var codeRequest = ( CodeRunRequest ) request; - logSourceCode( request.code ); + logSourceCode( codeRequest.code ); - var deps = request.code.extractDependencies(); + var deps = codeRequest.code.extractDependencies(); logger.debug( "Dependencies to grab: {}", deps ); @@ -113,16 +135,18 @@ private static void handleClient( // run this synchronously, which means only one program can run per daemon at a time try { - runArgs.accept( request.code, request.args, classpath ); + runArgs.accept( codeRequest.code, codeRequest.args, classpath ); } catch ( Throwable t ) { t.printStackTrace( out ); } finally { System.setOut( System.out ); System.setErr( System.err ); } + + return true; } - private static CodeRunRequest parseRequest( BufferedReader in, PrintStream out ) throws IOException { + private static Request parseRequest( BufferedReader in, PrintStream out ) throws IOException { String firstLine = null; String inputLine; StringBuilder messageBuilder = new StringBuilder( 1024 ); @@ -137,7 +161,7 @@ private static CodeRunRequest parseRequest( BufferedReader in, PrintStream out ) if ( firstLine == null ) { out.println( "Communication error (input line is null)" ); - return null; + return StatelessRequest.DO_NOTHING; } JavaCode code; @@ -156,8 +180,7 @@ private static CodeRunRequest parseRequest( BufferedReader in, PrintStream out ) if ( input.equals( STOP_OPTION ) ) { logger.info( "--stop option received, stopping JGrab Daemon" ); out.println( "=== JGrab Daemon stopped ===" ); - System.exit( 0 ); - return null; + return StatelessRequest.DIE; } if ( input.equals( VERSION_OPTION ) ) { @@ -165,14 +188,14 @@ private static CodeRunRequest parseRequest( BufferedReader in, PrintStream out ) System.setOut( out ); System.setErr( out ); JGrabRunner.printVersion(); - return null; + return StatelessRequest.DO_NOTHING; } if ( input.startsWith( JGrabOptions.SNIPPET_OPTION ) ) { String snippet = input.substring( JGrabOptions.SNIPPET_OPTION.length() ); if ( snippet.isEmpty() ) { out.println( "ERROR: no snippet provided to execute" ); - return null; + return StatelessRequest.DO_NOTHING; } else { code = new StringJavaCode( snippet ); args = new String[ 0 ]; @@ -185,16 +208,6 @@ private static CodeRunRequest parseRequest( BufferedReader in, PrintStream out ) return new CodeRunRequest( code, args ); } - private static final class CodeRunRequest { - public final JavaCode code; - public final String[] args; - - public CodeRunRequest( JavaCode code, String[] args ) { - this.code = code; - this.args = args; - } - } - private static void logSourceCode( Object source ) { logger.trace( "Source code:\n------------------------------------\n" + "{}\n------------------------------------\n", source );