diff --git a/src/com/flaremicro/crossjeeves/ScriptProcessor.java b/src/com/flaremicro/crossjeeves/ScriptProcessor.java index 6723d39..1c20a31 100644 --- a/src/com/flaremicro/crossjeeves/ScriptProcessor.java +++ b/src/com/flaremicro/crossjeeves/ScriptProcessor.java @@ -1,10 +1,14 @@ package com.flaremicro.crossjeeves; +import java.io.BufferedInputStream; import java.io.File; +import java.io.FileInputStream; import java.io.FileWriter; import java.io.IOException; import java.io.StringReader; +import java.util.ArrayList; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.Map.Entry; import java.util.Properties; @@ -25,7 +29,10 @@ import com.flaremicro.crossjeeves.net.NetworkHandler; import com.flaremicro.crossjeeves.net.StreamForwarder; import com.flaremicro.crossjeeves.net.packet.Packet; import com.flaremicro.crossjeeves.net.packet.Packet3Clone; +import com.flaremicro.crossjeeves.net.packet.Packet4FileData; +import com.flaremicro.crossjeeves.net.packet.Packet5Artifact; import com.flaremicro.crossjeeves.net.packet.Packet7LogEntry; +import com.flaremicro.util.FileUtils; import com.flaremicro.util.Util; import com.flaremicro.util.ZipUtils; @@ -39,11 +46,11 @@ public class ScriptProcessor { private File workspace = null; private File workingDirectory = null; private boolean terminated = false; - + public ScriptProcessor(NetworkHandler netHandler, File workspace, Properties properties) { this.netHandler = netHandler; this.properties = properties; - this.workspace = workspace; + this.workspace = workspace; } public String requireAttribute(Element e, String attribute) throws ScriptProcessingException { @@ -70,7 +77,7 @@ public class ScriptProcessor { public void processScript(String script) throws ScriptProcessingException { workingDirectory = workspace; - if(terminated) + if (terminated) throw new ScriptProcessingException("Processor has been terminated"); DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance(); try @@ -89,14 +96,14 @@ public class ScriptProcessor { } catch (Exception ex) { - throw new ScriptProcessingException("Script threw an exception during processing" ,ex); + throw new ScriptProcessingException("Script threw an exception during processing", ex); } } private void processScriptNodes(NodeList nodeList) throws ScriptProcessingException { for (int i = 0; i < nodeList.getLength(); i++) { - if(terminated) + if (terminated) throw new ScriptProcessingException("Processor has been terminated"); Node node = nodeList.item(i); Element e = getElement(node); @@ -189,7 +196,7 @@ public class ScriptProcessor { File file = netHandler.waitForFile(id); if (file == null) throw new ScriptProcessingException("File failed to transfer"); - if(!ZipUtils.unzipDirectory(file, workspace)) + if (!ZipUtils.unzipDirectory(file, workspace)) throw new ScriptProcessingException("File failed to decompress"); } @@ -229,7 +236,6 @@ public class ScriptProcessor { processBuilder.directory(workspace); Map environment = processBuilder.environment(); - for (Entry set : currentEnvironment.entrySet()) { environment.put(set.getKey(), set.getValue()); @@ -243,15 +249,15 @@ public class ScriptProcessor { StreamForwarder stdOutForwarder = new StreamForwarder(netHandler, runningProcess.getInputStream(), Packet7LogEntry.STD_OUT); StreamForwarder stdErrForwarder = new StreamForwarder(netHandler, runningProcess.getErrorStream(), Packet7LogEntry.STD_ERR); - + stdOutForwarder.startAsync(); stdErrForwarder.startAsync(); - + int exit = runningProcess.waitFor(); - + stdOutForwarder.end(); stdErrForwarder.end(); - + return exit; } catch (IOException ex) @@ -296,9 +302,94 @@ public class ScriptProcessor { throw new ScriptProcessingException("Process returned exit code " + returnCode); } - private void artifactProcessor(Element e) { - // TODO Auto-generated method stub - + private void artifactProcessor(Element e) throws ScriptProcessingException { + NodeList nodeList = e.getChildNodes(); + for (int i = 0; i < nodeList.getLength(); i++) + { + Element var = getElement(nodeList.item(i)); + if (var == null) + continue; + String name = var.getTagName().trim(); + if (name.equalsIgnoreCase("workspace")) + { + System.out.println("Collecting workspace artifacts..."); + List filePaths = new ArrayList(); + NodeList workspaceNodes = var.getChildNodes(); + for (int j = 0; j < workspaceNodes.getLength(); j++) + { + workspaceNodes.item(j); + Element fileElement = getElement(workspaceNodes.item(j)); + if (fileElement == null) + continue; + String fileCommand = fileElement.getTagName().trim(); + if (fileCommand.equalsIgnoreCase("file")) + { + filePaths.add(fileElement.getTextContent()); + } + else if (fileCommand.equalsIgnoreCase("files")) + { + String resolver = this.getAttribute(fileElement, "resolver", "wildcard").trim(); + try + { + if (resolver.equalsIgnoreCase("wildcard")) + { + FileUtils.resolvePathsWildcard(workspace, fileElement.getTextContent(), filePaths); + } + else if (resolver.equalsIgnoreCase("regex")) + { + FileUtils.resolvePathsRegex(workspace, fileElement.getTextContent(), filePaths); + } + else + { + throw new ScriptProcessingException("Invalid file resolver: " + resolver); + } + } + catch (Exception ex) + { + throw new ScriptProcessingException("Failed to process files directive " + fileElement.getTextContent() + " with resolver " + resolver, ex); + } + } + } + if (filePaths.size() <= 0) + { + System.out.println("No workspace artifacts!"); + continue; + } + BufferedInputStream fileStream = null; + try + { + System.out.println("Sending artifacts..."); + long fileId = random.nextLong(); + File zipFile = File.createTempFile("workspace-" + fileId, ".zip"); + if (!ZipUtils.zipList(workspace, filePaths, zipFile)) + throw new ScriptProcessingException("Failed to compress workspace artifacts"); + Packet packet = new Packet5Artifact(fileId, zipFile.length(), "."); + netHandler.enqueue(packet); + fileStream = new BufferedInputStream(new FileInputStream(zipFile)); + int read; + byte[] buffer = new byte[4096]; + while ((read = fileStream.read(buffer)) > -1) + { + if (read == 0) + continue; + Packet4FileData dataPacket = new Packet4FileData(fileId, (short) read, buffer); + netHandler.enqueue(dataPacket); + } + packet = new Packet4FileData(fileId, (short) 0, new byte[] {}); + netHandler.enqueue(packet); + netHandler.waitForPacketSend(packet); + System.out.println("Sent!"); + } + catch (IOException ex) + { + throw new ScriptProcessingException("Failed to upload workspace artifacts", ex); + } + finally + { + Util.cleanClose(fileStream); + } + } + } } public void processAsync(final String script) { @@ -328,9 +419,9 @@ public class ScriptProcessor { public void terminate() { terminated = true; - if(runningProcess != null) + if (runningProcess != null) runningProcess.destroy(); - if(this.executionThread != null) + if (this.executionThread != null) { this.executionThread.interrupt(); } diff --git a/src/com/flaremicro/crossjeeves/net/ClientHandler.java b/src/com/flaremicro/crossjeeves/net/ClientHandler.java index 3b9fa83..2506024 100644 --- a/src/com/flaremicro/crossjeeves/net/ClientHandler.java +++ b/src/com/flaremicro/crossjeeves/net/ClientHandler.java @@ -8,6 +8,8 @@ import java.io.FileInputStream; import java.io.IOException; import java.net.InetAddress; import java.net.Socket; +import java.util.ArrayList; +import java.util.List; import com.flaremicro.crossjeeves.net.packet.Packet; import com.flaremicro.crossjeeves.net.packet.Packet0Identify; @@ -24,24 +26,38 @@ import com.flaremicro.util.ZipUtils; public class ClientHandler extends NetworkHandler { private String script; - + + private List runningThreads = new ArrayList(); + public ClientHandler(Socket socket, String script) throws IOException { super(socket); this.script = script; } - + private void init() { System.out.println("CrossJeeves connected! Sending Identify packet..."); enqueue(new Packet0Identify(0)); this.beginWriteThread(); this.beginReading(); } - + @Override public void handlePacket(Packet packet) { disconnect(INVALID_PACKET_RECIEVED.id, "Recieved invalid packet " + packet.getId() + " (Unknown)"); } + @Override + public void doDisconnect(int err) { + super.doDisconnect(err); + synchronized (runningThreads) + { + for (Thread t : runningThreads) + { + t.interrupt(); + } + } + } + @Override public void handlePacket(Packet0Identify packet) { if (packet.getProtocolVersion() != Packet.PROTOCOL_VERSION) @@ -57,7 +73,7 @@ public class ClientHandler extends NetworkHandler { @Override public void handlePacket(Packet1Status packet) { - if((packet.getFlags() & Packet1Status.BUSY) != 0) + if ((packet.getFlags() & Packet1Status.BUSY) != 0) { disconnect(OUTDATED_AGENT.id, "Agent is too busy"); } @@ -120,8 +136,50 @@ public class ClientHandler extends NetworkHandler { } @Override - public void handlePacket(Packet5Artifact packet) { + public void handlePacket(final Packet5Artifact packet) { + try + { + System.out.println("Getting artifacts!"); + final String workspace = System.getenv("WORKSPACE"); + if (workspace == null || workspace.trim().length() <= 0) + { + disconnect(INVALID_SYSTEM_STATE.id, "Workspace is not defined"); + } + this.beginFile(packet.getFileId(), packet.getFileSize()); + + Thread thread = new Thread(new Runnable() { + public void run() { + File file = waitForFile(packet.getFileId()); + System.out.println("Got artifacts!"); + if (file == null) + { + disconnect(FILE_DOWNLOAD_FAILURE.id, "Failed to download transferred file"); + } + if (!ZipUtils.unzipDirectory(file, new File(workspace))) + { + disconnect(FILE_DOWNLOAD_FAILURE.id, "Failed to extract transferred file"); + } + else + { + System.out.println("Extracted artifacts!"); + } + synchronized (runningThreads) + { + runningThreads.remove(this); + } + } + }); + synchronized (runningThreads) + { + runningThreads.add(thread); + } + thread.start(); + } + catch (IOException e) + { + disconnect(FILE_DOWNLOAD_FAILURE.id, "Failed to create file for transfer"); + } } @Override @@ -152,10 +210,9 @@ public class ClientHandler extends NetworkHandler { @Override public void handlePacket(Packet7LogEntry packet) { - if(packet.getStdOutput() == Packet7LogEntry.STD_ERR) - System.err.println(packet.getLogEntry()); - else - System.out.println(packet.getLogEntry()); + if (packet.getStdOutput() == Packet7LogEntry.STD_ERR) + System.err.println("[AGENT] " + packet.getLogEntry()); + else System.out.println("[AGENT] " + packet.getLogEntry()); } } diff --git a/src/com/flaremicro/crossjeeves/net/NetworkHandler.java b/src/com/flaremicro/crossjeeves/net/NetworkHandler.java index dd243ac..fcfc06b 100644 --- a/src/com/flaremicro/crossjeeves/net/NetworkHandler.java +++ b/src/com/flaremicro/crossjeeves/net/NetworkHandler.java @@ -46,6 +46,22 @@ public abstract class NetworkHandler { this.out = new DataOutputStream(socket.getOutputStream()); } + public void waitForPacketSend(Packet packet) { + synchronized (packet) + { + if (!outbox.contains(packet)) + return; + try + { + packet.wait(); + } + catch (InterruptedException e) + { + e.printStackTrace(); + } + } + } + public File waitForFile(long fileId) { try { @@ -103,6 +119,11 @@ public abstract class NetworkHandler { Util.cleanClose(in); Util.cleanClose(out); Util.cleanClose(socket); + for (FileTransferInfo transferInfo : downloadQueue.values()) + { + Util.cleanClose(transferInfo.outputStream); + transferInfo.file.delete(); + } } public int getExitCode() { @@ -161,6 +182,10 @@ public abstract class NetworkHandler { try { packet.sendPacket(out); + synchronized (packet) + { + packet.notifyAll(); + } } catch (IOException e) { @@ -200,7 +225,7 @@ public abstract class NetworkHandler { public abstract void handlePacket(Packet5Artifact packet); public abstract void handlePacket(Packet6Disconnect packet); - + public abstract void handlePacket(Packet7LogEntry packet); public abstract void handlePacket(Packet127KeepAlive packet); diff --git a/src/com/flaremicro/crossjeeves/net/ServerHandler.java b/src/com/flaremicro/crossjeeves/net/ServerHandler.java index be159b8..9a825ee 100644 --- a/src/com/flaremicro/crossjeeves/net/ServerHandler.java +++ b/src/com/flaremicro/crossjeeves/net/ServerHandler.java @@ -20,6 +20,7 @@ import com.flaremicro.crossjeeves.net.packet.Packet5Artifact; import com.flaremicro.crossjeeves.net.packet.Packet4FileData; import com.flaremicro.crossjeeves.net.packet.Packet6Disconnect; import com.flaremicro.crossjeeves.net.packet.Packet7LogEntry; +import com.flaremicro.util.FileUtils; public class ServerHandler extends NetworkHandler { private CrossJeevesHost host; @@ -52,6 +53,7 @@ public class ServerHandler extends NetworkHandler { scriptProcessor.terminate(); host.removeConnection(this); super.doDisconnect(exitCode); + FileUtils.deleteDirectory(workspace); } @Override diff --git a/src/com/flaremicro/crossjeeves/net/StreamForwarder.java b/src/com/flaremicro/crossjeeves/net/StreamForwarder.java index f6da059..1a5828f 100644 --- a/src/com/flaremicro/crossjeeves/net/StreamForwarder.java +++ b/src/com/flaremicro/crossjeeves/net/StreamForwarder.java @@ -44,7 +44,7 @@ public class StreamForwarder { } public void start() { - if (!isRunning && parent != null) + if (!isRunning && parent == null) { isRunning = true; parent = new Thread(this); diff --git a/src/com/flaremicro/crossjeeves/net/packet/Packet5Artifact.java b/src/com/flaremicro/crossjeeves/net/packet/Packet5Artifact.java index 5250e49..844471a 100644 --- a/src/com/flaremicro/crossjeeves/net/packet/Packet5Artifact.java +++ b/src/com/flaremicro/crossjeeves/net/packet/Packet5Artifact.java @@ -8,24 +8,28 @@ import com.flaremicro.crossjeeves.net.NetworkHandler; public class Packet5Artifact extends Packet{ private long fileId; + private long fileSize; private String relativeFile; public Packet5Artifact(){ } - public Packet5Artifact(long fileId, String relativeFile){ + public Packet5Artifact(long fileId, long fileSize, String relativeFile){ this.fileId = fileId; + this.fileSize = fileSize; this.relativeFile = relativeFile; } public void recievePacket(DataInputStream in) throws IOException { fileId = in.readLong(); + fileSize = in.readLong(); relativeFile = in.readUTF(); } public void sendPacket(DataOutputStream out) throws IOException { super.sendPacket(out); out.writeLong(fileId); + out.writeLong(fileSize); out.writeUTF(relativeFile); } @@ -33,6 +37,12 @@ public class Packet5Artifact extends Packet{ networkHandler.handlePacket(this); } + + public long getFileSize() + { + return fileSize; + } + public long getFileId() { return fileId; diff --git a/src/com/flaremicro/util/FileUtils.java b/src/com/flaremicro/util/FileUtils.java new file mode 100644 index 0000000..473d7f4 --- /dev/null +++ b/src/com/flaremicro/util/FileUtils.java @@ -0,0 +1,55 @@ +package com.flaremicro.util; + +import java.io.File; +import java.io.IOException; +import java.util.List; +import java.util.regex.Pattern; + +public class FileUtils { + + public static void resolvePathsWildcard(File baseDir, String wildcard, List matchedPaths) + { + throw new RuntimeException("Resolving with wildcards is not yet supported"); + } + + public static void resolvePathsRegex(File baseDir, String regex, List matchedPaths) throws IOException { + Pattern pattern = Pattern.compile(regex); + resolvePathsRecursiveRegex(baseDir, baseDir, pattern, matchedPaths); + } + + private static void resolvePathsRecursiveRegex(File baseDir, File currDir, Pattern pattern, List matchedPaths) throws IOException { + String baseDirPath = baseDir.getCanonicalPath(); + for (File file : currDir.listFiles()) + { + String filePath = file.getCanonicalPath(); + //Likely a link, ignore + if (!filePath.startsWith(baseDirPath)) + continue; + + filePath = filePath.substring(baseDirPath.length()); + + if (file.isDirectory()) + { + resolvePathsRecursiveRegex(baseDir, file, pattern, matchedPaths); + } + else + { + if (pattern.matcher(filePath).matches()) + { + matchedPaths.add(filePath); + } + } + } + } + + public static boolean deleteDirectory(File directory) { + File[] allContents = directory.listFiles(); + if (allContents != null) { + for (File file : allContents) { + deleteDirectory(file); + } + } + directory.deleteOnExit(); + return directory.delete(); + } +} diff --git a/src/com/flaremicro/util/ZipUtils.java b/src/com/flaremicro/util/ZipUtils.java index 74e7a55..57794c4 100644 --- a/src/com/flaremicro/util/ZipUtils.java +++ b/src/com/flaremicro/util/ZipUtils.java @@ -6,6 +6,7 @@ import java.io.FileInputStream; import java.io.FileOutputStream; import java.io.IOException; import java.util.Enumeration; +import java.util.List; import java.util.zip.ZipEntry; import java.util.zip.ZipException; import java.util.zip.ZipFile; @@ -13,17 +14,68 @@ import java.util.zip.ZipOutputStream; public class ZipUtils { - public static void zipDirectory(File sourceDir, File zipFile) throws IOException { - FileOutputStream fos = new FileOutputStream(zipFile); - ZipOutputStream zos = new ZipOutputStream(fos); + public static boolean zipList(File baseDir, List filePaths, File zipFile) { + FileOutputStream fos = null; + ZipOutputStream zos = null; + FileInputStream fis = null; try { - zipFilesRecursively(sourceDir, sourceDir, zos); + fos = new FileOutputStream(zipFile); + zos = new ZipOutputStream(fos); + + for (String filePath : filePaths) + { + File file = new File(baseDir, filePath); + if (file.isDirectory() || !file.exists()) + continue; + fis = new FileInputStream(file); + String zipEntryName = filePath.replace("\\", "/"); + ZipEntry zipEntry = new ZipEntry(zipEntryName); + zos.putNextEntry(zipEntry); + byte[] buffer = new byte[4096]; + int bytesRead; + while ((bytesRead = fis.read(buffer)) != -1) + { + zos.write(buffer, 0, bytesRead); + } + zos.closeEntry(); + fis.close(); + fis = null; + } + + return true; + } + catch (IOException ex) + { + ex.printStackTrace(); } finally { - zos.close(); + Util.cleanClose(fis); + Util.cleanClose(zos); } + return false; + } + + public static boolean zipDirectory(File sourceDir, File zipFile) { + FileOutputStream fos = null; + ZipOutputStream zos = null; + try + { + fos = new FileOutputStream(zipFile); + zos = new ZipOutputStream(fos); + zipFilesRecursively(sourceDir, sourceDir, zos); + return true; + } + catch (IOException ex) + { + ex.printStackTrace(); + } + finally + { + Util.cleanClose(zos); + } + return false; } private static void zipFilesRecursively(File rootDir, File source, ZipOutputStream zos) throws IOException { diff --git a/test/com/flaremicro/crossjeeves/net/packet/test/Packet5ArtifactTest.java b/test/com/flaremicro/crossjeeves/net/packet/test/Packet5ArtifactTest.java index f203252..4485eb0 100644 --- a/test/com/flaremicro/crossjeeves/net/packet/test/Packet5ArtifactTest.java +++ b/test/com/flaremicro/crossjeeves/net/packet/test/Packet5ArtifactTest.java @@ -22,13 +22,15 @@ public class Packet5ArtifactTest extends PacketTestBase { @Test public void testWrite() throws IOException { - Packet5Artifact packet = new Packet5Artifact(1, "file.txt"); + Packet5Artifact packet = new Packet5Artifact(1, 2, "file.txt"); assertEquals(packet.getFileId(), 1); + assertEquals(packet.getFileSize(), 2); assertEquals("file.txt", packet.getRelativeFile()); packet.sendPacket(output()); assertEquals(5, input().readByte()); assertEquals(1, input().readLong()); + assertEquals(2, input().readLong()); assertEquals("file.txt", input().readUTF()); } @@ -38,18 +40,20 @@ public class Packet5ArtifactTest extends PacketTestBase { Packet5Artifact packet = new Packet5Artifact(); output().writeLong(1); + output().writeLong(2); output().writeUTF("file.txt"); packet.recievePacket(input()); assertEquals(1, packet.getFileId()); + assertEquals(2, packet.getFileSize()); assertEquals("file.txt", packet.getRelativeFile()); } @Test public void testProcess() throws IOException { - Packet5Artifact packet = new Packet5Artifact(1, "file.txt"); + Packet5Artifact packet = new Packet5Artifact(1, 2, "file.txt"); packet.processPacket(handler); Mockito.verify(handler, Mockito.times(1)).handlePacket(packet); }