Adds envVars propagation client -> server

clientserver/src/mill/clientserver/Client.java

@@ -7,10 +7,7 @@
 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 Client {
     static void initServer(String lockBase, boolean setJnaNoSys) throws IOException,URISyntaxException{
@@ -51,6 +48,7 @@ static void initServer(String lockBase, boolean setJnaNoSys) throws IOException,
     }
     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");
         }
@@ -82,7 +80,8 @@ public void run() {
                         System.in,
                         System.out,
                         System.err,
-                        args
+                        args,
+                        env
                 );
                 System.exit(exitCode);
             }
@@ -97,10 +96,12 @@ public static int run(String lockBase,
                           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;

clientserver/src/mill/clientserver/ClientServer.java

@@ -4,6 +4,8 @@
 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 static String[] parseArgs(InputStream argStream) throws IOException {
         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());
+        }
     }
-  }
-}
+
+    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);
+    }
+
+}

clientserver/src/mill/clientserver/Server.scala

@@ -3,6 +3,7 @@ package mill.clientserver
 import java.io._
 import java.net.Socket
 
+import scala.collection.JavaConverters._
 import org.scalasbt.ipcsocket._
 
 trait ServerMain[T]{
@@ -21,7 +22,8 @@ trait ServerMain[T]{
             mainInteractive: Boolean,
             stdin: InputStream,
             stdout: PrintStream,
-            stderr: PrintStream): (Boolean, Option[T])
+            stderr: PrintStream,
+            env : Map[String, String]): (Boolean, Option[T])
 }
 
 
@@ -76,6 +78,7 @@ class Server[T](lockBase: String,
     val argStream = new FileInputStream(lockBase + "/run")
     val interactive = argStream.read() != 0;
     val args = ClientServer.parseArgs(argStream)
+    val env = ClientServer.parseMap(argStream)
     argStream.close()
 
     var done = false
@@ -89,7 +92,9 @@ class Server[T](lockBase: String,
           sm.stateCache,
           interactive,
           socketIn,
-          stdout, stderr
+          stdout,
+          stderr,
+          env.asScala.toMap
         )
 
         sm.stateCache = newStateCache

clientserver/test/src/mill/clientserver/ClientServerTests.scala

@@ -2,20 +2,29 @@ package mill.clientserver
 import java.io._
 import java.nio.file.Path
 
+import scala.collection.JavaConverters._
 import utest._
 class EchoServer extends ServerMain[Int]{
   def main0(args: Array[String],
             stateCache: Option[Int],
             mainInteractive: Boolean,
             stdin: InputStream,
             stdout: PrintStream,
-            stderr: PrintStream) = {
+            stderr: PrintStream,
+            env: Map[String, String]) = {
 
     val reader = new BufferedReader(new InputStreamReader(stdin))
     val str = reader.readLine()
-    stdout.println(str + args(0))
+    if (args.nonEmpty){
+      stdout.println(str + args(0))
+    }
+    env.toSeq.sortBy(_._1).foreach{
+      case (key, value) => stdout.println(s"$key=$value")
+    }
     stdout.flush()
-    stderr.println(str.toUpperCase + args(0))
+    if (args.nonEmpty){
+      stderr.println(str.toUpperCase + args(0))
+    }
     stderr.flush()
     (true, None)
   }
@@ -35,37 +44,39 @@ object ClientServerTests extends TestSuite{
     (tmpDir, locks)
   }
 
+  def spawnEchoServer(tmpDir : Path, locks: Locks): Unit = {
+    new Thread(() => new Server(
+      tmpDir.toString,
+      new EchoServer(),
+      () => (),
+      1000,
+      locks
+    ).run()).start()
+  }
+
+  def runClientAux(tmpDir : Path, locks: Locks)
+                  (env : Map[String, String], args: Array[String]) = {
+    val (in, out, err) = initStreams()
+    Server.lockBlock(locks.clientLock){
+      Client.run(
+        tmpDir.toString,
+        () => spawnEchoServer(tmpDir, locks),
+        locks,
+        in,
+        out,
+        err,
+        args,
+        env.asJava
+      )
+      Thread.sleep(100)
+      (new String(out.toByteArray), new String(err.toByteArray))
+    }
+  }
+
   def tests = Tests{
     'hello - {
       val (tmpDir, locks) = init()
-
-      def spawnEchoServer(): Unit = {
-        new Thread(() => new Server(
-          tmpDir.toString,
-          new EchoServer(),
-          () => (),
-          1000,
-          locks
-        ).run()).start()
-      }
-
-
-      def runClient(arg: String) = {
-        val (in, out, err) = initStreams()
-        Server.lockBlock(locks.clientLock){
-          Client.run(
-            tmpDir.toString,
-            () => spawnEchoServer(),
-            locks,
-            in,
-            out,
-            err,
-            Array(arg)
-          )
-          Thread.sleep(100)
-          (new String(out.toByteArray), new String(err.toByteArray))
-        }
-      }
+      def runClient(s: String) = runClientAux(tmpDir, locks)(Map.empty, Array(s))
 
       // Make sure the simple "have the client start a server and
       // exchange one message" workflow works from end to end.
@@ -116,5 +127,83 @@ object ClientServerTests extends TestSuite{
         err3 == "HELLO World\n"
       )
     }
+
+    'envVars - {
+      val (tmpDir, locks) = init()
+
+      def runClient(env : Map[String, String]) = runClientAux(tmpDir, locks)(env, Array())
+
+      // Make sure the simple "have the client start a server and
+      // exchange one message" workflow works from end to end.
+
+      assert(
+        locks.clientLock.probe(),
+        locks.serverLock.probe(),
+        locks.processLock.probe()
+      )
+
+      def longString(s : String) = Array.fill(1000)(s).mkString
+      val b1000 = longString("b")
+      val c1000 = longString("c")
+      val a1000 = longString("a")
+
+      val env = Map(
+        "a" -> a1000,
+        "b" -> b1000,
+        "c" -> c1000
+      )
+
+
+      val (out1, err1) = runClient(env)
+      val expected = s"a=$a1000\nb=$b1000\nc=$c1000\n"
+
+      assert(
+        out1 == expected,
+        err1 == ""
+      )
+
+      // Give a bit of time for the server to release the lock and
+      // re-acquire it to signal to the client that it's done
+      Thread.sleep(100)
+
+      assert(
+        locks.clientLock.probe(),
+        !locks.serverLock.probe(),
+        !locks.processLock.probe()
+      )
+
+      val path = List(
+        "/Users/foo/Library/Haskell/bin",
+        "/usr/local/git/bin",
+        "/sw/bin/",
+        "/usr/local/bin",
+        "/usr/local/",
+        "/usr/local/sbin",
+        "/usr/local/mysql/bin",
+        "/usr/local/bin",
+        "/usr/bin",
+        "/bin",
+        "/usr/sbin",
+        "/sbin",
+        "/opt/X11/bin",
+        "/usr/local/MacGPG2/bin",
+        "/Library/TeX/texbin",
+        "/usr/local/bin/",
+        "/Users/foo/bin",
+        "/Users/foo/go/bin",
+        "~/.bloop"
+      )
+
+      val pathEnvVar = path.mkString(":")
+      val (out2, err2) = runClient(Map("PATH" -> pathEnvVar))
+
+      val expected2 = s"PATH=$pathEnvVar\n"
+
+      assert(
+        out2 == expected2,
+        err2 == ""
+      )
+
+    }
   }
 }