Add more net code

This commit is contained in:
Flare Microsystems
2024-11-11 08:46:26 -08:00
parent d8f7a36bc2
commit b328918684
12 changed files with 197 additions and 73 deletions

View File

@@ -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;
}
@@ -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
{
@@ -191,17 +191,16 @@ 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");
try
{
@@ -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();
}
}
}

View File

@@ -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;
@@ -51,28 +54,28 @@ public class ClientHandler extends NetworkHandler {
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,8 +90,15 @@ 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
@@ -103,4 +113,10 @@ public class ClientHandler extends NetworkHandler {
}
@Override
public void handlePacket(Packet6Disconnect packet) {
// TODO Auto-generated method stub
}
}

View File

@@ -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>();
@@ -36,15 +43,14 @@ public abstract class NetworkHandler {
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(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();
}
}
}

View File

@@ -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
}

View File

@@ -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());
}

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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;
@@ -30,8 +31,9 @@ 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() {