Add support for AAD auth, enabled by default (#19)

* Base work for aad auth * Add read only tests * Add AadToken tests * Typo
cmbrose · Aug 16, 2024 · 4ecd157 · 4ecd157
1 parent 2e153a9
commit 4ecd157
Show file tree

Hide file tree

Showing 12 changed files with 247 additions and 63 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,10 @@
 # Release Notes
 All notable changes and release history of the "cosmos-db" module will be documented in this file.
 
+## 1.19
+* Changed to Entra Id authentication by default. 
+* Adds a `-enableMasterKeyAuth` flag to `Use-CosmosDbInternalFlag` which reverts auth to use master keys.
+
 ## 1.18
 * Fixes a bug in commands like `Search-CosmosDbRecords` and `Get-AllCosmosDbRecords` which might run for long enough that their auth tokens expire and aren't refreshed. Auth tokens will now be refreshed every 10 min as these commands run.
 * Adds a `-enableAuthHeaderReuse` flag to `Use-CosmosDbInternalFlag` which disables the 10 minute refresh period and forces auth header refreshes for every API call.

diff --git a/README.md b/README.md
@@ -266,6 +266,7 @@ Use-CosmosDbInternalFlag -EnableFiddlerDebugging $true
 | EnableCaching | Enables/disables caching certain values like DB keys, partition ranges, etc. Improves performance of nearly all operations. | No - default is enabled |
 | EnablePartitionKeyRangeSearches | **[Experimental]** Enables/disables filtering `Search` queries to relevant partition ranges instead of a full scan. Improves performance of `Search` commands. | No - default is disabled |
 | EnableAuthHeaderReuse | Enables/disables reusing auth headers for commands which use continuation tokens, like `Search-CosmosDbRecords` or `Get-AllCosmosDbRecords`. | No - default is enabled |
+| EnableMasterKeyAuth | Enables/disables using master key based authentication instead of Entra Id. | No - default is disabled |
 
 ## Error Handling
 

diff --git a/cosmos-db/cosmos-db.psd1 b/cosmos-db/cosmos-db.psd1
@@ -11,7 +11,7 @@
     # RootModule = ''
 
     # Version number of this module.
-    ModuleVersion     = '1.18'
+    ModuleVersion     = '1.19'
 
     # Supported PSEditions
     # CompatiblePSEditions = @()

diff --git a/cosmos-db/cosmos-db.psm1 b/cosmos-db/cosmos-db.psm1
@@ -15,6 +15,7 @@ $API_VERSION = "2018-12-31"
 $AUTHORIZATION_HEADER_REFRESH_THRESHOLD = [System.TimeSpan]::FromMinutes(10);
 
 $MASTER_KEY_CACHE = @{}
+$AAD_TOKEN_CACHE = @{}
 $SIGNATURE_HASH_CACHE = @{}
 $PARTITION_KEY_RANGE_CACHE = @{}
 
@@ -57,13 +58,23 @@ Function Get-CacheValue([string]$key, [hashtable]$cache) {
     return $null
 }
 
-Function Set-CacheValue([string]$key, $value, [hashtable]$cache, [int]$expirationHours) {
+Function Set-CacheValue([string]$key, $value, [hashtable]$cache, [int]$expirationHours, [int]$expirationMinutes) {
     $cache[$key] = @{ 
-        Expiration = [datetime]::UtcNow.AddHours($expirationHours);
+        Expiration = [datetime]::UtcNow.AddHours($expirationHours).AddMinutes($expirationMinutes);
         Value      = $value 
     }
 }
 
+Function RequireWriteableKey() {
+    # Entra Id auth does not support read only tokens (because it's just handled by RBAC).
+    # So if the user requested to be in read only mode, we have to manage it ourselves.
+
+    $readonlyEnabled = $env:COSMOS_DB_FLAG_ENABLE_READONLY_KEYS -eq 1
+    if ($readonlyEnabled) {
+        throw "Operation not allowed in readonly mode"
+    }
+}
+
 Function Get-Base64Masterkey([string]$ResourceGroup, [string]$Database, [string]$SubscriptionId) {
     $readonly = $env:COSMOS_DB_FLAG_ENABLE_READONLY_KEYS -eq 1
 
@@ -99,6 +110,31 @@ Function Get-Base64MasterkeyWithoutCaching([string]$ResourceGroup, [string]$Data
     $masterKey
 }
 
+Function Get-AADToken([string]$ResourceGroup, [string]$Database, [string]$SubscriptionId) {
+    $cacheKey = "aadtoken"
+    $cacheResult = Get-CacheValue -Key $cacheKey -Cache $AAD_TOKEN_CACHE
+    if ($cacheResult) {
+        return $cacheResult
+    }
+
+    $oauth = Get-AADTokenWithoutCaching
+    $token = $oauth.AccessToken
+
+    $expirationUtcTimestamp = $oauth.expires_on
+    $expiration = [System.DateTimeOffset]::FromUnixTimeSeconds($expirationUtcTimestamp).DateTime
+    $remainingTime = $expiration - [System.DateTime]::UtcNow
+    $cacheExpirationMinutes = ($remainingTime - $AUTHORIZATION_HEADER_REFRESH_THRESHOLD).Minutes
+
+    Set-CacheValue -Key $cacheKey -Value $token -Cache $AAD_TOKEN_CACHE -ExpirationMinutes $cacheExpirationMinutes
+
+    $token
+}
+
+# This is just to support testing caching with Get-AADToken and isn't meant to be used directly
+Function Get-AADTokenWithoutCaching() {
+    az account get-access-token --resource "https://cosmos.azure.com" | ConvertFrom-Json
+}
+
 Function Get-Signature([string]$verb, [string]$resourceType, [string]$resourceUrl, [string]$now) {
     $parts = @(
         $verb.ToLower(),
@@ -128,19 +164,25 @@ Function Get-Base64EncryptedSignatureHash([string]$masterKey, [string]$signature
     $base64Hash
 }
 
-Function Get-EncodedAuthString([string]$signatureHash) {
-    $authString = "type=master&ver=1.0&sig=$signatureHash"
+Function Get-EncodedAuthString([string]$signatureHash, [string]$type) {
+    $authString = "type=$type&ver=1.0&sig=$signatureHash"
     [uri]::EscapeDataString($authString)
 }
 
 Function Get-AuthorizationHeader([string]$ResourceGroup, [string]$SubscriptionId, [string]$Database, [string]$verb, [string]$resourceType, [string]$resourceUrl, [string]$now) {            
-    $masterKey = Get-Base64Masterkey -ResourceGroup $ResourceGroup -Database $Database -SubscriptionId $SubscriptionId
+    if ($env:COSMOS_DB_FLAG_ENABLE_MASTER_KEY_AUTH -eq 1) {
+        $masterKey = Get-Base64Masterkey -ResourceGroup $ResourceGroup -Database $Database -SubscriptionId $SubscriptionId
 
-    $signature = Get-Signature -verb $verb -resourceType $resourceType -resourceUrl $resourceUrl -now $now
+        $signature = Get-Signature -verb $verb -resourceType $resourceType -resourceUrl $resourceUrl -now $now
+
+        $signatureHash = Get-Base64EncryptedSignatureHash -masterKey $masterKey -signature $signature
+
+        return Get-EncodedAuthString -signatureHash $signatureHash -type "master"
+    }
 
-    $signatureHash = Get-Base64EncryptedSignatureHash -masterKey $masterKey -signature $signature
+    $token = Get-AADToken
 
-    Get-EncodedAuthString -signatureHash $signatureHash
+    return Get-EncodedAuthString -signatureHash $token -type "aad"
 }
 
 Function Get-CommonHeaders([string]$now, [string]$encodedAuthString, [string]$contentType = "application/json", [bool]$isQuery = $false, [string]$PartitionKey = $null, [string]$Etag = $null) {
@@ -268,7 +310,7 @@ Function Invoke-CosmosDbApiRequestWithContinuation([string]$verb, [string]$url,
         $headers["x-ms-continuation"] = $continuationToken
 
         $authHeaderReuseDisabled = $env:COSMOS_DB_FLAG_ENABLE_AUTH_HEADER_REUSE -eq 0
-        $authHeaderExpired = [System.DateTime]::Parse($authHeaders.now) + $AUTHORIZATION_HEADER_REFRESH_THRESHOLD -lt [System.DateTime]::UtcNow
+        $authHeaderExpired = [System.DateTimeOffset]::Parse($authHeaders.now) + $AUTHORIZATION_HEADER_REFRESH_THRESHOLD -lt [System.DateTime]::UtcNow
 
         if ($authHeaderReuseDisabled -or $authHeaderExpired) {
             $authHeaders = Invoke-Command -ScriptBlock $refreshAuthHeaders
@@ -746,6 +788,8 @@ Function New-CosmosDbRecord {
     )
 
     begin {
+        RequireWriteableKey
+
         $baseUrl = Get-BaseDatabaseUrl $Database
         $collectionsUrl = Get-CollectionsUrl $Container $Collection
         $docsUrl = "$collectionsUrl/$DOCS_TYPE"
@@ -833,6 +877,8 @@ Function Update-CosmosDbRecord {
     )
 
     begin {
+        RequireWriteableKey
+
         $baseUrl = Get-BaseDatabaseUrl $Database
     }
     process {
@@ -929,6 +975,8 @@ Function Remove-CosmosDbRecord {
     )
 
     begin {
+        RequireWriteableKey
+
         $baseUrl = Get-BaseDatabaseUrl $Database        
     }
     process {
@@ -1005,7 +1053,8 @@ Function Use-CosmosDbInternalFlag
     $enableFiddlerDebugging = $null,
     $enableCaching = $null,
     $enablePartitionKeyRangeSearches = $null,
-    $enableAuthHeaderReuse = $null
+    $enableAuthHeaderReuse = $null,
+    $enableMasterKeyAuth = $null
 ) {
     if ($null -ne $enableFiddlerDebugging) {
         $env:AZURE_CLI_DISABLE_CONNECTION_VERIFICATION = if ($enableFiddlerDebugging) { 1 } else { 0 }
@@ -1022,6 +1071,10 @@ Function Use-CosmosDbInternalFlag
     if ($null -ne $enableAuthHeaderReuse) {
         $env:COSMOS_DB_FLAG_ENABLE_AUTH_HEADER_REUSE = if ($enableAuthHeaderReuse) { 1 } else { 0 }
     }
+
+    if ($null -ne $enableMasterKeyAuth) {
+        $env:COSMOS_DB_FLAG_ENABLE_MASTER_KEY_AUTH = if ($enableMasterKeyAuth) { 1 } else { 0 }
+    }
 }
 
 Function Use-CosmosDbReadonlyKeys

diff --git a/tests/Get-AadToken.Tests.ps1 b/tests/Get-AadToken.Tests.ps1
@@ -0,0 +1,67 @@
+Get-Module cosmos-db | Remove-Module -Force
+Import-Module $PSScriptRoot\..\cosmos-db\cosmos-db.psm1 -Force
+
+InModuleScope cosmos-db {
+    Describe "Get-AadToken" {                    
+        BeforeAll {
+            $MOCK_TOKEN = @{
+                accessToken = "MOCK_TOKEN"
+                expires_on  = [System.DateTimeOffset]::UtcNow.AddHours(1).ToUnixTimeSeconds()
+            }
+
+            Mock Get-AadTokenWithoutCaching {
+                return $MOCK_TOKEN
+            }
+        }
+
+        BeforeEach {
+            # This is defined in the main module
+            $AAD_TOKEN_CACHE = @{}
+        }
+
+        It "Only calls the core logic once with caching enabled" {
+            Use-CosmosDbInternalFlag -EnableCaching $true
+
+            $key = Get-AadToken
+            $key | Should -Be $MOCK_TOKEN.accessToken | Out-Null
+
+            $key = Get-AadToken
+            $key | Should -Be $MOCK_TOKEN.accessToken | Out-Null
+
+            Assert-MockCalled Get-AadTokenWithoutCaching -Times 1 -Exactly
+        }
+
+        It "Calls the core logic for each call with caching disabled" {  
+            Use-CosmosDbInternalFlag -EnableCaching $false
+
+            $key1 = Get-AadToken
+            $key1 | Should -Be $MOCK_TOKEN.accessToken | Out-Null
+
+            $key2 = Get-AadToken
+            $key2 | Should -Be $MOCK_TOKEN.accessToken | Out-Null
+
+            Assert-MockCalled Get-AadTokenWithoutCaching -Times 2 -Exactly
+        }
+
+        It "Respects token expiration" {
+            Use-CosmosDbInternalFlag -EnableCaching $true
+
+            $MOCK_TOKEN = @{
+                accessToken = "MOCK_TOKEN"
+                expires_on  = [System.DateTimeOffset]::UtcNow.ToUnixTimeSeconds()
+            }
+
+            Mock Get-AadTokenWithoutCaching {
+                return $MOCK_TOKEN
+            }
+
+            $key1 = Get-AadToken
+            $key1 | Should -Be $MOCK_TOKEN.accessToken | Out-Null
+
+            $key2 = Get-AadToken
+            $key2 | Should -Be $MOCK_TOKEN.accessToken | Out-Null
+
+            Assert-MockCalled Get-AadTokenWithoutCaching -Times 2 -Exactly
+        }
+    }
+}
diff --git a/tests/Get-AuthorizationHeader.Tests.ps1 b/tests/Get-AuthorizationHeader.Tests.ps1
@@ -13,6 +13,7 @@ InModuleScope cosmos-db {
             $MOCK_RESOURCE_TYPE = "MOCK_RESOURCE_TYPE"
             $MOCK_VERB = "MOCK_VERB"
             $MOCK_NOW = "MOCK_NOW"
+            $MOCK_AAD_TOKEN = "MOCK_AAD_TOKEN"
 
             $MOCK_MASTER_KEY_BYTES = [System.Text.Encoding]::UTF8.GetBytes('gVkYp3s6v9y$B&E)H@MbQeThWmZq4t7w')
 
@@ -25,23 +26,43 @@ InModuleScope cosmos-db {
 
                 [System.Convert]::ToBase64String($MOCK_MASTER_KEY_BYTES)
             }
+
+            Mock Get-AADToken {
+                return $MOCK_AAD_TOKEN
+            }
+        }
+
+        AfterEach {
+            $env:COSMOS_DB_FLAG_ENABLE_MASTER_KEY_AUTH = $null
         }
 
-        It "Returns the correct signature hashed with the master key" {    
-            $result = Get-AuthorizationHeader  -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB -Verb $MOCK_VERB -ResourceType $MOCK_RESOURCE_TYPE -ResourceUrl $MOCK_RESOURCE_URL -Now $MOCK_NOW
+        It "Returns the correct signature hashed with the master key" {
+            Use-CosmosDbInternalFlag -enableMasterKeyAuth $true
+
+            $result = Get-AuthorizationHeader -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB -Verb $MOCK_VERB -ResourceType $MOCK_RESOURCE_TYPE -ResourceUrl $MOCK_RESOURCE_URL -Now $MOCK_NOW
 
             $expectedSignature = "$($MOCK_VERB.ToLower())`n$($MOCK_RESOURCE_TYPE.ToLower())`n$MOCK_RESOURCE_URL`n$($MOCK_NOW.ToLower())`n`n"
 
             $hasher = New-Object System.Security.Cryptography.HMACSHA256 -Property @{ Key = $MOCK_MASTER_KEY_BYTES }
-            $sigBinary=[System.Text.Encoding]::UTF8.GetBytes($expectedSignature)
-            $hashBytes=$hasher.ComputeHash($sigBinary)
-            $expectedBase64Hash=[System.Convert]::ToBase64String($hashBytes)
+            $sigBinary = [System.Text.Encoding]::UTF8.GetBytes($expectedSignature)
+            $hashBytes = $hasher.ComputeHash($sigBinary)
+            $expectedBase64Hash = [System.Convert]::ToBase64String($hashBytes)
 
             $expectedHeader = [uri]::EscapeDataString("type=master&ver=1.0&sig=$expectedBase64Hash")
 
             $result | Should -Be $expectedHeader
 
             Assert-MockCalled Get-Base64Masterkey -Times 1
         }
+
+        It "Returns the correct signature with for entra id auth" {           
+            $result = Get-AuthorizationHeader -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB -Verb $MOCK_VERB -ResourceType $MOCK_RESOURCE_TYPE -ResourceUrl $MOCK_RESOURCE_URL -Now $MOCK_NOW
+
+            $expectedHeader = [uri]::EscapeDataString("type=aad&ver=1.0&sig=$MOCK_AAD_TOKEN")
+
+            $result | Should -Be $expectedHeader
+
+            Assert-MockCalled Get-AADToken -Times 1
+        }
     }
 }
diff --git a/tests/Get-Base64Masterkey.Tests.ps1 b/tests/Get-Base64Masterkey.Tests.ps1
@@ -26,6 +26,10 @@ InModuleScope cosmos-db {
             }
         }
 
+        AfterAll {
+            $env:COSMOS_DB_FLAG_ENABLE_READONLY_KEYS = $null
+        }
+
         BeforeEach {
             # This is defined in the main module
             $MASTER_KEY_CACHE = @{}

diff --git a/tests/Get-CosmosDbRecordContent.Tests.ps1 b/tests/Get-CosmosDbRecordContent.Tests.ps1
@@ -9,6 +9,10 @@ InModuleScope cosmos-db {
             . $PSScriptRoot\Utils.ps1
         }
 
+        AfterAll {
+            $env:COSMOS_DB_FLAG_ENABLE_READONLY_KEYS = $null
+        }
+
         It "Returns the Content of a successful response" {
             $content = @{ 
                 Key1   = "Value1";

diff --git a/tests/Invoke-CosmosDbApiRequestWithContinuation.Tests.ps1 b/tests/Invoke-CosmosDbApiRequestWithContinuation.Tests.ps1
@@ -387,7 +387,6 @@ InModuleScope cosmos-db {
         }
 
         It "Does not refresh auth headers before age threshold" {
-            # Force each call to refresh
             $AUTHORIZATION_HEADER_REFRESH_THRESHOLD = [System.TimeSpan]::FromHours(1)
 
             $continuationTokens = @($null, "100", "200", "300")