Prepend prefix for MAR operations (#1741)

PowerShell · Oct 30, 2024 · de278b4 · de278b4
1 parent 8a339b0
commit de278b4
Show file tree

Hide file tree

Showing 8 changed files with 111 additions and 7 deletions.
diff --git a/src/code/ContainerRegistryServerAPICalls.cs b/src/code/ContainerRegistryServerAPICalls.cs
@@ -286,7 +286,8 @@ public override Stream InstallPackage(string packageName, string packageVersion,
                 return results;
             }
 
-            results = InstallVersion(packageName, packageVersion, out errRecord);
+            string packageNameForInstall = PrependMARPrefix(packageName);
+            results = InstallVersion(packageNameForInstall, packageVersion, out errRecord);
             return results;
         }
 
@@ -1601,13 +1602,14 @@ private Hashtable[] FindPackagesWithVersionHelper(string packageName, VersionTyp
             string registryUrl = Repository.Uri.ToString();
             string packageNameLowercase = packageName.ToLower();
 
+            string packageNameForFind = PrependMARPrefix(packageNameLowercase);
             string containerRegistryAccessToken = GetContainerRegistryAccessToken(out errRecord);
             if (errRecord != null)
             {
                 return emptyHashResponses;
             }
 
-            var foundTags = FindContainerRegistryImageTags(packageNameLowercase, "*", containerRegistryAccessToken, out errRecord);
+            var foundTags = FindContainerRegistryImageTags(packageNameForFind, "*", containerRegistryAccessToken, out errRecord);
             if (errRecord != null || foundTags == null)
             {
                 return emptyHashResponses;
@@ -1616,7 +1618,7 @@ private Hashtable[] FindPackagesWithVersionHelper(string packageName, VersionTyp
             List<Hashtable> latestVersionResponse = new List<Hashtable>();
             List<JToken> allVersionsList = foundTags["tags"].ToList();
 
-            SortedDictionary<NuGet.Versioning.SemanticVersion, string> sortedQualifyingPkgs = GetPackagesWithRequiredVersion(allVersionsList, versionType, versionRange, requiredVersion, packageNameLowercase, includePrerelease, out errRecord);
+            SortedDictionary<NuGet.Versioning.SemanticVersion, string> sortedQualifyingPkgs = GetPackagesWithRequiredVersion(allVersionsList, versionType, versionRange, requiredVersion, packageNameForFind, includePrerelease, out errRecord);
             if (errRecord != null)
             {
                 return emptyHashResponses;
@@ -1627,7 +1629,7 @@ private Hashtable[] FindPackagesWithVersionHelper(string packageName, VersionTyp
             foreach (var pkgVersionTag in pkgsInDescendingOrder)
             {
                 string exactTagVersion = pkgVersionTag.Value.ToString();
-                Hashtable metadata = GetContainerRegistryMetadata(packageNameLowercase, exactTagVersion, containerRegistryAccessToken, out errRecord);
+                Hashtable metadata = GetContainerRegistryMetadata(packageNameForFind, exactTagVersion, containerRegistryAccessToken, out errRecord);
                 if (errRecord != null || metadata.Count == 0)
                 {
                     return emptyHashResponses;
@@ -1694,6 +1696,18 @@ private Hashtable[] FindPackagesWithVersionHelper(string packageName, VersionTyp
             return sortedPkgs;
         }
 
+        private string PrependMARPrefix(string packageName)
+        {
+            string prefix = string.IsNullOrEmpty(InternalHooks.MARPrefix) ? PSRepositoryInfo.MARPrefix : InternalHooks.MARPrefix;
+
+            // If the repostitory is MAR and its not a wildcard search, we need to prefix the package name with MAR prefix.
+            string updatedPackageName = Repository.IsMARRepository() && packageName.Trim() != "*"
+                                            ? string.Concat(prefix, packageName)
+                                            : packageName;
+
+            return updatedPackageName;
+        }
+
         #endregion
     }
 }
diff --git a/src/code/InternalHooks.cs b/src/code/InternalHooks.cs
@@ -15,6 +15,8 @@ public class InternalHooks
 
         internal static string AllowedUri;
 
+        internal static string MARPrefix;
+
         public static void SetTestHook(string property, object value)
         {
             var fieldInfo = typeof(InternalHooks).GetField(property, BindingFlags.Static | BindingFlags.NonPublic);

diff --git a/src/code/PSRepositoryInfo.cs b/src/code/PSRepositoryInfo.cs
@@ -11,6 +11,10 @@ namespace Microsoft.PowerShell.PSResourceGet.UtilClasses
     /// </summary>
     public sealed class PSRepositoryInfo
     {
+        #region constants
+        internal const string MARPrefix = "psresource/";
+        #endregion
+
         #region Enums
 
         public enum APIVersion
@@ -95,5 +99,14 @@ public enum RepositoryProviderType
         public bool IsAllowedByPolicy { get; set; }
 
         #endregion
+
+        #region Methods
+
+        internal bool IsMARRepository()
+        {
+            return (ApiVersion == APIVersion.ContainerRegistry && Uri.Host.Contains("mcr.microsoft.com"));
+        }
+
+        #endregion
     }
 }
diff --git a/src/code/PublishHelper.cs b/src/code/PublishHelper.cs
@@ -368,6 +368,17 @@ internal void PushResource(string Repository, string modulePrefix, bool SkipDepe
                     return;
                 }
 
+                if (repository.IsMARRepository())
+                {
+                    _cmdletPassedIn.WriteError(new ErrorRecord(
+                        new PSInvalidOperationException($"Repository '{repository.Name}' is a MAR repository and cannot be published to."),
+                        "MARRepositoryPublishError",
+                        ErrorCategory.PermissionDenied,
+                        this));
+
+                    return;
+                }
+
                 _networkCredential = Utils.SetNetworkCredential(repository, _networkCredential, _cmdletPassedIn);
 
                 // Check if dependencies already exist within the repo if:

diff --git a/src/code/RepositorySettings.cs b/src/code/RepositorySettings.cs
@@ -862,7 +862,7 @@ private static PSRepositoryInfo.APIVersion GetRepoAPIVersion(Uri repoUri)
                 // repositories with Uri Scheme "temp" may have PSPath Uri's like: "Temp:\repo" and we should consider them as local repositories.
                 return PSRepositoryInfo.APIVersion.Local;
             }
-            else if (repoUri.AbsoluteUri.EndsWith(".azurecr.io") || repoUri.AbsoluteUri.EndsWith(".azurecr.io/"))
+            else if (repoUri.AbsoluteUri.EndsWith(".azurecr.io") || repoUri.AbsoluteUri.EndsWith(".azurecr.io/") || repoUri.AbsoluteUri.Contains("mcr.microsoft.com"))
             {
                 return PSRepositoryInfo.APIVersion.ContainerRegistry;
             }
@@ -876,7 +876,7 @@ private static RepositoryProviderType GetRepositoryProviderType(Uri repoUri)
         {
             string absoluteUri = repoUri.AbsoluteUri;
             // We want to use contains instead of EndsWith to accomodate for trailing '/'
-            if (absoluteUri.Contains("azurecr.io")){
+            if (absoluteUri.Contains("azurecr.io") || absoluteUri.Contains("mcr.microsoft.com")){
                 return RepositoryProviderType.ACR;
             }
             // TODO: add a regex for this match

diff --git a/test/FindPSResourceTests/FindPSResourceContainerRegistryServer.Tests.ps1 b/test/FindPSResourceTests/FindPSResourceContainerRegistryServer.Tests.ps1
@@ -228,3 +228,21 @@ Describe 'Test HTTP Find-PSResource for ACR Server Protocol' -tags 'CI' {
         $res.Type.ToString() | Should -Be "Script"
     }
 }
+
+Describe 'Test Find-PSResource for MAR Repository' -tags 'CI' {
+    BeforeAll {
+        [Microsoft.PowerShell.PSResourceGet.UtilClasses.InternalHooks]::SetTestHook("MARPrefix", "azure-powershell/");
+        Register-PSResourceRepository -Name "MAR" -Uri "https://mcr.microsoft.com" -ApiVersion "ContainerRegistry"
+    }
+
+    AfterAll {
+        [Microsoft.PowerShell.PSResourceGet.UtilClasses.InternalHooks]::SetTestHook("MARPrefix", $null);
+        Unregister-PSResourceRepository -Name "MAR"
+    }
+
+    It "Should find resource given specific Name, Version null" {
+        $res = Find-PSResource -Name "Az.Accounts" -Repository "MAR"
+        $res.Name | Should -Be "Az.Accounts"
+        $res.Version | Should -Be "3.0.4"
+    }
+}
diff --git a/test/InstallPSResourceTests/InstallPSResourceContainerRegistryServer.Tests.ps1 b/test/InstallPSResourceTests/InstallPSResourceContainerRegistryServer.Tests.ps1
@@ -137,7 +137,7 @@ Describe 'Test Install-PSResource for ACR scenarios' -tags 'CI' {
 
     It "Install resource with a dependency (should install both parent and dependency)" {
         Install-PSResource -Name $testModuleParentName -Repository $ACRRepoName -TrustRepository
-        
+
         $parentPkg = Get-InstalledPSResource $testModuleParentName
         $parentPkg.Name | Should -Be $testModuleParentName
         $parentPkg.Version | Should -Be "1.0.0"
@@ -307,3 +307,28 @@ Describe 'Test Install-PSResource for V3Server scenarios' -tags 'ManualValidatio
         Set-PSResourceRepository PoshTestGallery -Trusted
     }
 }
+
+Describe 'Test Install-PSResource for MAR Repository' -tags 'CI' {
+    BeforeAll {
+        [Microsoft.PowerShell.PSResourceGet.UtilClasses.InternalHooks]::SetTestHook("MARPrefix", "azure-powershell/");
+        Register-PSResourceRepository -Name "MAR" -Uri "https://mcr.microsoft.com" -ApiVersion "ContainerRegistry"
+    }
+
+    AfterAll {
+        [Microsoft.PowerShell.PSResourceGet.UtilClasses.InternalHooks]::SetTestHook("MARPrefix", $null);
+        Unregister-PSResourceRepository -Name "MAR"
+    }
+
+    It "Should find resource given specific Name, Version null" {
+        try {
+            $pkg = Install-PSResource -Name "Az.Accounts" -Repository "MAR" -PassThru -TrustRepository -Reinstall
+            $pkg.Name | Should -Be "Az.Accounts"
+            $pkg.Version | Should -Be "3.0.4"
+        }
+        finally {
+            if ($pkg) {
+                Uninstall-PSResource -Name "Az.Accounts" -Version "3.0.4"
+            }
+        }
+    }
+}
diff --git a/test/PublishPSResourceTests/PublishPSResourceContainerRegistryServer.Tests.ps1 b/test/PublishPSResourceTests/PublishPSResourceContainerRegistryServer.Tests.ps1
@@ -512,3 +512,24 @@ Describe "Test Publish-PSResource" -tags 'CI' {
         $results[0].Version | Should -Be $version
     }
 }
+
+Describe 'Test Publish-PSResource for MAR Repository' -tags 'CI' {
+    BeforeAll {
+        [Microsoft.PowerShell.PSResourceGet.UtilClasses.InternalHooks]::SetTestHook("MARPrefix", "azure-powershell/");
+        Register-PSResourceRepository -Name "MAR" -Uri "https://mcr.microsoft.com" -ApiVersion "ContainerRegistry"
+    }
+
+    AfterAll {
+        [Microsoft.PowerShell.PSResourceGet.UtilClasses.InternalHooks]::SetTestHook("MARPrefix", $null);
+        Unregister-PSResourceRepository -Name "MAR"
+    }
+
+    It "Should find resource given specific Name, Version null" {
+        $fileName = "NonExistent.psd1"
+        $modulePath = New-Item -Path "$TestDrive\NonExistent" -ItemType Directory -Force
+        $psd1Path = Join-Path -Path $modulePath -ChildPath $fileName
+        New-ModuleManifest -Path $psd1Path -ModuleVersion "1.0.0" -Description "NonExistent module"
+
+        { Publish-PSResource -Path $modulePath -Repository "MAR" -ErrorAction Stop } | Should -Throw -ErrorId "MARRepositoryPublishError,Microsoft.PowerShell.PSResourceGet.Cmdlets.PublishPSResource"
+    }
+}