112 lines
4.8 KiB
C#
112 lines
4.8 KiB
C#
using System.Net.Http.Headers;
|
|
using System.Security.Cryptography;
|
|
using System.Text.Json;
|
|
|
|
namespace Zcbot.WindowsNode;
|
|
|
|
internal sealed class JobInputDownloader(NodeConfig config, JobInboxStore inbox)
|
|
{
|
|
internal async Task DownloadAsync(RecoverableJob job, CancellationToken cancellationToken)
|
|
{
|
|
var transfers = inbox.InputTransfers(job);
|
|
for (var index = 0; index < transfers.Count; index++)
|
|
{
|
|
inbox.WriteState(
|
|
job,
|
|
"downloading_inputs",
|
|
Math.Max(0, index * 5 / transfers.Count),
|
|
$"正在下载输入 {index + 1}/{transfers.Count}");
|
|
await DownloadOneAsync(job, transfers[index], cancellationToken);
|
|
}
|
|
}
|
|
|
|
private async Task DownloadOneAsync(
|
|
RecoverableJob job,
|
|
JsonElement transfer,
|
|
CancellationToken cancellationToken)
|
|
{
|
|
var downloadPath = transfer.GetProperty("download_path").GetString()!;
|
|
if (!downloadPath.StartsWith("/v1/software-jobs/", StringComparison.Ordinal)
|
|
|| !Uri.TryCreate(downloadPath, UriKind.Relative, out var relativeUri))
|
|
{
|
|
throw new InvalidDataException("Job input download path is invalid.");
|
|
}
|
|
var expectedSize = transfer.GetProperty("size_bytes").GetInt64();
|
|
var expectedSha256 = transfer.GetProperty("sha256").GetString()!;
|
|
var destination = inbox.InputPath(job, transfer);
|
|
if (File.Exists(destination))
|
|
{
|
|
await VerifyExistingAsync(destination, expectedSize, expectedSha256, cancellationToken);
|
|
return;
|
|
}
|
|
|
|
Directory.CreateDirectory(Path.GetDirectoryName(destination)!);
|
|
var temporaryPath = destination + ".tmp-" + Guid.NewGuid().ToString("N");
|
|
using var client = new HttpClient { BaseAddress = config.ServerUrl };
|
|
client.DefaultRequestHeaders.Authorization =
|
|
new AuthenticationHeaderValue("Bearer", config.NodeToken);
|
|
client.DefaultRequestHeaders.Add("X-Node-Id", config.NodeId.ToString());
|
|
try
|
|
{
|
|
using var response = await client.GetAsync(
|
|
relativeUri, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
|
|
response.EnsureSuccessStatusCode();
|
|
if (response.Content.Headers.ContentLength is long contentLength
|
|
&& contentLength != expectedSize)
|
|
{
|
|
throw new InvalidDataException("Job input size header does not match the manifest.");
|
|
}
|
|
await using var source = await response.Content.ReadAsStreamAsync(cancellationToken);
|
|
await using var target = new FileStream(
|
|
temporaryPath, FileMode.CreateNew, FileAccess.Write, FileShare.None,
|
|
bufferSize: 64 * 1024, FileOptions.Asynchronous | FileOptions.WriteThrough);
|
|
using var hash = IncrementalHash.CreateHash(HashAlgorithmName.SHA256);
|
|
var buffer = new byte[64 * 1024];
|
|
long total = 0;
|
|
while (true)
|
|
{
|
|
var count = await source.ReadAsync(buffer, cancellationToken);
|
|
if (count == 0) break;
|
|
total += count;
|
|
if (total > expectedSize)
|
|
{
|
|
throw new InvalidDataException("Job input exceeded its declared size.");
|
|
}
|
|
hash.AppendData(buffer, 0, count);
|
|
await target.WriteAsync(buffer.AsMemory(0, count), cancellationToken);
|
|
}
|
|
await target.FlushAsync(cancellationToken);
|
|
target.Flush(flushToDisk: true);
|
|
var actualSha256 = Convert.ToHexString(hash.GetHashAndReset()).ToLowerInvariant();
|
|
if (total != expectedSize || actualSha256 != expectedSha256)
|
|
{
|
|
throw new InvalidDataException("Job input digest does not match the manifest.");
|
|
}
|
|
target.Close();
|
|
File.Move(temporaryPath, destination, overwrite: false);
|
|
}
|
|
finally
|
|
{
|
|
if (File.Exists(temporaryPath)) File.Delete(temporaryPath);
|
|
}
|
|
}
|
|
|
|
private static async Task VerifyExistingAsync(
|
|
string path, long expectedSize, string expectedSha256, CancellationToken cancellationToken)
|
|
{
|
|
var info = new FileInfo(path);
|
|
if (info.Length != expectedSize)
|
|
{
|
|
throw new InvalidDataException("Existing job input size does not match the manifest.");
|
|
}
|
|
await using var stream = new FileStream(
|
|
path, FileMode.Open, FileAccess.Read, FileShare.Read, 64 * 1024, FileOptions.Asynchronous);
|
|
var digest = Convert.ToHexString(
|
|
await SHA256.HashDataAsync(stream, cancellationToken)).ToLowerInvariant();
|
|
if (digest != expectedSha256)
|
|
{
|
|
throw new InvalidDataException("Existing job input digest does not match the manifest.");
|
|
}
|
|
}
|
|
}
|