Skip to content

Commit a4730af

Browse files
authored
Fix thread-safety in ContainerRegistryServerAPICalls.InstallPackageAsync
InstallPackageAsync was calling the synchronous InstallPackage() which writes to _cmdletPassedIn.WriteDebug, causing cross-thread cmdlet stream writes when used from Parallel.ForEach. - Refactor InstallPackageAsync to not call InstallPackage(); inline the null-version check and enqueue all messages via the provided queues - Add a queue-aware InstallVersion overload that uses ConcurrentQueue parameters instead of _cmdletPassedIn.Write* calls, for use by the async install path - Keep the original InstallVersion(out ErrorRecord) overload intact for the synchronous path
1 parent fbb940f commit a4730af

1 file changed

Lines changed: 81 additions & 3 deletions

File tree

src/code/ContainerRegistryServerAPICalls.cs

Lines changed: 81 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -368,12 +368,20 @@ public override Stream InstallPackage(string packageName, string packageVersion,
368368
public override Task<Stream> InstallPackageAsync(string packageName, string packageVersion, bool includePrerelease, ConcurrentQueue<ErrorRecord> errorMsgs, ConcurrentQueue<string> warningMsgs, ConcurrentQueue<string> debugMsgs, ConcurrentQueue<string> verboseMsgs)
369369
{
370370
debugMsgs.Enqueue("In ContainerRegistryServerAPICalls::InstallPackageAsync()");
371-
Stream results = InstallPackage(packageName, packageVersion, includePrerelease, out ErrorRecord errRecord);
372-
if (errRecord != null)
371+
Stream results = new MemoryStream();
372+
if (string.IsNullOrEmpty(packageVersion))
373373
{
374-
errorMsgs.Enqueue(errRecord);
374+
errorMsgs.Enqueue(new ErrorRecord(
375+
exception: new ArgumentNullException($"Package version could not be found for {packageName}"),
376+
"PackageVersionNullOrEmptyError",
377+
ErrorCategory.InvalidArgument,
378+
_cmdletPassedIn));
379+
380+
return Task.FromResult(results);
375381
}
376382

383+
string packageNameForInstall = PrependMARPrefix(packageName);
384+
results = InstallVersion(packageNameForInstall, packageVersion, errorMsgs, debugMsgs, verboseMsgs);
377385
return Task.FromResult(results);
378386
}
379387

@@ -445,6 +453,76 @@ private Stream InstallVersion(
445453
return responseContent.ReadAsStreamAsync().Result;
446454
}
447455

456+
/// <summary>
457+
/// Installs a package with version specified using concurrent queues for output instead of cmdlet streams.
458+
/// Used by the async install path to avoid cross-thread cmdlet stream writes.
459+
/// </summary>
460+
private Stream InstallVersion(
461+
string packageName,
462+
string packageVersion,
463+
ConcurrentQueue<ErrorRecord> errorMsgs,
464+
ConcurrentQueue<string> debugMsgs,
465+
ConcurrentQueue<string> verboseMsgs)
466+
{
467+
debugMsgs.Enqueue("In ContainerRegistryServerAPICalls::InstallVersion()");
468+
string packageNameLowercase = packageName.ToLower();
469+
string tempPath = Path.Combine(Path.GetTempPath(), Guid.NewGuid().ToString());
470+
try
471+
{
472+
Directory.CreateDirectory(tempPath);
473+
}
474+
catch (Exception e)
475+
{
476+
errorMsgs.Enqueue(new ErrorRecord(
477+
exception: e,
478+
"InstallVersionTempDirCreationError",
479+
ErrorCategory.InvalidResult,
480+
_cmdletPassedIn));
481+
482+
return null;
483+
}
484+
485+
string containerRegistryAccessToken = GetContainerRegistryAccessToken(needCatalogAccess: false, isPushOperation: false, out ErrorRecord errRecord);
486+
if (errRecord != null)
487+
{
488+
errorMsgs.Enqueue(errRecord);
489+
return null;
490+
}
491+
492+
verboseMsgs.Enqueue($"Getting manifest for {packageNameLowercase} - {packageVersion}");
493+
var manifest = GetContainerRegistryRepositoryManifest(packageNameLowercase, packageVersion, containerRegistryAccessToken, out errRecord);
494+
if (errRecord != null)
495+
{
496+
errorMsgs.Enqueue(errRecord);
497+
return null;
498+
}
499+
string digest = GetDigestFromManifest(manifest, out errRecord);
500+
if (errRecord != null)
501+
{
502+
errorMsgs.Enqueue(errRecord);
503+
return null;
504+
}
505+
506+
verboseMsgs.Enqueue($"Downloading blob for {packageNameLowercase} - {packageVersion}");
507+
HttpContent responseContent;
508+
try
509+
{
510+
responseContent = GetContainerRegistryBlobAsync(packageNameLowercase, digest, containerRegistryAccessToken).Result;
511+
}
512+
catch (Exception e)
513+
{
514+
errorMsgs.Enqueue(new ErrorRecord(
515+
exception: e,
516+
"InstallVersionGetContainerRegistryBlobAsyncError",
517+
ErrorCategory.InvalidResult,
518+
_cmdletPassedIn));
519+
520+
return null;
521+
}
522+
523+
return responseContent.ReadAsStreamAsync().Result;
524+
}
525+
448526
#endregion
449527

450528
#region Authentication and Token Methods

0 commit comments

Comments
 (0)