Skip to content

Commit

Permalink
support authenticated package download (#920)
Browse files Browse the repository at this point in the history
  • Loading branch information
ericsciple authored and TingluoHuang committed Apr 21, 2021
1 parent 4ec52d2 commit 29eae8d
Show file tree
Hide file tree
Showing 4 changed files with 48 additions and 13 deletions.
12 changes: 6 additions & 6 deletions src/Runner.Common/RunnerServer.cs
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,8 @@ public interface IRunnerServer : IRunnerService
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<PackageMetadata> GetPackageAsync(string packageType, string platform, string version, 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, bool includeToken, CancellationToken cancellationToken);

// agent update
Task<TaskAgent> UpdateAgentUpdateStateAsync(int agentPoolId, int agentId, string currentState);
Expand Down Expand Up @@ -317,16 +317,16 @@ public Task<TaskAgentMessage> GetAgentMessageAsync(Int32 poolId, Guid sessionId,
//-----------------------------------------------------------------
// 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)
Expand Down
24 changes: 17 additions & 7 deletions src/Runner.Listener/SelfUpdater.cs
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ private async Task<bool> UpdateNeeded(string targetVersion, CancellationToken to
// 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}.");
Expand All @@ -121,7 +121,7 @@ private async Task<bool> UpdateNeeded(string targetVersion, CancellationToken to
}
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}.");
Expand Down Expand Up @@ -211,12 +211,22 @@ private async Task DownloadLatestRunner(CancellationToken token)

//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).
await result.CopyToAsync(fs, 81920, downloadCts.Token);
await fs.FlushAsync(downloadCts.Token);
if (!string.IsNullOrEmpty(_targetPackage.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");
Expand Down
15 changes: 15 additions & 0 deletions src/Sdk/DTGenerated/Generated/TaskAgentHttpClientBase.cs
Original file line number Diff line number Diff line change
Expand Up @@ -587,25 +587,34 @@ public virtual async Task SendMessageAsync(
/// <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)]
public virtual Task<PackageMetadata> GetPackageAsync(
string packageType,
string platform,
string version,
bool? includeToken = null,
object userState = null,
CancellationToken cancellationToken = default)
{
HttpMethod httpMethod = new HttpMethod("GET");
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);
}
Expand All @@ -616,13 +625,15 @@ public virtual Task<PackageMetadata> GetPackageAsync(
/// <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)]
public virtual Task<List<PackageMetadata>> GetPackagesAsync(
string packageType,
string platform = null,
int? top = null,
bool? includeToken = null,
object userState = null,
CancellationToken cancellationToken = default)
{
Expand All @@ -635,6 +646,10 @@ public virtual Task<List<PackageMetadata>> GetPackagesAsync(
{
queryParams.Add("$top", top.Value.ToString(CultureInfo.InvariantCulture));
}
if (includeToken != null)
{
queryParams.Add("includeToken", includeToken.Value.ToString());
}

return SendAsync<List<PackageMetadata>>(
httpMethod,
Expand Down
10 changes: 10 additions & 0 deletions src/Sdk/DTWebApi/WebApi/PackageMetadata.cs
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,16 @@ public String DownloadUrl
set;
}

/// <summary>
/// Auth token to download the package
/// </summary>
[DataMember]
public String Token
{
get;
set;
}

/// <summary>
/// MD5 hash as a base64 string
/// </summary>
Expand Down

0 comments on commit 29eae8d

Please sign in to comment.