merge

2 files changed, 76 insertions, 20 deletions

diff --git a/client/src/mill/client/ClientServer.java b/client/src/mill/client/ClientServer.java
index c30fc221..468f8ab3 100644
--- a/client/src/mill/client/ClientServer.java
+++ b/client/src/mill/client/ClientServer.java
@@ -4,6 +4,8 @@ package mill.client;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.util.HashMap;
+import java.util.Map;
 
 public class ClientServer {
     public static boolean isWindows = System.getProperty("os.name").toLowerCase().startsWith("windows");
@@ -20,23 +22,76 @@ public class ClientServer {
         int argsLength = argStream.read();
         String[] args = new String[argsLength];
         for (int i = 0; i < args.length; i++) {
-            int n = argStream.read();
-            byte[] arr = new byte[n];
-            argStream.read(arr);
-            args[i] = new String(arr);
+            args[i] = readString(argStream);
         }
         return args;
     }
     public static void writeArgs(Boolean interactive,
                                  String[] args,
-                                 OutputStream argStream) throws IOException{
-    argStream.write(interactive ? 1 : 0);
-    argStream.write(args.length);
-    int i = 0;
-    while (i < args.length){
-      argStream.write(args[i].length());
-      argStream.write(args[i].getBytes());
-      i += 1;
+                                 OutputStream argStream) throws IOException {
+        argStream.write(interactive ? 1 : 0);
+        argStream.write(args.length);
+        int i = 0;
+        while (i < args.length) {
+            writeString(argStream, args[i]);
+            i += 1;
+        }
+    }
+
+    /**
+     * This allows the mill client to pass the environment as he sees it to the
+     * server (as the server remains alive over the course of several runs and
+     * does not see the environment changes the client would)
+     */
+    public static void writeMap(Map<String, String> map, OutputStream argStream) throws IOException {
+        argStream.write(map.size());
+        for (Map.Entry<String, String> kv : map.entrySet()) {
+            writeString(argStream, kv.getKey());
+            writeString(argStream, kv.getValue());
+        }
     }
-  }
-}
\ No newline at end of file
+
+    public static Map<String, String> parseMap(InputStream argStream) throws IOException {
+        Map<String, String> env = new HashMap<>();
+        int mapLength = argStream.read();
+        for (int i = 0; i < mapLength; i++) {
+            String key = readString(argStream);
+            String value = readString(argStream);
+            env.put(key, value);
+        }
+        return env;
+    }
+
+    private static String readString(InputStream inputStream) throws IOException {
+        // Result is between 0 and 255, hence the loop.
+        int read = inputStream.read();
+        int bytesToRead = read;
+        while(read == 255){
+            read = inputStream.read();
+            bytesToRead += read;
+        }
+        byte[] arr = new byte[bytesToRead];
+        int readTotal = 0;
+        while (readTotal < bytesToRead) {
+            read = inputStream.read(arr, readTotal, bytesToRead - readTotal);
+            readTotal += read;
+        }
+        return new String(arr);
+    }
+
+    private static void writeString(OutputStream outputStream, String string) throws IOException {
+        // When written, an int > 255 gets splitted. This logic performs the
+        // split beforehand so that the reading side knows that there is still
+        // more metadata to come before it's able to read the actual data.
+        // Could do with rewriting using logical masks / shifts.
+        byte[] bytes = string.getBytes();
+        int toWrite = bytes.length;
+        while(toWrite >= 255){
+            outputStream.write(255);
+            toWrite = toWrite - 255;
+        }
+        outputStream.write(toWrite);
+        outputStream.write(bytes);
+    }
+
+}
diff --git a/client/src/mill/client/Main.java b/client/src/mill/client/Main.java
index a26b653e..109a9a9d 100644
--- a/client/src/mill/client/Main.java
+++ b/client/src/mill/client/Main.java
@@ -7,10 +7,7 @@ import java.net.Socket;
 import java.net.URISyntaxException;
 import java.net.URL;
 import java.nio.channels.FileChannel;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Iterator;
-import java.util.Properties;
+import java.util.*;
 
 public class Main {
     static void initServer(String lockBase, boolean setJnaNoSys) throws IOException,URISyntaxException{
@@ -52,6 +49,7 @@ public class Main {
     }
     public static void main(String[] args) throws Exception{
         boolean setJnaNoSys = System.getProperty("jna.nosys") == null;
+        Map<String, String> env = System.getenv();
         if (setJnaNoSys) {
             System.setProperty("jna.nosys", "true");
         }
@@ -83,7 +81,8 @@ public class Main {
                         System.in,
                         System.out,
                         System.err,
-                        args
+                        args,
+                        env
                 );
                 System.exit(exitCode);
             }
@@ -98,10 +97,12 @@ public class Main {
                           InputStream stdin,
                           OutputStream stdout,
                           OutputStream stderr,
-                          String[] args) throws Exception{
+                          String[] args,
+                          Map<String, String> env) throws Exception{
 
         FileOutputStream f = new FileOutputStream(lockBase + "/run");
         ClientServer.writeArgs(System.console() != null, args, f);
+        ClientServer.writeMap(env, f);
         f.close();
 
         boolean serverInit = false;