diff --git a/pom.xml b/pom.xml index cf3d0e0..82d9e04 100644 --- a/pom.xml +++ b/pom.xml @@ -114,6 +114,16 @@ 4.12 test + + net.java.dev.jna + jna + 4.4.0 + + + net.java.dev.jna + jna-platform + 4.4.0 + UTF-8 diff --git a/src/main/java/com/profesorfalken/jpowershell/PowerShell.java b/src/main/java/com/profesorfalken/jpowershell/PowerShell.java index f3580c3..a9ef09f 100644 --- a/src/main/java/com/profesorfalken/jpowershell/PowerShell.java +++ b/src/main/java/com/profesorfalken/jpowershell/PowerShell.java @@ -25,7 +25,9 @@ import java.io.IOException; import java.io.OutputStreamWriter; import java.io.PrintWriter; +import java.lang.reflect.Field; import java.nio.charset.Charset; +import java.time.Instant; import java.util.Date; import java.util.Map; import java.util.concurrent.Callable; @@ -34,9 +36,14 @@ import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.logging.Level; import java.util.logging.Logger; +import com.sun.jna.Pointer; +import com.sun.jna.platform.win32.Kernel32; +import com.sun.jna.platform.win32.WinNT; + /** * Allows to open a session into PowerShell console and launch different * commands.
@@ -58,6 +65,7 @@ public class PowerShell { //Threaded session variables private boolean closed = false; private ExecutorService threadpool; + private boolean blocked = false; //Config values private int maxThreads = 3; @@ -150,10 +158,14 @@ public static PowerShell openSession() throws PowerShellNotAvailableException { * @return PowerShellResponse the information returned by powerShell */ public PowerShellResponse executeCommand(String command) { + + Instant instant = Instant.now(); + long timeStampMillis = instant.toEpochMilli(); + Callable commandProcessor = new PowerShellCommandProcessor("standard", - p.getInputStream(), this.maxWait, this.waitPause, this.scriptMode); + p.getInputStream(), this.maxWait, this.waitPause, this.scriptMode, timeStampMillis); Callable commandProcessorError = new PowerShellCommandProcessor("error", - p.getErrorStream(), this.maxWait, this.waitPause, false); + p.getErrorStream(), this.maxWait, this.waitPause, false, timeStampMillis); String commandOutput = ""; boolean isError = false; @@ -169,25 +181,47 @@ public PowerShellResponse executeCommand(String command) { //Launch command commandWriter.println(command); + long elapsedTime = 0; + try { - while (!result.isDone() && !resultError.isDone()) { + while (!result.isDone() && !resultError.isDone() && (elapsedTime < maxWait)) { Thread.sleep(this.waitPause); + elapsedTime = Instant.now().toEpochMilli() - timeStampMillis; } - if (result.isDone()) { - if (((PowerShellCommandProcessor) commandProcessor).isTimeout()) { - timeout = true; - } else { - commandOutput = result.get(); - } - } else { + + if (elapsedTime >= maxWait) { + //Command timed out but may have some output + timeout = true; + commandOutput = result.get(waitPause, TimeUnit.MILLISECONDS); + } + + if (result.isDone()) { + //Command finished successfully + + do { + //stderr might need another wait cycle to finish, as it finishes after the stdout. + //If it is ready then, we can wait for the output to be read and the future to be done. + Thread.sleep(this.waitPause); + } while (((PowerShellCommandProcessor) commandProcessorError).isReady() && !resultError.isDone() && !timeout); + commandOutput = result.get(); + } + + if (resultError.isDone() && !timeout) { + //Command failed isError = true; commandOutput = resultError.get(); - } + } + } catch (InterruptedException ex) { + isError = true; Logger.getLogger(PowerShell.class.getName()).log(Level.SEVERE, "Unexpected error when processing PowerShell command", ex); } catch (ExecutionException ex) { + isError = true; Logger.getLogger(PowerShell.class.getName()).log(Level.SEVERE, "Unexpected error when processing PowerShell command", ex); - } finally { + } catch (TimeoutException e) { + Logger.getLogger(PowerShell.class.getName()).log(Level.WARNING, "Failed to read data due to timeout, close will be forced!"); + this.blocked = true; + } finally { //issue #2. Close and cancel processors/threads - Thanks to r4lly for helping me here ((PowerShellCommandProcessor) commandProcessor).close(); ((PowerShellCommandProcessor) commandProcessorError).close(); @@ -343,6 +377,39 @@ public PowerShellResponse executeScript(BufferedReader srcReader, String params) * Closes all the resources used to maintain the PowerShell context */ public void close() { + + if (blocked) { + Logger.getLogger(PowerShell.class.getName()).log(Level.INFO, "Shell is blocked, closing of PowerShell will be forced!"); + + int pid = -1; + + Field f; + try { + f = p.getClass().getDeclaredField("handle"); + f.setAccessible(true); + long handLong = f.getLong(p); + + Kernel32 kernel = Kernel32.INSTANCE; + WinNT.HANDLE handle = new WinNT.HANDLE(); + handle.setPointer(Pointer.createConstant(handLong)); + pid = kernel.GetProcessId(handle); + + } catch (Exception e1) { + Logger.getLogger(PowerShell.class.getName()).log(Level.SEVERE, "Unexpected error while getting process handle", e1); + } + + Logger.getLogger(PowerShell.class.getName()).log(Level.INFO, "Killing process with PID: " + pid); + + try { + Runtime.getRuntime().exec("taskkill.exe /f /PID " + pid + " /T" ); + this.closed = true; + } catch (IOException e) { + Logger.getLogger(PowerShell.class.getName()).log(Level.SEVERE, "Unexpected error while killing powershell process", e); + } + + + } + if (!this.closed) { try { Future closeTask = threadpool.submit(new Callable() { @@ -363,6 +430,7 @@ public String call() throws Exception { Logger.getLogger(PowerShell.class.getName()).log(Level.SEVERE, "Unexpected error when when closing streams", ex); } commandWriter.close(); + if (this.threadpool != null) { try { this.threadpool.shutdownNow(); @@ -379,6 +447,7 @@ public String call() throws Exception { private void waitUntilClose(Future task) throws InterruptedException { int closingTime = 0; + while (!task.isDone()) { if (closingTime > maxWait) { Logger.getLogger(PowerShell.class.getName()).log(Level.SEVERE, "Unexpected error when closing PowerShell: TIMEOUT!"); diff --git a/src/main/java/com/profesorfalken/jpowershell/PowerShellCommandProcessor.java b/src/main/java/com/profesorfalken/jpowershell/PowerShellCommandProcessor.java index 38a24d6..1f46497 100644 --- a/src/main/java/com/profesorfalken/jpowershell/PowerShellCommandProcessor.java +++ b/src/main/java/com/profesorfalken/jpowershell/PowerShellCommandProcessor.java @@ -1,5 +1,5 @@ /* - * Copyright 2016 Javier Garcia Alonso. + * Copyright 2016-2017 Javier Garcia Alonso. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -19,6 +19,7 @@ import java.io.IOException; import java.io.InputStream; import java.io.InputStreamReader; +import java.time.Instant; import java.util.concurrent.Callable; import java.util.logging.Level; import java.util.logging.Logger; @@ -45,19 +46,27 @@ class PowerShellCommandProcessor implements Callable { private final long maxWait; private final int waitPause; + private long commandStart; + private boolean ready = false; + /** * Constructor that takes the output and the input of the PowerShell session * * @param commandWriter the input to the PowerShell console * @param inputStream the stream needed to read the command output */ - public PowerShellCommandProcessor(String name, InputStream inputStream, long maxWait, int waitPause, boolean scriptMode) { + public PowerShellCommandProcessor(String name, InputStream inputStream, long maxWait, int waitPause, boolean scriptMode, long commandStart) { this.reader = new BufferedReader(new InputStreamReader( inputStream)); this.name = name; this.maxWait = maxWait; this.waitPause = waitPause; this.scriptMode = scriptMode; + this.commandStart = commandStart; + } + + public boolean isReady() { + return this.ready; } /** @@ -112,17 +121,23 @@ private void readData(StringBuilder powerShellOutput) throws IOException { //Checks when we can start reading the output. Timeout if it takes too long in order to avoid hangs private boolean startReading() throws IOException, InterruptedException { - int timeWaiting = 0; + long commandStart = this.commandStart; + long elapsedTime = Instant.now().toEpochMilli() - commandStart; while (!this.reader.ready()) { Thread.sleep(this.waitPause); - timeWaiting += this.waitPause; - if ((timeWaiting > this.maxWait) || this.closed) { - this.timeout = timeWaiting > this.maxWait; + elapsedTime = Instant.now().toEpochMilli() - commandStart; + + if ((elapsedTime > maxWait)) { + this.timeout = true; + return false; + } else if (this.closed) { return false; - } + } } - return true; + this.ready = true; + return true; + } //Checks when we the reader can continue to read.