Support override of powershell module feeds

eng/common/scripts/Helpers/PSModule-Helpers.ps1

@@ -1,4 +1,3 @@
-$DefaultPSRepositoryUrl = "https://www.powershellgallery.com/api/v2"
 $global:CurrentUserModulePath = ""
 
 function Update-PSModulePathForCI()
@@ -47,6 +46,65 @@ function Update-PSModulePathForCI()
   }
 }
 
+function Get-ModuleRepositories([string]$moduleName) {
+  $DefaultPSRepositoryUrl = "https://www.powershellgallery.com/api/v2"
+  # List of modules+versions we want to replace with internal feed sources for reliability, security, etc.
+  $packageFeedOverrides = @{
+    'powershell-yaml' = 'https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-tools/nuget/v2'
+  }
+
+  $repoUrls = if ($packageFeedOverrides.Contains("${moduleName}")) {
+    @($packageFeedOverrides["${moduleName}"], $DefaultPSRepositoryUrl)
+  } else {
+    @($DefaultPSRepositoryUrl)
+  }
+
+  return $repoUrls
+}
+
+function moduleIsInstalled([string]$moduleName, [string]$version) {
+  $modules = (Get-Module -ListAvailable $moduleName)
+  if ($version -as [Version]) {
+    $modules = $modules.Where({ [Version]$_.Version -ge [Version]$version })
+    if ($modules.Count -gt 0)
+    {
+      Write-Host "Using module $($modules[0].Name) with version $($modules[0].Version)."
+      return $modules[0]
+    }
+  }
+  return $null
+}
+
+function installModule([string]$moduleName, [string]$version, $repoUrl) {
+  $repo = (Get-PSRepository).Where({ $_.SourceLocation -eq $repoUrl })
+  if ($repo.Count -eq 0)
+  {
+    Register-PSRepository -Name $repoUrl -SourceLocation $repoUrl -InstallationPolicy Trusted
+    $repo = (Get-PSRepository).Where({ $_.SourceLocation -eq $repoUrl })
+    if ($repo.Count -eq 0) {
+      throw "Failed to register package repository $repoUrl."
+    }
+  }
+
+  if ($repo.InstallationPolicy -ne "Trusted") {
+    Set-PSRepository -Name $repo.Name -InstallationPolicy "Trusted"
+  }
+
+  Write-Host "Installing module $moduleName with min version $version from $repoUrl"
+  # Install under CurrentUser scope so that the end up under $CurrentUserModulePath for caching
+  Install-Module $moduleName -MinimumVersion $version -Repository $repo.Name -Scope CurrentUser -Force
+  # Ensure module installed
+  $modules = (Get-Module -ListAvailable $moduleName)
+  if ($version -as [Version]) {
+    $modules = $modules.Where({ [Version]$_.Version -ge [Version]$version })
+  }
+  if ($modules.Count -eq 0) {
+    throw "Failed to install module $moduleName with version $version"
+  }
+
+  return $modules[0]
+}
+
 # Manual test at eng/common-tests/psmodule-helpers/Install-Module-Parallel.ps1
 # If we want to use another default repository other then PSGallery we can update the default parameters
 function Install-ModuleIfNotInstalled()
@@ -55,69 +113,45 @@ function Install-ModuleIfNotInstalled()
   param(
     [string]$moduleName,
     [string]$version,
-    [string]$repositoryUrl = $DefaultPSRepositoryUrl
+    [string]$repositoryUrl
   )
 
-  # Check installed modules
-  $modules = (Get-Module -ListAvailable $moduleName)
-  if ($version -as [Version]) {
-    $modules = $modules.Where({ [Version]$_.Version -ge [Version]$version })
-  }
+  # Check installed modules before after acquiring lock to avoid a big queue
+  $module = moduleIsInstalled -moduleName $moduleName -version $version
+  if ($module) { return $module }
 
-  if ($modules.Count -eq 0)
-  {
-    # Use double-checked locking to avoid locking when module is already installed
+  try {
     $mutex = New-Object System.Threading.Mutex($false, "Install-ModuleIfNotInstalled")
     $null = $mutex.WaitOne()
 
-    try {
-      # Check installed modules again after acquiring lock
-      $modules = (Get-Module -ListAvailable $moduleName)
-      if ($version -as [Version]) {
-        $modules = $modules.Where({ [Version]$_.Version -ge [Version]$version })
-      }
-
-      if ($modules.Count -eq 0)
-      {
-        $repositories = (Get-PSRepository).Where({ $_.SourceLocation -eq $repositoryUrl })
-        if ($repositories.Count -eq 0)
-        {
-          Register-PSRepository -Name $repositoryUrl -SourceLocation $repositoryUrl -InstallationPolicy Trusted
-          $repositories = (Get-PSRepository).Where({ $_.SourceLocation -eq $repositoryUrl })
-          if ($repositories.Count -eq 0) {
-            Write-Error "Failed to register package repository $repositoryUrl."
-            return
-          }
-        }
-        $repository = $repositories[0]
-
-        if ($repository.InstallationPolicy -ne "Trusted") {
-          Set-PSRepository -Name $repository.Name -InstallationPolicy "Trusted"
-        }
-
-        Write-Host "Installing module $moduleName with min version $version from $repositoryUrl"
-        # Install under CurrentUser scope so that the end up under $CurrentUserModulePath for caching
-        Install-Module $moduleName -MinimumVersion $version -Repository $repository.Name -Scope CurrentUser -Force
-
-        # Ensure module installed
-        $modules = (Get-Module -ListAvailable $moduleName)
-        if ($version -as [Version]) {
-          $modules = $modules.Where({ [Version]$_.Version -ge [Version]$version })
-        }
-
-        if ($modules.Count -eq 0) {
-          Write-Error "Failed to install module $moduleName with version $version"
-          return
+    # Check installed modules again after acquiring lock, in case it has been installed
+    $module = moduleIsInstalled -moduleName $moduleName -version $version
+    if ($module) { return $module }
+
+    $repoUrls = Get-ModuleRepositories $moduleName
+
+    foreach ($url in $repoUrls) {
+      try {
+        $module = installModule -moduleName $moduleName -version $version -repoUrl $url
+      } catch {
+        if ($url -ne $repoUrls[-1]) {
+          Write-Warning "Failed to install powershell module from '$url'. Retrying with fallback repository"
+          Write-Warning $_
+          continue
+        } else {
+          Write-Warning "Failed to install powershell module from $url"
+          throw
         }
       }
+      break
     }
-    finally {
-      $mutex.ReleaseMutex()
-    }
+
+    Write-Host "Using module '$($module.Name)' with version '$($module.Version)'."
+  } finally {
+    $mutex.ReleaseMutex()
   }
 
-  Write-Host "Using module $($modules[0].Name) with version $($modules[0].Version)."
-  return $modules[0]
+  return $module
 }
 
 if ($null -ne $env:SYSTEM_TEAMPROJECTID) {