Adds envVars propagation client -> server

Since Mill now executes in a long-lived JVM, the builds do not have a
chance to use environment variables as inputs. This propagates the
environment variables from the client all the way down to the context
available to the tasks as a `Map[String, String]` so that they can be
used as inputs should the user choose to do so.

https://github.com/lihaoyi/mill/issues/257

10 files changed, 111 insertions, 30 deletions

diff --git a/clientserver/src/mill/clientserver/Client.java b/clientserver/src/mill/clientserver/Client.java
index c4cbc265..c570a65c 100644
--- a/clientserver/src/mill/clientserver/Client.java
+++ b/clientserver/src/mill/clientserver/Client.java
@@ -101,6 +101,7 @@ public class Client {
 
         FileOutputStream f = new FileOutputStream(lockBase + "/run");
         ClientServer.writeArgs(System.console() != null, args, f);
+        ClientServer.writeEnv(f);
         f.close();
 
         boolean serverInit = false;
diff --git a/clientserver/src/mill/clientserver/ClientServer.java b/clientserver/src/mill/clientserver/ClientServer.java
index e2e63dcf..45c99ea0 100644
--- a/clientserver/src/mill/clientserver/ClientServer.java
+++ b/clientserver/src/mill/clientserver/ClientServer.java
@@ -4,6 +4,8 @@ package mill.clientserver;
 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,77 @@ 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 writeEnv(OutputStream argStream) throws IOException {
+        Map<String, String> env = System.getenv();
+        argStream.write(env.size());
+        for (Map.Entry<String, String> kv : env.entrySet()) {
+            writeString(argStream, kv.getKey());
+            writeString(argStream, kv.getValue());
+        }
     }
-  }
-}
\ No newline at end of file
+
+    public static Map<String, String> parseEnv(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/clientserver/src/mill/clientserver/Server.scala b/clientserver/src/mill/clientserver/Server.scala
index 24827ac2..d9e67489 100644
--- a/clientserver/src/mill/clientserver/Server.scala
+++ b/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.parseEnv(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
diff --git a/clientserver/test/src/mill/clientserver/ClientServerTests.scala b/clientserver/test/src/mill/clientserver/ClientServerTests.scala
index 2c9a57b0..0b55d4bb 100644
--- a/clientserver/test/src/mill/clientserver/ClientServerTests.scala
+++ b/clientserver/test/src/mill/clientserver/ClientServerTests.scala
@@ -9,7 +9,8 @@ class EchoServer extends ServerMain[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()
diff --git a/core/src/mill/eval/Evaluator.scala b/core/src/mill/eval/Evaluator.scala
index 33141c0a..7b3634ad 100644
--- a/core/src/mill/eval/Evaluator.scala
+++ b/core/src/mill/eval/Evaluator.scala
@@ -2,6 +2,8 @@ package mill.eval
 
 import java.net.URLClassLoader
 
+import scala.collection.JavaConverters._
+
 import mill.util.Router.EntryPoint
 import ammonite.ops._
 import ammonite.runtime.SpecialClassLoader
@@ -32,7 +34,8 @@ case class Evaluator[T](home: Path,
                         rootModule: mill.define.BaseModule,
                         log: Logger,
                         classLoaderSig: Seq[(Either[String, Path], Long)] = Evaluator.classLoaderSig,
-                        workerCache: mutable.Map[Segments, (Int, Any)] = mutable.Map.empty){
+                        workerCache: mutable.Map[Segments, (Int, Any)] = mutable.Map.empty,
+                        env : Map[String, String] = Evaluator.defaultEnv){
   val classLoaderSignHash = classLoaderSig.hashCode()
   def evaluate(goals: Agg[Task[_]]): Evaluator.Results = {
     mkdir(outPath)
@@ -271,7 +274,8 @@ case class Evaluator[T](home: Path,
                 }
             },
             multiLogger,
-            home
+            home,
+            env
           )
 
           val out = System.out
@@ -335,6 +339,8 @@ object Evaluator{
   // in directly) we are forced to pass it in via a ThreadLocal
   val currentEvaluator = new ThreadLocal[mill.eval.Evaluator[_]]
 
+  val defaultEnv: Map[String, String] = System.getenv().asScala.toMap
+
   case class Paths(out: Path,
                    dest: Path,
                    meta: Path,
diff --git a/core/src/mill/util/Ctx.scala b/core/src/mill/util/Ctx.scala
index 99818194..88a8baec 100644
--- a/core/src/mill/util/Ctx.scala
+++ b/core/src/mill/util/Ctx.scala
@@ -36,7 +36,8 @@ object Ctx{
 class Ctx(val args: IndexedSeq[_],
           dest0: () => Path,
           val log: Logger,
-          val home: Path)
+          val home: Path,
+          val env : Map[String, String])
   extends Ctx.Dest
   with Ctx.Log
   with Ctx.Args
diff --git a/main/src/mill/Main.scala b/main/src/mill/Main.scala
index e026dfe0..c9ec00ca 100644
--- a/main/src/mill/Main.scala
+++ b/main/src/mill/Main.scala
@@ -2,6 +2,8 @@ package mill
 
 import java.io.{InputStream, PrintStream}
 
+import scala.collection.JavaConverters._
+
 import ammonite.main.Cli._
 import ammonite.ops._
 import ammonite.util.Util
@@ -16,13 +18,15 @@ object ServerMain extends mill.clientserver.ServerMain[Evaluator.State]{
             mainInteractive: Boolean,
             stdin: InputStream,
             stdout: PrintStream,
-            stderr: PrintStream) = Main.main0(
+            stderr: PrintStream,
+            env : Map[String, String]) = Main.main0(
     args,
     stateCache,
     mainInteractive,
     DummyInputStream,
     stdout,
-    stderr
+    stderr,
+    env
   )
 }
 object Main {
@@ -38,7 +42,8 @@ object Main {
       ammonite.Main.isInteractive(),
       System.in,
       System.out,
-      System.err
+      System.err,
+      System.getenv().asScala.toMap
     )
     System.exit(if(result) 0 else 1)
   }
@@ -48,7 +53,8 @@ object Main {
             mainInteractive: Boolean,
             stdin: InputStream,
             stdout: PrintStream,
-            stderr: PrintStream): (Boolean, Option[Evaluator.State]) = {
+            stderr: PrintStream,
+            env: Map[String, String]): (Boolean, Option[Evaluator.State]) = {
     import ammonite.main.Cli
 
     val removed = Set("predef-code", "no-home-predef")
@@ -116,7 +122,8 @@ object Main {
           val runner = new mill.main.MainRunner(
             config.copy(colored = Some(mainInteractive)),
             stdout, stderr, stdin,
-            stateCache
+            stateCache,
+            env
           )
 
           if (mill.clientserver.ClientServer.isJava9OrAbove) {
diff --git a/main/src/mill/main/MainRunner.scala b/main/src/mill/main/MainRunner.scala
index efebd5a5..fed664fd 100644
--- a/main/src/mill/main/MainRunner.scala
+++ b/main/src/mill/main/MainRunner.scala
@@ -20,7 +20,8 @@ class MainRunner(val config: ammonite.main.Cli.Config,
                  outprintStream: PrintStream,
                  errPrintStream: PrintStream,
                  stdIn: InputStream,
-                 stateCache0: Option[Evaluator.State] = None)
+                 stateCache0: Option[Evaluator.State] = None,
+                 env : Map[String, String])
   extends ammonite.MainRunner(
     config, outprintStream, errPrintStream,
     stdIn, outprintStream, errPrintStream
@@ -75,7 +76,8 @@ class MainRunner(val config: ammonite.main.Cli.Config,
             errPrintStream,
             errPrintStream,
             stdIn
-          )
+          ),
+          env
         )
 
         result match{
diff --git a/main/src/mill/main/RunScript.scala b/main/src/mill/main/RunScript.scala
index 77930cc8..75042dea 100644
--- a/main/src/mill/main/RunScript.scala
+++ b/main/src/mill/main/RunScript.scala
@@ -29,7 +29,8 @@ object RunScript{
                 instantiateInterpreter: => Either[(Res.Failing, Seq[(Path, Long)]), ammonite.interp.Interpreter],
                 scriptArgs: Seq[String],
                 stateCache: Option[Evaluator.State],
-                log: Logger)
+                log: Logger,
+                env : Map[String, String])
   : (Res[(Evaluator[Any], Seq[PathRef], Either[String, Seq[Js.Value]])], Seq[(Path, Long)]) = {
 
     val (evalState, interpWatched) = stateCache match{
@@ -53,7 +54,8 @@ object RunScript{
 
     val evalRes =
       for(s <- evalState)
-      yield new Evaluator[Any](home, wd / 'out, wd / 'out, s.rootModule, log, s.classLoaderSig, s.workerCache)
+      yield new Evaluator[Any](home, wd / 'out, wd / 'out, s.rootModule, log,
+        s.classLoaderSig, s.workerCache, env)
 
     val evaluated = for{
       evaluator <- evalRes
diff --git a/main/test/src/mill/util/ScriptTestSuite.scala b/main/test/src/mill/util/ScriptTestSuite.scala
index a2f2676a..f88007c5 100644
--- a/main/test/src/mill/util/ScriptTestSuite.scala
+++ b/main/test/src/mill/util/ScriptTestSuite.scala
@@ -15,7 +15,7 @@ abstract class ScriptTestSuite(fork: Boolean) extends TestSuite{
   val stdIn = new ByteArrayInputStream(Array())
   lazy val runner = new mill.main.MainRunner(
     ammonite.main.Cli.Config(wd = workspacePath),
-    stdOutErr, stdOutErr, stdIn
+    stdOutErr, stdOutErr, stdIn, None, Map.empty
   )
   def eval(s: String*) = {
     if (!fork) runner.runScript(workspacePath / "build.sc", s.toList)