From 761785620f6e9d2c65557dde0dd1b929e85dedcf Mon Sep 17 00:00:00 2001
From: Tingluo Huang <tingluohuang@github.com>
Date: Wed, 22 Jun 2022 10:59:29 -0400
Subject: [PATCH] Support pass runner JitConfig as arg. (#1925)

---
 src/Runner.Common/Constants.cs         |  1 +
 src/Runner.Listener/CommandSettings.cs |  7 +++++++
 src/Runner.Listener/Runner.cs          | 26 ++++++++++++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/src/Runner.Common/Constants.cs b/src/Runner.Common/Constants.cs
index f7aca5158..1ee60714e 100644
--- a/src/Runner.Common/Constants.cs
+++ b/src/Runner.Common/Constants.cs
@@ -90,6 +90,7 @@ namespace GitHub.Runner.Common
                 public static class Args
                 {
                     public static readonly string Auth = "auth";
+                    public static readonly string JitConfig = "jitconfig";
                     public static readonly string Labels = "labels";
                     public static readonly string MonitorSocketAddress = "monitorsocketaddress";
                     public static readonly string Name = "name";
diff --git a/src/Runner.Listener/CommandSettings.cs b/src/Runner.Listener/CommandSettings.cs
index 5067cf921..368b97b5e 100644
--- a/src/Runner.Listener/CommandSettings.cs
+++ b/src/Runner.Listener/CommandSettings.cs
@@ -39,6 +39,7 @@ namespace GitHub.Runner.Listener
                     Constants.Runner.CommandLine.Flags.RunAsService,
                     Constants.Runner.CommandLine.Flags.Unattended,
                     Constants.Runner.CommandLine.Args.Auth,
+                    Constants.Runner.CommandLine.Args.JitConfig,
                     Constants.Runner.CommandLine.Args.Labels,
                     Constants.Runner.CommandLine.Args.MonitorSocketAddress,
                     Constants.Runner.CommandLine.Args.Name,
@@ -213,6 +214,12 @@ namespace GitHub.Runner.Listener
                 validator: Validators.AuthSchemeValidator);
         }
 
+        public string GetJitConfig()
+        {
+            return GetArg(
+                name: Constants.Runner.CommandLine.Args.JitConfig);
+        }
+
         public string GetRunnerName()
         {
             return GetArgOrPrompt(
diff --git a/src/Runner.Listener/Runner.cs b/src/Runner.Listener/Runner.cs
index f7505bc08..cf12c91ab 100644
--- a/src/Runner.Listener/Runner.cs
+++ b/src/Runner.Listener/Runner.cs
@@ -1,9 +1,11 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.IO.Compression;
 using System.Linq;
 using System.Reflection;
 using System.Runtime.CompilerServices;
+using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using GitHub.DistributedTask.WebApi;
@@ -192,6 +194,30 @@ namespace GitHub.Runner.Listener
                     return Constants.Runner.ReturnCode.Success;
                 }
 
+                var base64JitConfig = command.GetJitConfig();
+                if (!string.IsNullOrEmpty(base64JitConfig))
+                {
+                    try
+                    {
+                        var decodedJitConfig = Encoding.UTF8.GetString(Convert.FromBase64String(base64JitConfig));
+                        var jitConfig = StringUtil.ConvertFromJson<Dictionary<string, string>>(decodedJitConfig);
+                        foreach (var config in jitConfig)
+                        {
+                            var configFile = Path.Combine(HostContext.GetDirectory(WellKnownDirectory.Root), config.Key);
+                            var configContent = Encoding.UTF8.GetString(Convert.FromBase64String(config.Value));
+                            File.WriteAllText(configFile, configContent, Encoding.UTF8);
+                            File.SetAttributes(configFile, File.GetAttributes(configFile) | FileAttributes.Hidden);
+                            Trace.Info($"Save {configContent.Length} chars to '{configFile}'.");
+                        }
+                    }
+                    catch (Exception ex)
+                    {
+                        Trace.Error(ex);
+                        _term.WriteError(ex.Message);
+                        return Constants.Runner.ReturnCode.TerminatedError;
+                    }
+                }
+
                 RunnerSettings settings = configManager.LoadSettings();
 
                 var store = HostContext.GetService<IConfigurationStore>();