diff --git a/Wabbajack.BuildServer/Controllers/UploadedFiles.cs b/Wabbajack.BuildServer/Controllers/UploadedFiles.cs index 62735b77..1fc16684 100644 --- a/Wabbajack.BuildServer/Controllers/UploadedFiles.cs +++ b/Wabbajack.BuildServer/Controllers/UploadedFiles.cs @@ -1,4 +1,5 @@ using System; +using System.Collections.Concurrent; using System.Collections.Generic; using System.IO; using System.Linq; @@ -26,6 +27,7 @@ namespace Wabbajack.BuildServer.Controllers { public class UploadedFiles : AControllerBase { + private static ConcurrentDictionary _writeLocks = new ConcurrentDictionary(); private AppSettings _settings; public UploadedFiles(ILogger logger, DBContext db, AppSettings settings) : base(logger, db) @@ -39,6 +41,9 @@ namespace Wabbajack.BuildServer.Controllers { var guid = Guid.NewGuid(); var key = Encoding.UTF8.GetBytes($"{Path.GetFileNameWithoutExtension(Name)}|{guid.ToString()}|{Path.GetExtension(Name)}").ToHex(); + + _writeLocks.GetOrAdd(key, new AsyncLock()); + System.IO.File.Create(Path.Combine("public", "files", key)).Close(); Utils.Log($"Starting Ingest for {key}"); return Ok(key); @@ -52,18 +57,27 @@ namespace Wabbajack.BuildServer.Controllers if (!Key.All(a => HexChars.Contains(a))) return BadRequest("NOT A VALID FILENAME"); Utils.Log($"Writing at position {Offset} in ingest file {Key}"); + + var ms = new MemoryStream(); + await Request.Body.CopyToAsync(ms); + ms.Position = 0; + + long position; + using (var _ = await _writeLocks[Key].Wait()) await using (var file = System.IO.File.Open(Path.Combine("public", "files", Key), FileMode.Open, FileAccess.Write, FileShare.ReadWrite)) { file.Position = Offset; - await Request.Body.CopyToAsync(file); - return Ok(file.Position.ToString()); + await ms.CopyToAsync(file); + position = file.Position; } + return Ok(position); } [HttpPut] - [Route("upload_file/{Key}/finish")] - public async Task UploadFileFinish(string Key) + [Route("upload_file/{Key}/finish/{xxHashAsHex}")] + public async Task UploadFileFinish(string Key, string xxHashAsHex) { + var expectedHash = xxHashAsHex.FromHex().ToBase64(); var user = User.FindFirstValue(ClaimTypes.Name); if (!Key.All(a => HexChars.Contains(a))) return BadRequest("NOT A VALID FILENAME"); @@ -74,8 +88,14 @@ namespace Wabbajack.BuildServer.Controllers var final_path = Path.Combine("public", "files", final_name); System.IO.File.Move(Path.Combine("public", "files", Key), final_path); var hash = await final_path.FileHashAsync(); - + if (expectedHash != hash) + { + System.IO.File.Delete(final_path); + return BadRequest($"Bad Hash, Expected: {expectedHash} Got: {hash}"); + } + + _writeLocks.TryRemove(Key, out var _); var record = new UploadedFile { Id = parts[1], diff --git a/Wabbajack.Lib/FileUploader/AuthorAPI.cs b/Wabbajack.Lib/FileUploader/AuthorAPI.cs index be56db50..553b3049 100644 --- a/Wabbajack.Lib/FileUploader/AuthorAPI.cs +++ b/Wabbajack.Lib/FileUploader/AuthorAPI.cs @@ -1,4 +1,6 @@ using System; +using System.Collections; +using System.Collections.Generic; using System.IO; using System.Linq; using System.Net; @@ -38,6 +40,9 @@ namespace Wabbajack.Lib.FileUploader var handler = new HttpClientHandler {MaxConnectionsPerServer = MAX_CONNECTIONS}; var client = new HttpClient(handler); var fsize = new FileInfo(filename).Length; + + var hash_task = filename.FileHashAsync(); + client.DefaultRequestHeaders.Add("X-API-KEY", AuthorAPI.GetAPIKey()); var response = await client.PutAsync(UploadURL+$"/{Path.GetFileName(filename)}/start", new StringContent("")); if (!response.IsSuccessStatusCode) @@ -46,12 +51,18 @@ namespace Wabbajack.Lib.FileUploader return; } + IEnumerable Blocks(long fsize) + { + for (long block = 0; block * BLOCK_SIZE < fsize; block ++) + yield return block; + } + var key = await response.Content.ReadAsStringAsync(); long sent = 0; using (var iqueue = new WorkQueue(MAX_CONNECTIONS)) { iqueue.Report("Starting Upload", 1); - await Enumerable.Range(0, (int)(fsize / BLOCK_SIZE)) + await Blocks(fsize) .PMap(iqueue, async block_idx => { if (tcs.Task.IsFaulted) return; @@ -93,7 +104,8 @@ namespace Wabbajack.Lib.FileUploader if (!tcs.Task.IsFaulted) { progressFn(1.0); - response = await client.PutAsync(UploadURL + $"/{key}/finish", new StringContent("")); + var hash = (await hash_task).FromBase64().ToHex(); + response = await client.PutAsync(UploadURL + $"/{key}/finish/{hash}", new StringContent("")); if (response.IsSuccessStatusCode) tcs.SetResult(await response.Content.ReadAsStringAsync()); else