From b3289186844370724d8aa0c1173ab39a048d2a60 Mon Sep 17 00:00:00 2001 From: Flare Microsystems Date: Mon, 11 Nov 2024 08:46:26 -0800 Subject: [PATCH] Add more net code --- .../crossjeeves/ScriptProcessor.java | 64 +++++++---- .../crossjeeves/net/ClientHandler.java | 46 +++++--- .../crossjeeves/net/NetworkHandler.java | 102 +++++++++++++++--- .../crossjeeves/net/ServerHandler.java | 36 +++++-- .../crossjeeves/net/packet/Packet.java | 2 + .../net/packet/test/Packet0IdentifyTest.java | 2 - .../net/packet/test/Packet1StatusTest.java | 2 - .../net/packet/test/Packet2ScriptTest.java | 2 - .../net/packet/test/Packet3CloneTest.java | 2 - .../net/packet/test/Packet4FileDataTest.java | 2 - .../net/packet/test/Packet5ArtifactTest.java | 2 - .../net/packet/test/PacketTestBase.java | 8 +- 12 files changed, 197 insertions(+), 73 deletions(-) diff --git a/src/com/flaremicro/crossjeeves/ScriptProcessor.java b/src/com/flaremicro/crossjeeves/ScriptProcessor.java index 23c4d4d..de83261 100644 --- a/src/com/flaremicro/crossjeeves/ScriptProcessor.java +++ b/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(); - - public ScriptProcessor(NetworkHandler netHandler, Properties properties) - { + private Thread executionThread = null; + + 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 environment = processBuilder.environment(); - - for(Entry set : currentEnvironment.entrySet()) + + for (Entry set : currentEnvironment.entrySet()) { environment.put(set.getKey(), set.getValue()); } - - if(runningProcess != null) - throw new ScriptProcessingException("Previous process not terminated"); - try + + if (runningProcess != null) + throw new ScriptProcessingException("Previous process not terminated"); + 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()); - if(returnCode != 0) + int returnCode = beginProcess(properties.getProperty(shellName + "-executor", shellName), file.getAbsolutePath()); + 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(); + } + } } diff --git a/src/com/flaremicro/crossjeeves/net/ClientHandler.java b/src/com/flaremicro/crossjeeves/net/ClientHandler.java index bdde0fb..52063f2 100644 --- a/src/com/flaremicro/crossjeeves/net/ClientHandler.java +++ b/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 + } } diff --git a/src/com/flaremicro/crossjeeves/net/NetworkHandler.java b/src/com/flaremicro/crossjeeves/net/NetworkHandler.java index c3706eb..4057f73 100644 --- a/src/com/flaremicro/crossjeeves/net/NetworkHandler.java +++ b/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 downloadQueue = new HashMap(); + private HashMap downloadQueue = new HashMap(); private HashMap downloadComplete = new HashMap(); private BlockingQueue outbox = new LinkedBlockingQueue(); @@ -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? - protected void disconnect() { + public void disconnect(int code, String message) { + 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(); + } + } } diff --git a/src/com/flaremicro/crossjeeves/net/ServerHandler.java b/src/com/flaremicro/crossjeeves/net/ServerHandler.java index 2a24712..4f20fc8 100644 --- a/src/com/flaremicro/crossjeeves/net/ServerHandler.java +++ b/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 { - - public ServerHandler(Socket socket) throws IOException { + Properties properties; + 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 } diff --git a/src/com/flaremicro/crossjeeves/net/packet/Packet.java b/src/com/flaremicro/crossjeeves/net/packet/Packet.java index 5319b4c..c3a8e38 100644 --- a/src/com/flaremicro/crossjeeves/net/packet/Packet.java +++ b/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()); } diff --git a/test/com/flaremicro/crossjeeves/net/packet/test/Packet0IdentifyTest.java b/test/com/flaremicro/crossjeeves/net/packet/test/Packet0IdentifyTest.java index 346a069..491e511 100644 --- a/test/com/flaremicro/crossjeeves/net/packet/test/Packet0IdentifyTest.java +++ b/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 diff --git a/test/com/flaremicro/crossjeeves/net/packet/test/Packet1StatusTest.java b/test/com/flaremicro/crossjeeves/net/packet/test/Packet1StatusTest.java index 4245524..bed7ddf 100644 --- a/test/com/flaremicro/crossjeeves/net/packet/test/Packet1StatusTest.java +++ b/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 diff --git a/test/com/flaremicro/crossjeeves/net/packet/test/Packet2ScriptTest.java b/test/com/flaremicro/crossjeeves/net/packet/test/Packet2ScriptTest.java index 1974b6e..4fbc43e 100644 --- a/test/com/flaremicro/crossjeeves/net/packet/test/Packet2ScriptTest.java +++ b/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 diff --git a/test/com/flaremicro/crossjeeves/net/packet/test/Packet3CloneTest.java b/test/com/flaremicro/crossjeeves/net/packet/test/Packet3CloneTest.java index 43a8ec5..fdb7d1f 100644 --- a/test/com/flaremicro/crossjeeves/net/packet/test/Packet3CloneTest.java +++ b/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 diff --git a/test/com/flaremicro/crossjeeves/net/packet/test/Packet4FileDataTest.java b/test/com/flaremicro/crossjeeves/net/packet/test/Packet4FileDataTest.java index ce7f931..8e5ba9a 100644 --- a/test/com/flaremicro/crossjeeves/net/packet/test/Packet4FileDataTest.java +++ b/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 diff --git a/test/com/flaremicro/crossjeeves/net/packet/test/Packet5ArtifactTest.java b/test/com/flaremicro/crossjeeves/net/packet/test/Packet5ArtifactTest.java index e4031ed..f203252 100644 --- a/test/com/flaremicro/crossjeeves/net/packet/test/Packet5ArtifactTest.java +++ b/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 diff --git a/test/com/flaremicro/crossjeeves/net/packet/test/PacketTestBase.java b/test/com/flaremicro/crossjeeves/net/packet/test/PacketTestBase.java index 5c196eb..fa5d267 100644 --- a/test/com/flaremicro/crossjeeves/net/packet/test/PacketTestBase.java +++ b/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() { - assertEquals(0, buffer.size()); + + @After + public void assertBufferEmpty() { + assertEquals("Buffer should be empty:" ,0, buffer.size()); } protected DataInputStream input() {