From 009ea1a27f2c14e9f849d40ee8aceb194e8926f2 Mon Sep 17 00:00:00 2001 From: Caleb Brose <5447118+cmbrose@users.noreply.github.com> Date: Mon, 18 Mar 2024 22:38:41 +0000 Subject: [PATCH 1/3] Add readonly key support --- .vscode/launch.json | 16 +++++ README.md | 20 ++++++ cosmos-db/cosmos-db.psd1 | 4 +- cosmos-db/cosmos-db.psm1 | 75 ++++++++++++++------ tests/Get-Base64Masterkey.Tests.ps1 | 102 ++++++++++++++++++++++++++++ 5 files changed, 196 insertions(+), 21 deletions(-) create mode 100644 .vscode/launch.json create mode 100644 tests/Get-Base64Masterkey.Tests.ps1 diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 0000000..c80f1a7 --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,16 @@ +{ + // Use IntelliSense to learn about possible attributes. + // Hover to view descriptions of existing attributes. + // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 + "version": "0.2.0", + "configurations": [ + { + "name": "PowerShell: Run Pester Tests", + "type": "PowerShell", + "request": "launch", + "script": "Invoke-Pester", + "createTemporaryIntegratedConsole": true, + "attachDotnetDebugger": true + } + ] +} \ No newline at end of file diff --git a/README.md b/README.md index 23dcf34..0588ddf 100644 --- a/README.md +++ b/README.md @@ -228,6 +228,26 @@ Remove-CosmosDbRecord ... | PartitionKey | The partition key of the resource | No - defaults to `Id`
Must be set if the collection uses a different parition scheme | | GetPartitionKeyBlock | Callback to get the `PartitionKey` from `Object` - useful in pipelines | No - used only if `PartitionKey` is not set | +### Use-CosmosDbReadonlyKeys + +Enables or disables using readonly DB keys in commands. By default this is disabled (i.e. writable keys are used). + +#### Examples + +```powershell +# Enable using readonly keys +Use-CosmosDbReadonlyKeys + +# Disable using readonly keys (use writable keys) +Use-CosmosDbReadonlyKeys -Disable +``` + +#### Parameters + +| Name | Usage | Required | +| - | - | - | +| Disable | Disables readonly keys if set (enables writable keys) | No - default is `$false` (which will enable readonly keys) | + ### Use-CosmosDbInternalFlag Enables or disables internal flags in the module, normally should only be used for debugging or dogfooding diff --git a/cosmos-db/cosmos-db.psd1 b/cosmos-db/cosmos-db.psd1 index 5502223..03a9cd0 100644 --- a/cosmos-db/cosmos-db.psd1 +++ b/cosmos-db/cosmos-db.psd1 @@ -11,7 +11,7 @@ # RootModule = '' # Version number of this module. - ModuleVersion = '1.15' + ModuleVersion = '1.16' # Supported PSEditions # CompatiblePSEditions = @() @@ -71,7 +71,7 @@ FunctionsToExport = 'Get-CosmosDbRecord', 'Get-AllCosmosDbRecords', 'Search-CosmosDbRecords', 'New-CosmosDbRecord', 'Update-CosmosDbRecord', 'Remove-CosmosDbRecord', - 'Get-CosmosDbRecordContent', 'Use-CosmosDbInternalFlag' + 'Get-CosmosDbRecordContent', 'Use-CosmosDbReadonlyKeys', 'Use-CosmosDbInternalFlag' # Cmdlets to export from this module, for best performance, do not use wildcards and do not delete the entry, use an empty array if there are no cmdlets to export. CmdletsToExport = '*' diff --git a/cosmos-db/cosmos-db.psm1 b/cosmos-db/cosmos-db.psm1 index 9ef8cca..ebfc83d 100644 --- a/cosmos-db/cosmos-db.psm1 +++ b/cosmos-db/cosmos-db.psm1 @@ -30,7 +30,7 @@ Function Get-DocumentsUrl([string]$Container, [string]$Collection, [string]$Reco $encodedRecordId = [uri]::EscapeDataString($RecordId) return @{ - ApiUrl = "$collectionsUrl/$DOCS_TYPE/$encodedRecordId"; + ApiUrl = "$collectionsUrl/$DOCS_TYPE/$encodedRecordId"; ResourceUrl = "$collectionsUrl/$DOCS_TYPE/$RecordId"; } } @@ -61,20 +61,36 @@ Function Set-CacheValue([string]$key, $value, [hashtable]$cache, [int]$expiratio } Function Get-Base64Masterkey([string]$ResourceGroup, [string]$Database, [string]$SubscriptionId) { - $cacheKey = "$SubscriptionId/$ResourceGroup/$Database" + $readonly = $env:COSMOS_DB_FLAG_ENABLE_READONLY_KEYS -eq 1 + + $cacheKey = "$SubscriptionId/$ResourceGroup/$Database/$readonly" $cacheResult = Get-CacheValue -Key $cacheKey -Cache $MASTER_KEY_CACHE if ($cacheResult) { return $cacheResult } - if ($SubscriptionId) { - $masterKey = az cosmosdb keys list --name $Database --query primaryMasterKey --output tsv --resource-group $ResourceGroup --subscription $SubscriptionId + $masterKey = Get-Base64MasterkeyWithoutCaching -ResourceGroup $ResourceGroup -Database $Database -SubscriptionId $SubscriptionId -Readonly $readonly + + Set-CacheValue -Key $cacheKey -Value $masterKey -Cache $MASTER_KEY_CACHE -ExpirationHours 6 + + $masterKey +} + +# This is just to support testing caching with Get-Base64Masterkey and isn't meant to be used directly +Function Get-Base64MasterkeyWithoutCaching([string]$ResourceGroup, [string]$Database, [string]$SubscriptionId, [bool]$Readonly) { + $query = if ($readonly) { + "primaryReadonlyMasterKey" } else { - $masterKey = az cosmosdb keys list --name $Database --query primaryMasterKey --output tsv --resource-group $ResourceGroup + "primaryMasterKey" } - Set-CacheValue -Key $cacheKey -Value $masterKey -Cache $MASTER_KEY_CACHE -ExpirationHours 6 + if ($SubscriptionId) { + $masterKey = az cosmosdb keys list --name $Database --query $query --output tsv --resource-group $ResourceGroup --subscription $SubscriptionId + } + else { + $masterKey = az cosmosdb keys list --name $Database --query $query --output tsv --resource-group $ResourceGroup + } $masterKey } @@ -140,9 +156,9 @@ Function Get-CommonHeaders([string]$now, [string]$encodedAuthString, [string]$co $headers["x-ms-documentdb-partitionkey"] = "[`"$PartitionKey`"]" } - if ($Etag) { - $headers["If-Match"] = $Etag - } + if ($Etag) { + $headers["If-Match"] = $Etag + } $headers } @@ -172,7 +188,7 @@ Function Get-ExceptionResponseOrThrow($err) { if ($err.Exception.Response) { $msg = @{ - StatusCode = $err.Exception.Response.StatusCode; + StatusCode = $err.Exception.Response.StatusCode; RawResponse = $err.Exception.Response; } @@ -180,7 +196,8 @@ Function Get-ExceptionResponseOrThrow($err) { # In PS Core, the body is eaten and put into this message # See: https://stackoverflow.com/questions/18771424/how-to-get-powershell-invoke-restmethod-to-return-body-of-http-500-code-response $msg.Content = $err.ErrorDetails.Message - } else { + } + else { # In Desktop we can re-read the content stream $result = $err.Exception.Response.GetResponseStream() $reader = New-Object System.IO.StreamReader($result) @@ -191,7 +208,8 @@ Function Get-ExceptionResponseOrThrow($err) { } return [PSCustomObject]$msg - } else { + } + else { throw $err.Exception } } @@ -802,7 +820,8 @@ Function Update-CosmosDbRecord { if ($EnforceOptimisticConcurrency) { $headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -PartitionKey $requestPartitionKey -Etag $Object._etag - } else { + } + else { $headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -PartitionKey $requestPartitionKey } @@ -911,11 +930,12 @@ Function Get-CosmosDbRecordContent([parameter(ValueFromPipeline)]$RecordResponse process { $code = [int]$RecordResponse.StatusCode $content = - if ($RecordResponse.Content) { - $RecordResponse.Content | ConvertFrom-Json - } else { - $null - } + if ($RecordResponse.Content) { + $RecordResponse.Content | ConvertFrom-Json + } + else { + $null + } if ($code -lt 300) { if ($RecordResponse.Content) { @@ -925,6 +945,12 @@ Function Get-CosmosDbRecordContent([parameter(ValueFromPipeline)]$RecordResponse $null } } + elseif ($code -eq 401) { + if ($env:COSMOS_DB_FLAG_ENABLE_READONLY_KEYS -eq 1) { + throw "Unauthorized (used a readonly key)" + } + throw "Unauthorized" + } elseif ($code -eq 404) { if ($content.Message -like "*Owner resource does not exist*") { throw "Database does not exist" @@ -935,6 +961,7 @@ Function Get-CosmosDbRecordContent([parameter(ValueFromPipeline)]$RecordResponse elseif ($code -eq 429) { throw "Request rate limited" } + else { $message = $content.Message throw "Request failed with status code $code with message`n`n$message" @@ -961,6 +988,14 @@ Function Use-CosmosDbInternalFlag } } +Function Use-CosmosDbReadonlyKeys +( + [switch]$Disable +) { + $env:COSMOS_DB_FLAG_ENABLE_READONLY_KEYS = if ($Disable) { 0 } else { 1 } +} + + Export-ModuleMember -Function "Get-CosmosDbRecord" Export-ModuleMember -Function "Get-AllCosmosDbRecords" @@ -974,4 +1009,6 @@ Export-ModuleMember -Function "Remove-CosmosDbRecord" Export-ModuleMember -Function "Get-CosmosDbRecordContent" -Export-ModuleMember -Function "Use-CosmosDbInternalFlag" \ No newline at end of file +Export-ModuleMember -Function "Use-CosmosDbReadonlyKeys" + +Export-ModuleMember -Function "Use-CosmosDbInternalFlag" diff --git a/tests/Get-Base64Masterkey.Tests.ps1 b/tests/Get-Base64Masterkey.Tests.ps1 new file mode 100644 index 0000000..d039010 --- /dev/null +++ b/tests/Get-Base64Masterkey.Tests.ps1 @@ -0,0 +1,102 @@ +Get-Module cosmos-db | Remove-Module -Force +Import-Module $PSScriptRoot\..\cosmos-db\cosmos-db.psm1 -Force + +InModuleScope cosmos-db { + Describe "Get-Base64Masterkey" { + BeforeAll { + $MOCK_SUB = "MOCK_SUB" + $MOCK_RG = "MOCK_RG" + $MOCK_DB = "MOCK_DB" + $MOCK_KEY = "MOCK_KEY" + $MOCK_READONLY_KEY = "MOCK_READONLY_KEY" + + Mock Get-Base64MasterkeyWithoutCaching { + param($ResourceGroup, $Database, $SubscriptionId, $Readonly) + + $ResourceGroup | Should -Be $MOCK_RG | Out-Null + $Database | Should -Be $MOCK_DB | Out-Null + $SubscriptionId | Should -Be $MOCK_SUB | Out-Null + + if ($Readonly) { + $MOCK_READONLY_KEY + } + else { + $MOCK_KEY + } + } + } + + BeforeEach { + # This is defined in the main module + $MASTER_KEY_CACHE = @{} + } + + It "Only calls the core logic once with caching enabled" { + Use-CosmosDbReadonlyKeys -Disable + Use-CosmosDbInternalFlag -EnableCaching $true + + $key = Get-Base64Masterkey -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB + $key | Should -Be $MOCK_KEY | Out-Null + + $key = Get-Base64Masterkey -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB + $key | Should -Be $MOCK_KEY | Out-Null + + Assert-MockCalled Get-Base64MasterkeyWithoutCaching -Times 1 -Exactly + } + + It "Should return the readonly key when configured" { + Use-CosmosDbReadonlyKeys + Use-CosmosDbInternalFlag -EnableCaching $true + + $key = Get-Base64Masterkey -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB + $key | Should -Be $MOCK_READONLY_KEY | Out-Null + + Assert-MockCalled Get-Base64MasterkeyWithoutCaching -Times 1 -Exactly + } + + It "Should not use a cached key when switching between readonly and writable" { + Use-CosmosDbInternalFlag -EnableCaching $true + + Use-CosmosDbReadonlyKeys -Disable + + $key = Get-Base64Masterkey -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB + $key | Should -Be $MOCK_KEY | Out-Null + + $key = Get-Base64Masterkey -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB + $key | Should -Be $MOCK_KEY | Out-Null + + Assert-MockCalled Get-Base64MasterkeyWithoutCaching -Times 1 -Exactly + + Use-CosmosDbReadonlyKeys + + $key = Get-Base64Masterkey -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB + $key | Should -Be $MOCK_READONLY_KEY | Out-Null + + $key = Get-Base64Masterkey -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB + $key | Should -Be $MOCK_READONLY_KEY | Out-Null + + Assert-MockCalled Get-Base64MasterkeyWithoutCaching -Times 2 -Exactly + + Use-CosmosDbReadonlyKeys -Disable + + $key = Get-Base64Masterkey -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB + $key | Should -Be $MOCK_KEY | Out-Null + + # This one should reuse the existing writable key cached value and shouldn't increment the call counter + Assert-MockCalled Get-Base64MasterkeyWithoutCaching -Times 2 -Exactly + } + + It "Calls the core logic for each call with caching disabled" { + Use-CosmosDbReadonlyKeys -Disable + Use-CosmosDbInternalFlag -EnableCaching $false + + $key1 = Get-Base64Masterkey -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB + $key1 | Should -Be $MOCK_KEY | Out-Null + + $key2 = Get-Base64Masterkey -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB + $key2 | Should -Be $MOCK_KEY | Out-Null + + Assert-MockCalled Get-Base64MasterkeyWithoutCaching -Times 2 -Exactly + } + } +} \ No newline at end of file From 0d5bb9ba293951a09d9b8a082ea521882d4c994f Mon Sep 17 00:00:00 2001 From: Caleb Brose <5447118+cmbrose@users.noreply.github.com> Date: Tue, 19 Mar 2024 00:13:06 +0000 Subject: [PATCH 2/3] Update get content test --- tests/Get-CosmosDbRecordContent.Tests.ps1 | 38 +++++++++++++++++------ 1 file changed, 29 insertions(+), 9 deletions(-) diff --git a/tests/Get-CosmosDbRecordContent.Tests.ps1 b/tests/Get-CosmosDbRecordContent.Tests.ps1 index 6424ccf..7ddb7a0 100644 --- a/tests/Get-CosmosDbRecordContent.Tests.ps1 +++ b/tests/Get-CosmosDbRecordContent.Tests.ps1 @@ -11,8 +11,8 @@ InModuleScope cosmos-db { It "Returns the Content of a successful response" { $content = @{ - Key1 = "Value1"; - Key2 = 2; + Key1 = "Value1"; + Key2 = 2; Nested = @{ NestedKey1 = "NestedValue1"; NestedKey2 = 2; @@ -21,7 +21,7 @@ InModuleScope cosmos-db { $response = @{ StatusCode = 200; - Content = ($content | ConvertTo-Json -Depth 100) + Content = ($content | ConvertTo-Json -Depth 100) } $result = $response | Get-CosmosDbRecordContent @@ -42,9 +42,9 @@ InModuleScope cosmos-db { It "Throws not found for 404 - DB not found" { $response = @{ StatusCode = 404; - Content = (@{ - Message = "{`"Errors`":[`"Owner resource does not exist`"]}" - } | ConvertTo-Json) + Content = (@{ + Message = "{`"Errors`":[`"Owner resource does not exist`"]}" + } | ConvertTo-Json) } { $response | Get-CosmosDbRecordContent } | Should -Throw "Database does not exist" @@ -65,16 +65,36 @@ InModuleScope cosmos-db { { $response | Get-CosmosDbRecordContent } | Should -Throw "Request rate limited" } + + It "Throws unauthorized for 401" { + Use-CosmosDbReadonlyKeys -Disable + + $response = [pscustomobject] @{ + StatusCode = 401; + } + + { $response | Get-CosmosDbRecordContent } | Should -Throw "Unauthorized" + } + + It "Throws unauthorized for 401 with a message about readonly keys" { + Use-CosmosDbReadonlyKeys + + $response = [pscustomobject] @{ + StatusCode = 401; + } + + { $response | Get-CosmosDbRecordContent } | Should -Throw "Unauthorized (used a readonly key)" + } It "Throws useful error for unknown errors" { $errorMessage = "Mock error message" $response = [pscustomobject] @{ - StatusCode = 401; - Content = ( @{ Message = $errorMessage; } | ConvertTo-Json); + StatusCode = 1234; + Content = ( @{ Message = $errorMessage; } | ConvertTo-Json); } - { $response | Get-CosmosDbRecordContent } | Should -Throw "Request failed with status code 401 with message`n`n$errorMessage" + { $response | Get-CosmosDbRecordContent } | Should -Throw "Request failed with status code 1234 with message`n`n$errorMessage" } } } \ No newline at end of file From 87a1ab7626230fb93d3072a188e0aa5725834f83 Mon Sep 17 00:00:00 2001 From: Caleb Brose <5447118+cmbrose@users.noreply.github.com> Date: Tue, 19 Mar 2024 01:00:48 +0000 Subject: [PATCH 3/3] Fork the pester tester --- .github/workflows/ci.yml | 98 ++++++++++++++++++++-------------------- 1 file changed, 49 insertions(+), 49 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4b2f5ba..81d3f58 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,9 +2,9 @@ name: CI on: push: - branches: [ main ] + branches: [main] pull_request: - branches: [ main ] + branches: [main] workflow_dispatch: @@ -12,51 +12,51 @@ jobs: run-tests: runs-on: ubuntu-latest steps: - - name: checkout - uses: actions/checkout@v1 - - - name: test module - id: test_module - uses: zyborg/pester-tests-report@v1 - with: - include_paths: tests - report_name: cosmos-db tests - report_title: cosmos-db tests - gist_name: pester-tests-report.md - github_token: ${{ secrets.GITHUB_TOKEN }} - gist_token: ${{ secrets.GIST_TOKEN }} - gist_badge_label: Tests %ExecutedAt% + - name: checkout + uses: actions/checkout@v1 - - name: dump results - shell: pwsh - run: | - ipmo GitHubActions - $test_results_path = '${{ steps.test_module.outputs.test_results_path }}' - $error_message = '${{ steps.test_module.outputs.error_message }}' - $error_clixml_path = '${{ steps.test_module.outputs.error_clixml_path }}' - $result_clixml_path = '${{ steps.test_module.outputs.result_clixml_path }}' - $result_value = '${{ steps.test_module.outputs.result_value }}' - $total_count = '${{ steps.test_module.outputs.total_count }}' - $passed_count = '${{ steps.test_module.outputs.passed_count }}' - $failed_count = '${{ steps.test_module.outputs.failed_count }}' - Write-ActionInfo "Found these outputs from [test_module]:" - Write-ActionInfo " * test_results_path = $test_results_path" - Write-ActionInfo " * error_message = $error_message" - Write-ActionInfo " * error_clixml_path = $error_clixml_path" - Write-ActionInfo " * result_clixml_path = $result_clixml_path" - Write-ActionInfo " * result_value = $result_value" - Write-ActionInfo " * total_count = $total_count" - Write-ActionInfo " * passed_count = $passed_count" - Write-ActionInfo " * failed_count = $failed_count" - if ($error_clixml_path) { - $er = Import-Clixml -Path $error_clixml_path - Write-ActionInfo "Loaded up the ErrorRecord:" - $er - $er.Exception - } - if ($result_clixml_path) { - $pr = Import-Clixml -Path $result_clixml_path - Write-ActionInfo "Loaded up the Pester Result:" - $pr - } - exit $failed_count \ No newline at end of file + - name: test module + id: test_module + uses: cmbrose/pester-tests-report@v1 + with: + include_paths: tests + report_name: cosmos-db tests + report_title: cosmos-db tests + gist_name: pester-tests-report.md + github_token: ${{ secrets.GITHUB_TOKEN }} + gist_token: ${{ secrets.GIST_TOKEN }} + gist_badge_label: Tests %ExecutedAt% + + - name: dump results + shell: pwsh + run: | + ipmo GitHubActions + $test_results_path = '${{ steps.test_module.outputs.test_results_path }}' + $error_message = '${{ steps.test_module.outputs.error_message }}' + $error_clixml_path = '${{ steps.test_module.outputs.error_clixml_path }}' + $result_clixml_path = '${{ steps.test_module.outputs.result_clixml_path }}' + $result_value = '${{ steps.test_module.outputs.result_value }}' + $total_count = '${{ steps.test_module.outputs.total_count }}' + $passed_count = '${{ steps.test_module.outputs.passed_count }}' + $failed_count = '${{ steps.test_module.outputs.failed_count }}' + Write-ActionInfo "Found these outputs from [test_module]:" + Write-ActionInfo " * test_results_path = $test_results_path" + Write-ActionInfo " * error_message = $error_message" + Write-ActionInfo " * error_clixml_path = $error_clixml_path" + Write-ActionInfo " * result_clixml_path = $result_clixml_path" + Write-ActionInfo " * result_value = $result_value" + Write-ActionInfo " * total_count = $total_count" + Write-ActionInfo " * passed_count = $passed_count" + Write-ActionInfo " * failed_count = $failed_count" + if ($error_clixml_path) { + $er = Import-Clixml -Path $error_clixml_path + Write-ActionInfo "Loaded up the ErrorRecord:" + $er + $er.Exception + } + if ($result_clixml_path) { + $pr = Import-Clixml -Path $result_clixml_path + Write-ActionInfo "Loaded up the Pester Result:" + $pr + } + exit $failed_count