Add more net code

src/com/flaremicro/crossjeeves/ScriptProcessor.java

@@ -30,9 +30,9 @@ public class ScriptProcessor {
 	private Process runningProcess = null;
 	private NetworkHandler netHandler;
 	private Random random = new Random();
-	
+	private Thread executionThread = null;
-	public ScriptProcessor(NetworkHandler netHandler, Properties properties)
+
-	{
+	public ScriptProcessor(NetworkHandler netHandler, Properties properties) {
 		this.netHandler = netHandler;
 		this.properties = properties;
 	}
@@ -44,7 +44,7 @@ public class ScriptProcessor {
 		}
 		return e.getAttribute(attribute);
 	}
-	
+
 	public String getAttribute(Element e, String attribute, String fallback) throws ScriptProcessingException {
 		if (!e.hasAttribute(attribute))
 		{
@@ -156,7 +156,7 @@ public class ScriptProcessor {
 		Packet packet = new Packet3Clone(id, 0);
 		netHandler.enqueue(packet);
 		File file = netHandler.waitForFile(id);
-		if(file == null)
+		if (file == null)
 			throw new ScriptProcessingException("File failed to transfer");
 	}
 
@@ -176,7 +176,7 @@ public class ScriptProcessor {
 		}
 		finally
 		{
-			if(writer != null)
+			if (writer != null)
 			{
 				try
 				{
@@ -190,20 +190,19 @@ public class ScriptProcessor {
 		}
 		return tempFile;
 	}
-	
+
-	public int beginProcess(String... args) throws ScriptProcessingException
+	public int beginProcess(String... args) throws ScriptProcessingException {
-	{
 		ProcessBuilder processBuilder = new ProcessBuilder(args);
 		Map<String, String> environment = processBuilder.environment();
-		
+
-		for(Entry<String, String> set : currentEnvironment.entrySet())
+		for (Entry<String, String> set : currentEnvironment.entrySet())
 		{
 			environment.put(set.getKey(), set.getValue());
 		}
-		
+
-        if(runningProcess != null)
+		if (runningProcess != null)
-        	throw new ScriptProcessingException("Previous process not terminated");
+			throw new ScriptProcessingException("Previous process not terminated");
-        try
+		try
 		{
 			runningProcess = processBuilder.start();
 			return runningProcess.exitValue();
@@ -218,30 +217,29 @@ public class ScriptProcessor {
 	private void batchProcessor(Element e) throws ScriptProcessingException {
 		File file = writeScriptToTempFile(e.getTextContent(), ".bat");
 		int returnCode = beginProcess(properties.getProperty("win-cmd", "cmd.exe"), "/c", file.getAbsolutePath());
-		if(returnCode != 0)
+		if (returnCode != 0)
 			throw new ScriptProcessingException("Process returned exit code " + returnCode);
 	}
-        
 
 	private void pythonProcessor(Element e) throws ScriptProcessingException {
 		File file = writeScriptToTempFile(e.getTextContent(), ".py");
 		int returnCode = beginProcess(properties.getProperty("python", "python3"), file.getAbsolutePath());
-		if(returnCode != 0)
+		if (returnCode != 0)
 			throw new ScriptProcessingException("Process returned exit code " + returnCode);
 	}
 
 	private void shellProcessor(Element e) throws ScriptProcessingException {
 		File file = writeScriptToTempFile(e.getTextContent(), ".sh");
 		String shellName = getAttribute(e, "which", "bash");
-		int returnCode = beginProcess(properties.getProperty(shellName+"-executor", shellName), file.getAbsolutePath());
+		int returnCode = beginProcess(properties.getProperty(shellName + "-executor", shellName), file.getAbsolutePath());
-		if(returnCode != 0)
+		if (returnCode != 0)
 			throw new ScriptProcessingException("Process returned exit code " + returnCode);
 	}
 
 	private void powershellProcessor(Element e) throws DOMException, ScriptProcessingException {
 		File file = writeScriptToTempFile(e.getTextContent(), ".ps1");
 		int returnCode = beginProcess(properties.getProperty("win-ps", "powershell.exe"), "-ExecutionPolicy", "Bypass", "-File", file.getAbsolutePath());
-		if(returnCode != 0)
+		if (returnCode != 0)
 			throw new ScriptProcessingException("Process returned exit code " + returnCode);
 	}
 
@@ -249,4 +247,28 @@ public class ScriptProcessor {
 		// TODO Auto-generated method stub
 
 	}
+
+	public void processAsync(final String script) {
+		if (executionThread == null)
+		{
+			executionThread = new Thread(new Runnable() {
+				public void run() {
+					try
+					{
+						processScript(script);
+					}
+					catch (ScriptProcessingException e)
+					{
+						//TODO Disconnect
+						e.printStackTrace();
+					}
+					finally
+					{
+						executionThread = null;
+					}
+				}
+			});
+			executionThread.start();
+		}
+	}
 }

src/com/flaremicro/crossjeeves/net/ClientHandler.java

@@ -1,5 +1,7 @@
 package com.flaremicro.crossjeeves.net;
 
+import static com.flaremicro.crossjeeves.net.ErrorCodes.FILE_DOWNLOAD_FAILURE;
+
 import java.io.BufferedInputStream;
 import java.io.File;
 import java.io.FileInputStream;
@@ -12,8 +14,9 @@ import com.flaremicro.crossjeeves.net.packet.Packet127KeepAlive;
 import com.flaremicro.crossjeeves.net.packet.Packet1Status;
 import com.flaremicro.crossjeeves.net.packet.Packet2Script;
 import com.flaremicro.crossjeeves.net.packet.Packet3Clone;
-import com.flaremicro.crossjeeves.net.packet.Packet5Artifact;
 import com.flaremicro.crossjeeves.net.packet.Packet4FileData;
+import com.flaremicro.crossjeeves.net.packet.Packet5Artifact;
+import com.flaremicro.crossjeeves.net.packet.Packet6Disconnect;
 import com.flaremicro.util.Util;
 import com.flaremicro.util.ZipUtils;
 
@@ -26,53 +29,53 @@ public class ClientHandler extends NetworkHandler {
 	@Override
 	public void handlePacket(Packet packet) {
 		// TODO Auto-generated method stub
-		
+
 	}
 
 	@Override
 	public void handlePacket(Packet0Identify packet) {
 		// TODO Auto-generated method stub
-		
+
 	}
 
 	@Override
 	public void handlePacket(Packet1Status packet) {
 		// TODO Auto-generated method stub
-		
+
 	}
 
 	@Override
 	public void handlePacket(Packet2Script packet) {
 		// TODO Auto-generated method stub
-		
+
 	}
 
 	@Override
 	public void handlePacket(Packet3Clone packet) {
 		long fileId = packet.getFileID();
 		String workspace = System.getenv("WORKSPACE");
-		if(workspace == null || workspace.trim().length() <= 0)
+		if (workspace == null || workspace.trim().length() <= 0)
 		{
 			//disconnect
 		}
 		BufferedInputStream fileStream = null;
 		try
 		{
-			File workspaceZip = File.createTempFile("workspace-"+fileId, ".zip");
+			File workspaceZip = File.createTempFile("workspace-" + fileId, ".zip");
 			ZipUtils.zipDirectory(new File(workspace), workspaceZip);
 			packet = new Packet3Clone(fileId, workspaceZip.length());
 			enqueue(packet);
 			fileStream = new BufferedInputStream(new FileInputStream(workspaceZip));
 			int read;
 			byte[] buffer = new byte[4096];
-			while((read = fileStream.read(buffer)) > -1)
+			while ((read = fileStream.read(buffer)) > -1)
 			{
-				if(read == 0)
+				if (read == 0)
 					continue;
-				Packet4FileData dataPacket = new Packet4FileData(fileId, (short)read, buffer);
+				Packet4FileData dataPacket = new Packet4FileData(fileId, (short) read, buffer);
 				enqueue(dataPacket);
 			}
-			enqueue(new Packet4FileData(fileId, (short)0, new byte[]{}));
+			enqueue(new Packet4FileData(fileId, (short) 0, new byte[] {}));
 		}
 		catch (IOException e)
 		{
@@ -87,20 +90,33 @@ public class ClientHandler extends NetworkHandler {
 
 	@Override
 	public void handlePacket(Packet4FileData packet) {
-		// TODO Auto-generated method stub
+		try
-		
+		{
+			this.appendFile(packet.getFileId(), packet.getFileChunk());
+		}
+		catch (IOException e)
+		{
+			e.printStackTrace();
+			disconnect(FILE_DOWNLOAD_FAILURE, "Failed to download transferred file");
+		}
 	}
 
 	@Override
 	public void handlePacket(Packet5Artifact packet) {
 		// TODO Auto-generated method stub
-		
+
 	}
 
 	@Override
 	public void handlePacket(Packet127KeepAlive packet) {
 		// TODO Auto-generated method stub
-		
+
+	}
+
+	@Override
+	public void handlePacket(Packet6Disconnect packet) {
+		// TODO Auto-generated method stub
+
 	}
 
 }

src/com/flaremicro/crossjeeves/net/NetworkHandler.java

@@ -1,9 +1,13 @@
 package com.flaremicro.crossjeeves.net;
 
+import java.io.BufferedOutputStream;
+import java.io.Closeable;
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
 import java.io.File;
+import java.io.FileOutputStream;
 import java.io.IOException;
+import java.io.OutputStream;
 import java.net.Socket;
 import java.util.HashMap;
 import java.util.concurrent.BlockingQueue;
@@ -17,15 +21,18 @@ import com.flaremicro.crossjeeves.net.packet.Packet2Script;
 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.Packet6Disconnect;
 import com.flaremicro.util.Util;
 
+import static com.flaremicro.crossjeeves.net.ErrorCodes.*;
+
 public abstract class NetworkHandler {
 	private boolean connected;
 	private Socket socket;
 	private DataInputStream in;
 	private DataOutputStream out;
 	private Thread readThread;
-	private HashMap<Long, File> downloadQueue = new HashMap<Long, File>();
+	private HashMap<Long, FileTransferInfo> downloadQueue = new HashMap<Long, FileTransferInfo>();
 	private HashMap<Long, File> downloadComplete = new HashMap<Long, File>();
 
 	private BlockingQueue<Packet> outbox = new LinkedBlockingQueue<Packet>();
@@ -35,16 +42,15 @@ public abstract class NetworkHandler {
 		this.in = new DataInputStream(socket.getInputStream());
 		this.out = new DataOutputStream(socket.getOutputStream());
 	}
-	
+
-	public File waitForFile(long fileId)
+	public File waitForFile(long fileId) {
-	{
 		try
 		{
-			while(true)
+			while (true)
 			{
 				downloadComplete.wait();
 				File file = downloadComplete.get(fileId);
-				if(file != null)
+				if (file != null)
 					return file;
 			}
 		}
@@ -67,8 +73,12 @@ public abstract class NetworkHandler {
 		packet.processPacket(this);
 	}
 
-	//TODO error code?
+	public void disconnect(int code, String message) {
-	protected void disconnect() {
+		System.out.println("Disconnect code " + code + ": " + message);
+		doDisconnect();
+	}
+
+	protected void doDisconnect() {
 		connected = false;
 		readThread.interrupt();
 		Util.cleanClose(in);
@@ -93,7 +103,7 @@ public abstract class NetworkHandler {
 			catch (IOException e)
 			{
 				e.printStackTrace();
-				disconnect();
+				disconnect(READ_FAILED, e.toString());
 			}
 		}
 	}
@@ -114,14 +124,14 @@ public abstract class NetworkHandler {
 							}
 							catch (IOException e)
 							{
-								disconnect();
+								disconnect(WRITE_FAILED, e.toString());
 								e.printStackTrace();
 							}
 						}
 					}
 					catch (InterruptedException e)
 					{
-						disconnect();
+						disconnect(THREAD_INTERRUPTED, e.toString());
 						e.printStackTrace();
 					}
 				}
@@ -144,6 +154,74 @@ public abstract class NetworkHandler {
 
 	public abstract void handlePacket(Packet5Artifact packet);
 
-	public abstract void handlePacket(Packet127KeepAlive packet);
+	public abstract void handlePacket(Packet6Disconnect packet);
 
+	public abstract void handlePacket(Packet127KeepAlive packet);
+	
+	public void clearFile(long fileId)
+	{
+		FileTransferInfo fileTransferInfo = downloadQueue.get(fileId);
+		if(fileTransferInfo != null)
+		{
+			Util.cleanClose(fileTransferInfo);
+			fileTransferInfo.file.delete();
+		}
+		File file = downloadComplete.get(fileId);
+		if(file != null)
+		{
+			file.delete();
+		}
+	}
+	
+	public void beginFile(long fileId, long expectedSize) throws IOException
+	{
+		File file = File.createTempFile(fileId+"-cj-tmp", ".zip");
+		file.deleteOnExit();
+		OutputStream is = null;
+		try
+		{
+			is = new BufferedOutputStream(new FileOutputStream(file));
+			FileTransferInfo fti = new FileTransferInfo(file, is, expectedSize);
+			this.downloadQueue.put(fileId, fti);
+		}
+		catch(IOException ex)
+		{
+			Util.cleanClose(is);
+			throw ex;
+		}
+	}
+	
+	public boolean appendFile(long fileId, byte[] chunk) throws IOException
+	{
+		FileTransferInfo fileTransferInfo = downloadQueue.get(fileId);
+		if(fileTransferInfo == null)
+			return false;
+		if(chunk.length == 0)
+		{
+			Util.cleanClose(fileTransferInfo);
+			this.downloadComplete.put(fileId, fileTransferInfo.file);
+			this.downloadComplete.notifyAll();
+		}
+		else
+		{
+			fileTransferInfo.outputStream.write(chunk);
+		}
+		return true;
+	}
+	
+	class FileTransferInfo implements Closeable{
+		public final File file;
+		public final OutputStream outputStream;
+		public final long expectedSize;
+		FileTransferInfo(File file, OutputStream outputStream, long expectedSize)
+		{
+			this.file = file;
+			this.outputStream = outputStream;
+			this.expectedSize = expectedSize;
+		}
+		
+		public void close() throws IOException {
+			outputStream.close();
+		}
+	}
 }

src/com/flaremicro/crossjeeves/net/ServerHandler.java

@@ -1,8 +1,12 @@
 package com.flaremicro.crossjeeves.net;
 
+import static com.flaremicro.crossjeeves.net.ErrorCodes.*;
+
 import java.io.IOException;
 import java.net.Socket;
+import java.util.Properties;
 
+import com.flaremicro.crossjeeves.ScriptProcessor;
 import com.flaremicro.crossjeeves.net.packet.Packet;
 import com.flaremicro.crossjeeves.net.packet.Packet0Identify;
 import com.flaremicro.crossjeeves.net.packet.Packet127KeepAlive;
@@ -11,17 +15,18 @@ import com.flaremicro.crossjeeves.net.packet.Packet2Script;
 import com.flaremicro.crossjeeves.net.packet.Packet3Clone;
 import com.flaremicro.crossjeeves.net.packet.Packet5Artifact;
 import com.flaremicro.crossjeeves.net.packet.Packet4FileData;
+import com.flaremicro.crossjeeves.net.packet.Packet6Disconnect;
 
 public class ServerHandler extends NetworkHandler {
-
+	Properties properties;
-	public ServerHandler(Socket socket) throws IOException {
+	public ServerHandler(Socket socket, Properties properties) throws IOException {
 		super(socket);
+		this.properties = properties;
 	}
 
 	@Override
 	public void handlePacket(Packet packet) {
-		// TODO Auto-generated method stub
+		disconnect(INVALID_PACKET_RECIEVED, "Recieved invalid packet " + packet.getId() + " (Unknown)");
-		
 	}
 
 	@Override
@@ -37,28 +42,39 @@ public class ServerHandler extends NetworkHandler {
 
 	@Override
 	public void handlePacket(Packet2Script packet) {
-		// TODO Auto-generated method stub
+		ScriptProcessor scriptProcessor = new ScriptProcessor(this, properties);
-		
+		scriptProcessor.processAsync(packet.script);
 	}
 
 	@Override
 	public void handlePacket(Packet3Clone packet) {
+		disconnect(INVALID_PACKET_RECIEVED, "Recieved invalid packet " + packet.getId() + " (Clone)");
 	}
 
 	@Override
 	public void handlePacket(Packet4FileData packet) {
-		// TODO Auto-generated method stub
+		try
-		
+		{
+			this.appendFile(packet.getFileId(), packet.getFileChunk());
+		}
+		catch (IOException e)
+		{
+			e.printStackTrace();
+			disconnect(FILE_DOWNLOAD_FAILURE, "Failed to download transferred file");
+		}
 	}
 
 	@Override
 	public void handlePacket(Packet5Artifact packet) {
-		// TODO Auto-generated method stub
+		disconnect(INVALID_PACKET_RECIEVED, "Recieved invalid packet " + packet.getId() + " (Artifact)");
-		
 	}
 
 	@Override
 	public void handlePacket(Packet127KeepAlive packet) {
+	}
+
+	@Override
+	public void handlePacket(Packet6Disconnect packet) {
 		// TODO Auto-generated method stub
 		
 	}

src/com/flaremicro/crossjeeves/net/packet/Packet.java

@@ -23,6 +23,8 @@ public abstract class Packet {
 		registerPacket((byte) 3, new Packet3Clone());
 		registerPacket((byte) 4, new Packet4FileData());
 		registerPacket((byte) 5, new Packet5Artifact());
+		registerPacket((byte) 6, new Packet6Disconnect());
+		registerPacket((byte) 7, new Packet7LogEntry());
 		
 		registerPacket((byte) 127, new Packet127KeepAlive());
 	}

test/com/flaremicro/crossjeeves/net/packet/test/Packet0IdentifyTest.java

@@ -31,7 +31,6 @@ public class Packet0IdentifyTest extends PacketTestBase {
 		assertEquals(0, input().readByte());
 		assertEquals(Packet.PROTOCOL_VERSION, input().readInt());
 		assertEquals(1337, input().readInt());
-		assertBufferEmpty();
 	}
 	
 	@Test
@@ -46,7 +45,6 @@ public class Packet0IdentifyTest extends PacketTestBase {
 
 		assertEquals(Packet.PROTOCOL_VERSION, packet.getProtocolVersion());
 		assertEquals(1337, packet.getFlags());
-		assertBufferEmpty();
 	}
 	
 	@Test

test/com/flaremicro/crossjeeves/net/packet/test/Packet1StatusTest.java

@@ -29,7 +29,6 @@ public class Packet1StatusTest extends PacketTestBase {
 
 		assertEquals(1, input().readByte());
 		assertEquals(1337, input().readInt());
-		assertBufferEmpty();
 	}
 	
 	@Test
@@ -42,7 +41,6 @@ public class Packet1StatusTest extends PacketTestBase {
 		packet.recievePacket(input());
 
 		assertEquals(1337, packet.getFlags());
-		assertBufferEmpty();
 	}
 	
 	@Test

test/com/flaremicro/crossjeeves/net/packet/test/Packet2ScriptTest.java

@@ -29,7 +29,6 @@ public class Packet2ScriptTest extends PacketTestBase {
 
 		assertEquals(2, input().readByte());
 		assertEquals("1337", input().readUTF());
-		assertBufferEmpty();
 	}
 	
 	@Test
@@ -44,7 +43,6 @@ public class Packet2ScriptTest extends PacketTestBase {
 		packet.recievePacket(input());
 
 		assertEquals("1337", packet.script);
-		assertBufferEmpty();
 	}
 	
 	@Test

test/com/flaremicro/crossjeeves/net/packet/test/Packet3CloneTest.java

@@ -31,7 +31,6 @@ public class Packet3CloneTest extends PacketTestBase {
 		assertEquals(3, input().readByte());
 		assertEquals(1, input().readLong());
 		assertEquals(2, input().readLong());
-		assertBufferEmpty();
 	}
 	
 	@Test
@@ -46,7 +45,6 @@ public class Packet3CloneTest extends PacketTestBase {
 
 		assertEquals(1, packet.getFileID());
 		assertEquals(2, packet.getFileSize());
-		assertBufferEmpty();
 	}
 	
 	@Test

test/com/flaremicro/crossjeeves/net/packet/test/Packet4FileDataTest.java

@@ -36,7 +36,6 @@ public class Packet4FileDataTest extends PacketTestBase {
 		assertEquals(4, input().readShort());
 		input().readFully(buffer);
 		assertArrayEquals(new byte[]{1,2,3,4}, buffer);
-		assertBufferEmpty();
 	}
 	
 	@Test
@@ -52,7 +51,6 @@ public class Packet4FileDataTest extends PacketTestBase {
 
 		assertEquals(1, packet.getFileId());
 		assertArrayEquals(new byte[]{1,2,3,4}, packet.getFileChunk());
-		assertBufferEmpty();
 	}
 	
 	@Test

test/com/flaremicro/crossjeeves/net/packet/test/Packet5ArtifactTest.java

@@ -30,7 +30,6 @@ public class Packet5ArtifactTest extends PacketTestBase {
 		assertEquals(5, input().readByte());
 		assertEquals(1, input().readLong());
 		assertEquals("file.txt", input().readUTF());
-		assertBufferEmpty();
 	}
 	
 	@Test
@@ -45,7 +44,6 @@ public class Packet5ArtifactTest extends PacketTestBase {
 
 		assertEquals(1, packet.getFileId());
 		assertEquals("file.txt", packet.getRelativeFile());
-		assertBufferEmpty();
 	}
 	
 	@Test

test/com/flaremicro/crossjeeves/net/packet/test/PacketTestBase.java

@@ -6,6 +6,7 @@ import java.io.IOException;
 
 import static org.junit.Assert.assertEquals;
 
+import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.Before;
 
@@ -29,9 +30,10 @@ public abstract class PacketTestBase {
 	{
 		buffer.empty();
 	}
-
+	
-	public static void assertBufferEmpty() {
+	@After
-		assertEquals(0, buffer.size());
+	public void assertBufferEmpty() {
+		assertEquals("Buffer should be empty:" ,0, buffer.size());
 	}
 
 	protected DataInputStream input() {