support authenticated package download (#920)

src/Runner.Common/RunnerServer.cs

@@ -45,8 +45,8 @@ namespace GitHub.Runner.Common
         Task<TaskAgentJobRequest> FinishAgentRequestAsync(int poolId, long requestId, Guid lockToken, DateTime finishTime, TaskResult result, CancellationToken cancellationToken);
 
         // agent package
-        Task<List<PackageMetadata>> GetPackagesAsync(string packageType, string platform, int top, CancellationToken cancellationToken);
+        Task<List<PackageMetadata>> GetPackagesAsync(string packageType, string platform, int top, bool includeToken, CancellationToken cancellationToken);
-        Task<PackageMetadata> GetPackageAsync(string packageType, string platform, string version, CancellationToken cancellationToken);
+        Task<PackageMetadata> GetPackageAsync(string packageType, string platform, string version, bool includeToken, CancellationToken cancellationToken);
 
         // agent update
         Task<TaskAgent> UpdateAgentUpdateStateAsync(int agentPoolId, int agentId, string currentState);
@@ -317,16 +317,16 @@ namespace GitHub.Runner.Common
         //-----------------------------------------------------------------
         // Agent Package
         //-----------------------------------------------------------------
-        public Task<List<PackageMetadata>> GetPackagesAsync(string packageType, string platform, int top, CancellationToken cancellationToken)
+        public Task<List<PackageMetadata>> GetPackagesAsync(string packageType, string platform, int top, bool includeToken, CancellationToken cancellationToken)
         {
             CheckConnection(RunnerConnectionType.Generic);
-            return _genericTaskAgentClient.GetPackagesAsync(packageType, platform, top, cancellationToken: cancellationToken);
+            return _genericTaskAgentClient.GetPackagesAsync(packageType, platform, top, includeToken, cancellationToken: cancellationToken);
         }
 
-        public Task<PackageMetadata> GetPackageAsync(string packageType, string platform, string version, CancellationToken cancellationToken)
+        public Task<PackageMetadata> GetPackageAsync(string packageType, string platform, string version, bool includeToken, CancellationToken cancellationToken)
         {
             CheckConnection(RunnerConnectionType.Generic);
-            return _genericTaskAgentClient.GetPackageAsync(packageType, platform, version, cancellationToken: cancellationToken);
+            return _genericTaskAgentClient.GetPackageAsync(packageType, platform, version, includeToken, cancellationToken: cancellationToken);
         }
 
         public Task<TaskAgent> UpdateAgentUpdateStateAsync(int agentPoolId, int agentId, string currentState)

src/Runner.Listener/SelfUpdater.cs

@@ -110,7 +110,7 @@ namespace GitHub.Runner.Listener
             // old server won't send target version as part of update message.
             if (string.IsNullOrEmpty(targetVersion))
             {
-                var packages = await _runnerServer.GetPackagesAsync(_packageType, _platform, 1, token);
+                var packages = await _runnerServer.GetPackagesAsync(_packageType, _platform, 1, true, token);
                 if (packages == null || packages.Count == 0)
                 {
                     Trace.Info($"There is no package for {_packageType} and {_platform}.");
@@ -121,7 +121,7 @@ namespace GitHub.Runner.Listener
             }
             else
             {
-                _targetPackage = await _runnerServer.GetPackageAsync(_packageType, _platform, targetVersion, token);
+                _targetPackage = await _runnerServer.GetPackageAsync(_packageType, _platform, targetVersion, true, token);
                 if (_targetPackage == null)
                 {
                     Trace.Info($"There is no package for {_packageType} and {_platform} with version {targetVersion}.");
@@ -211,12 +211,22 @@ namespace GitHub.Runner.Listener
 
                             //open zip stream in async mode
                             using (HttpClient httpClient = new HttpClient(HostContext.CreateHttpClientHandler()))
-                            using (FileStream fs = new FileStream(archiveFile, FileMode.Create, FileAccess.Write, FileShare.None, bufferSize: 4096, useAsync: true))
-                            using (Stream result = await httpClient.GetStreamAsync(_targetPackage.DownloadUrl))
                             {
-                                //81920 is the default used by System.IO.Stream.CopyTo and is under the large object heap threshold (85k). 
+                                if (!string.IsNullOrEmpty(_targetPackage.Token))
-                                await result.CopyToAsync(fs, 81920, downloadCts.Token);
+                                {
-                                await fs.FlushAsync(downloadCts.Token);
+                                    Trace.Info($"Adding authorization token ({_targetPackage.Token.Length} chars)");
+                                    httpClient.DefaultRequestHeaders.Authorization = new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", _targetPackage.Token);
+                                }
+
+                                Trace.Info($"Downloading {_targetPackage.DownloadUrl}");
+
+                                using (FileStream fs = new FileStream(archiveFile, FileMode.Create, FileAccess.Write, FileShare.None, bufferSize: 4096, useAsync: true))
+                                using (Stream result = await httpClient.GetStreamAsync(_targetPackage.DownloadUrl))
+                                {
+                                    //81920 is the default used by System.IO.Stream.CopyTo and is under the large object heap threshold (85k). 
+                                    await result.CopyToAsync(fs, 81920, downloadCts.Token);
+                                    await fs.FlushAsync(downloadCts.Token);
+                                }
                             }
 
                             Trace.Info($"Download runner: finished download");

src/Sdk/DTGenerated/Generated/TaskAgentHttpClientBase.cs

@@ -587,6 +587,7 @@ namespace GitHub.DistributedTask.WebApi
         /// <param name="packageType"></param>
         /// <param name="platform"></param>
         /// <param name="version"></param>
+        /// <param name="includeToken"></param>
         /// <param name="userState"></param>
         /// <param name="cancellationToken">The cancellation token to cancel operation.</param>
         [EditorBrowsable(EditorBrowsableState.Never)]
@@ -594,6 +595,7 @@ namespace GitHub.DistributedTask.WebApi
             string packageType,
             string platform,
             string version,
+            bool? includeToken = null,
             object userState = null,
             CancellationToken cancellationToken = default)
         {
@@ -601,11 +603,18 @@ namespace GitHub.DistributedTask.WebApi
             Guid locationId = new Guid("8ffcd551-079c-493a-9c02-54346299d144");
             object routeValues = new { packageType = packageType, platform = platform, version = version };
 
+            List<KeyValuePair<string, string>> queryParams = new List<KeyValuePair<string, string>>();
+            if (includeToken != null)
+            {
+                queryParams.Add("includeToken", includeToken.Value.ToString());
+            }
+
             return SendAsync<PackageMetadata>(
                 httpMethod,
                 locationId,
                 routeValues: routeValues,
                 version: new ApiResourceVersion(5.1, 2),
+                queryParameters: queryParams,
                 userState: userState,
                 cancellationToken: cancellationToken);
         }
@@ -616,6 +625,7 @@ namespace GitHub.DistributedTask.WebApi
         /// <param name="packageType"></param>
         /// <param name="platform"></param>
         /// <param name="top"></param>
+        /// <param name="includeToken"></param>
         /// <param name="userState"></param>
         /// <param name="cancellationToken">The cancellation token to cancel operation.</param>
         [EditorBrowsable(EditorBrowsableState.Never)]
@@ -623,6 +633,7 @@ namespace GitHub.DistributedTask.WebApi
             string packageType,
             string platform = null,
             int? top = null,
+            bool? includeToken = null,
             object userState = null,
             CancellationToken cancellationToken = default)
         {
@@ -635,6 +646,10 @@ namespace GitHub.DistributedTask.WebApi
             {
                 queryParams.Add("$top", top.Value.ToString(CultureInfo.InvariantCulture));
             }
+            if (includeToken != null)
+            {
+                queryParams.Add("includeToken", includeToken.Value.ToString());
+            }
 
             return SendAsync<List<PackageMetadata>>(
                 httpMethod,

src/Sdk/DTWebApi/WebApi/PackageMetadata.cs

@@ -59,6 +59,16 @@ namespace GitHub.DistributedTask.WebApi
             set;
         }
 
+        /// <summary>
+        /// Auth token to download the package
+        /// </summary>
+        [DataMember]
+        public String Token
+        {
+            get;
+            set;
+        }
+
         /// <summary>
         /// MD5 hash as a base64 string
         /// </summary>