Skip to content

Commit 4ecd157

Browse files
cmbrosecmbrose
authored
Add support for AAD auth, enabled by default (#19)
* Base work for aad auth * Add read only tests * Add AadToken tests * Typo
1 parent 2e153a9 commit 4ecd157

12 files changed

+247
-63
lines changed

Diff for: ‎CHANGELOG.md

+4
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
# Release Notes
22
All notable changes and release history of the "cosmos-db" module will be documented in this file.
33

4+
## 1.19
5+
* Changed to Entra Id authentication by default.
6+
* Adds a `-enableMasterKeyAuth` flag to `Use-CosmosDbInternalFlag` which reverts auth to use master keys.
7+
48
## 1.18
59
* 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.
610
* Adds a `-enableAuthHeaderReuse` flag to `Use-CosmosDbInternalFlag` which disables the 10 minute refresh period and forces auth header refreshes for every API call.

Diff for: ‎README.md

+1
Original file line numberDiff line numberDiff line change
@@ -266,6 +266,7 @@ Use-CosmosDbInternalFlag -EnableFiddlerDebugging $true
266266
| EnableCaching | Enables/disables caching certain values like DB keys, partition ranges, etc. Improves performance of nearly all operations. | No - default is enabled |
267267
| 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 |
268268
| EnableAuthHeaderReuse | Enables/disables reusing auth headers for commands which use continuation tokens, like `Search-CosmosDbRecords` or `Get-AllCosmosDbRecords`. | No - default is enabled |
269+
| EnableMasterKeyAuth | Enables/disables using master key based authentication instead of Entra Id. | No - default is disabled |
269270

270271
## Error Handling
271272

Diff for: ‎cosmos-db/cosmos-db.psd1

+1-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
# RootModule = ''
1212

1313
# Version number of this module.
14-
ModuleVersion = '1.18'
14+
ModuleVersion = '1.19'
1515

1616
# Supported PSEditions
1717
# CompatiblePSEditions = @()

Diff for: ‎cosmos-db/cosmos-db.psm1

+63-10
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ $API_VERSION = "2018-12-31"
1515
$AUTHORIZATION_HEADER_REFRESH_THRESHOLD = [System.TimeSpan]::FromMinutes(10);
1616

1717
$MASTER_KEY_CACHE = @{}
18+
$AAD_TOKEN_CACHE = @{}
1819
$SIGNATURE_HASH_CACHE = @{}
1920
$PARTITION_KEY_RANGE_CACHE = @{}
2021

@@ -57,13 +58,23 @@ Function Get-CacheValue([string]$key, [hashtable]$cache) {
5758
return $null
5859
}
5960

60-
Function Set-CacheValue([string]$key, $value, [hashtable]$cache, [int]$expirationHours) {
61+
Function Set-CacheValue([string]$key, $value, [hashtable]$cache, [int]$expirationHours, [int]$expirationMinutes) {
6162
$cache[$key] = @{
62-
Expiration = [datetime]::UtcNow.AddHours($expirationHours);
63+
Expiration = [datetime]::UtcNow.AddHours($expirationHours).AddMinutes($expirationMinutes);
6364
Value = $value
6465
}
6566
}
6667

68+
Function RequireWriteableKey() {
69+
# Entra Id auth does not support read only tokens (because it's just handled by RBAC).
70+
# So if the user requested to be in read only mode, we have to manage it ourselves.
71+
72+
$readonlyEnabled = $env:COSMOS_DB_FLAG_ENABLE_READONLY_KEYS -eq 1
73+
if ($readonlyEnabled) {
74+
throw "Operation not allowed in readonly mode"
75+
}
76+
}
77+
6778
Function Get-Base64Masterkey([string]$ResourceGroup, [string]$Database, [string]$SubscriptionId) {
6879
$readonly = $env:COSMOS_DB_FLAG_ENABLE_READONLY_KEYS -eq 1
6980

@@ -99,6 +110,31 @@ Function Get-Base64MasterkeyWithoutCaching([string]$ResourceGroup, [string]$Data
99110
$masterKey
100111
}
101112

113+
Function Get-AADToken([string]$ResourceGroup, [string]$Database, [string]$SubscriptionId) {
114+
$cacheKey = "aadtoken"
115+
$cacheResult = Get-CacheValue -Key $cacheKey -Cache $AAD_TOKEN_CACHE
116+
if ($cacheResult) {
117+
return $cacheResult
118+
}
119+
120+
$oauth = Get-AADTokenWithoutCaching
121+
$token = $oauth.AccessToken
122+
123+
$expirationUtcTimestamp = $oauth.expires_on
124+
$expiration = [System.DateTimeOffset]::FromUnixTimeSeconds($expirationUtcTimestamp).DateTime
125+
$remainingTime = $expiration - [System.DateTime]::UtcNow
126+
$cacheExpirationMinutes = ($remainingTime - $AUTHORIZATION_HEADER_REFRESH_THRESHOLD).Minutes
127+
128+
Set-CacheValue -Key $cacheKey -Value $token -Cache $AAD_TOKEN_CACHE -ExpirationMinutes $cacheExpirationMinutes
129+
130+
$token
131+
}
132+
133+
# This is just to support testing caching with Get-AADToken and isn't meant to be used directly
134+
Function Get-AADTokenWithoutCaching() {
135+
az account get-access-token --resource "https://cosmos.azure.com" | ConvertFrom-Json
136+
}
137+
102138
Function Get-Signature([string]$verb, [string]$resourceType, [string]$resourceUrl, [string]$now) {
103139
$parts = @(
104140
$verb.ToLower(),
@@ -128,19 +164,25 @@ Function Get-Base64EncryptedSignatureHash([string]$masterKey, [string]$signature
128164
$base64Hash
129165
}
130166

131-
Function Get-EncodedAuthString([string]$signatureHash) {
132-
$authString = "type=master&ver=1.0&sig=$signatureHash"
167+
Function Get-EncodedAuthString([string]$signatureHash, [string]$type) {
168+
$authString = "type=$type&ver=1.0&sig=$signatureHash"
133169
[uri]::EscapeDataString($authString)
134170
}
135171

136172
Function Get-AuthorizationHeader([string]$ResourceGroup, [string]$SubscriptionId, [string]$Database, [string]$verb, [string]$resourceType, [string]$resourceUrl, [string]$now) {
137-
$masterKey = Get-Base64Masterkey -ResourceGroup $ResourceGroup -Database $Database -SubscriptionId $SubscriptionId
173+
if ($env:COSMOS_DB_FLAG_ENABLE_MASTER_KEY_AUTH -eq 1) {
174+
$masterKey = Get-Base64Masterkey -ResourceGroup $ResourceGroup -Database $Database -SubscriptionId $SubscriptionId
138175

139-
$signature = Get-Signature -verb $verb -resourceType $resourceType -resourceUrl $resourceUrl -now $now
176+
$signature = Get-Signature -verb $verb -resourceType $resourceType -resourceUrl $resourceUrl -now $now
177+
178+
$signatureHash = Get-Base64EncryptedSignatureHash -masterKey $masterKey -signature $signature
179+
180+
return Get-EncodedAuthString -signatureHash $signatureHash -type "master"
181+
}
140182

141-
$signatureHash = Get-Base64EncryptedSignatureHash -masterKey $masterKey -signature $signature
183+
$token = Get-AADToken
142184

143-
Get-EncodedAuthString -signatureHash $signatureHash
185+
return Get-EncodedAuthString -signatureHash $token -type "aad"
144186
}
145187

146188
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,
268310
$headers["x-ms-continuation"] = $continuationToken
269311

270312
$authHeaderReuseDisabled = $env:COSMOS_DB_FLAG_ENABLE_AUTH_HEADER_REUSE -eq 0
271-
$authHeaderExpired = [System.DateTime]::Parse($authHeaders.now) + $AUTHORIZATION_HEADER_REFRESH_THRESHOLD -lt [System.DateTime]::UtcNow
313+
$authHeaderExpired = [System.DateTimeOffset]::Parse($authHeaders.now) + $AUTHORIZATION_HEADER_REFRESH_THRESHOLD -lt [System.DateTime]::UtcNow
272314

273315
if ($authHeaderReuseDisabled -or $authHeaderExpired) {
274316
$authHeaders = Invoke-Command -ScriptBlock $refreshAuthHeaders
@@ -746,6 +788,8 @@ Function New-CosmosDbRecord {
746788
)
747789

748790
begin {
791+
RequireWriteableKey
792+
749793
$baseUrl = Get-BaseDatabaseUrl $Database
750794
$collectionsUrl = Get-CollectionsUrl $Container $Collection
751795
$docsUrl = "$collectionsUrl/$DOCS_TYPE"
@@ -833,6 +877,8 @@ Function Update-CosmosDbRecord {
833877
)
834878

835879
begin {
880+
RequireWriteableKey
881+
836882
$baseUrl = Get-BaseDatabaseUrl $Database
837883
}
838884
process {
@@ -929,6 +975,8 @@ Function Remove-CosmosDbRecord {
929975
)
930976

931977
begin {
978+
RequireWriteableKey
979+
932980
$baseUrl = Get-BaseDatabaseUrl $Database
933981
}
934982
process {
@@ -1005,7 +1053,8 @@ Function Use-CosmosDbInternalFlag
10051053
$enableFiddlerDebugging = $null,
10061054
$enableCaching = $null,
10071055
$enablePartitionKeyRangeSearches = $null,
1008-
$enableAuthHeaderReuse = $null
1056+
$enableAuthHeaderReuse = $null,
1057+
$enableMasterKeyAuth = $null
10091058
) {
10101059
if ($null -ne $enableFiddlerDebugging) {
10111060
$env:AZURE_CLI_DISABLE_CONNECTION_VERIFICATION = if ($enableFiddlerDebugging) { 1 } else { 0 }
@@ -1022,6 +1071,10 @@ Function Use-CosmosDbInternalFlag
10221071
if ($null -ne $enableAuthHeaderReuse) {
10231072
$env:COSMOS_DB_FLAG_ENABLE_AUTH_HEADER_REUSE = if ($enableAuthHeaderReuse) { 1 } else { 0 }
10241073
}
1074+
1075+
if ($null -ne $enableMasterKeyAuth) {
1076+
$env:COSMOS_DB_FLAG_ENABLE_MASTER_KEY_AUTH = if ($enableMasterKeyAuth) { 1 } else { 0 }
1077+
}
10251078
}
10261079

10271080
Function Use-CosmosDbReadonlyKeys

Diff for: ‎tests/Get-AadToken.Tests.ps1

+67
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
Get-Module cosmos-db | Remove-Module -Force
2+
Import-Module $PSScriptRoot\..\cosmos-db\cosmos-db.psm1 -Force
3+
4+
InModuleScope cosmos-db {
5+
Describe "Get-AadToken" {
6+
BeforeAll {
7+
$MOCK_TOKEN = @{
8+
accessToken = "MOCK_TOKEN"
9+
expires_on = [System.DateTimeOffset]::UtcNow.AddHours(1).ToUnixTimeSeconds()
10+
}
11+
12+
Mock Get-AadTokenWithoutCaching {
13+
return $MOCK_TOKEN
14+
}
15+
}
16+
17+
BeforeEach {
18+
# This is defined in the main module
19+
$AAD_TOKEN_CACHE = @{}
20+
}
21+
22+
It "Only calls the core logic once with caching enabled" {
23+
Use-CosmosDbInternalFlag -EnableCaching $true
24+
25+
$key = Get-AadToken
26+
$key | Should -Be $MOCK_TOKEN.accessToken | Out-Null
27+
28+
$key = Get-AadToken
29+
$key | Should -Be $MOCK_TOKEN.accessToken | Out-Null
30+
31+
Assert-MockCalled Get-AadTokenWithoutCaching -Times 1 -Exactly
32+
}
33+
34+
It "Calls the core logic for each call with caching disabled" {
35+
Use-CosmosDbInternalFlag -EnableCaching $false
36+
37+
$key1 = Get-AadToken
38+
$key1 | Should -Be $MOCK_TOKEN.accessToken | Out-Null
39+
40+
$key2 = Get-AadToken
41+
$key2 | Should -Be $MOCK_TOKEN.accessToken | Out-Null
42+
43+
Assert-MockCalled Get-AadTokenWithoutCaching -Times 2 -Exactly
44+
}
45+
46+
It "Respects token expiration" {
47+
Use-CosmosDbInternalFlag -EnableCaching $true
48+
49+
$MOCK_TOKEN = @{
50+
accessToken = "MOCK_TOKEN"
51+
expires_on = [System.DateTimeOffset]::UtcNow.ToUnixTimeSeconds()
52+
}
53+
54+
Mock Get-AadTokenWithoutCaching {
55+
return $MOCK_TOKEN
56+
}
57+
58+
$key1 = Get-AadToken
59+
$key1 | Should -Be $MOCK_TOKEN.accessToken | Out-Null
60+
61+
$key2 = Get-AadToken
62+
$key2 | Should -Be $MOCK_TOKEN.accessToken | Out-Null
63+
64+
Assert-MockCalled Get-AadTokenWithoutCaching -Times 2 -Exactly
65+
}
66+
}
67+
}

Diff for: ‎tests/Get-AuthorizationHeader.Tests.ps1

+26-5
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ InModuleScope cosmos-db {
1313
$MOCK_RESOURCE_TYPE = "MOCK_RESOURCE_TYPE"
1414
$MOCK_VERB = "MOCK_VERB"
1515
$MOCK_NOW = "MOCK_NOW"
16+
$MOCK_AAD_TOKEN = "MOCK_AAD_TOKEN"
1617

1718
$MOCK_MASTER_KEY_BYTES = [System.Text.Encoding]::UTF8.GetBytes('gVkYp3s6v9y$B&E)H@MbQeThWmZq4t7w')
1819

@@ -25,23 +26,43 @@ InModuleScope cosmos-db {
2526

2627
[System.Convert]::ToBase64String($MOCK_MASTER_KEY_BYTES)
2728
}
29+
30+
Mock Get-AADToken {
31+
return $MOCK_AAD_TOKEN
32+
}
33+
}
34+
35+
AfterEach {
36+
$env:COSMOS_DB_FLAG_ENABLE_MASTER_KEY_AUTH = $null
2837
}
2938

30-
It "Returns the correct signature hashed with the master key" {
31-
$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
39+
It "Returns the correct signature hashed with the master key" {
40+
Use-CosmosDbInternalFlag -enableMasterKeyAuth $true
41+
42+
$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
3243

3344
$expectedSignature = "$($MOCK_VERB.ToLower())`n$($MOCK_RESOURCE_TYPE.ToLower())`n$MOCK_RESOURCE_URL`n$($MOCK_NOW.ToLower())`n`n"
3445

3546
$hasher = New-Object System.Security.Cryptography.HMACSHA256 -Property @{ Key = $MOCK_MASTER_KEY_BYTES }
36-
$sigBinary=[System.Text.Encoding]::UTF8.GetBytes($expectedSignature)
37-
$hashBytes=$hasher.ComputeHash($sigBinary)
38-
$expectedBase64Hash=[System.Convert]::ToBase64String($hashBytes)
47+
$sigBinary = [System.Text.Encoding]::UTF8.GetBytes($expectedSignature)
48+
$hashBytes = $hasher.ComputeHash($sigBinary)
49+
$expectedBase64Hash = [System.Convert]::ToBase64String($hashBytes)
3950

4051
$expectedHeader = [uri]::EscapeDataString("type=master&ver=1.0&sig=$expectedBase64Hash")
4152

4253
$result | Should -Be $expectedHeader
4354

4455
Assert-MockCalled Get-Base64Masterkey -Times 1
4556
}
57+
58+
It "Returns the correct signature with for entra id auth" {
59+
$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
60+
61+
$expectedHeader = [uri]::EscapeDataString("type=aad&ver=1.0&sig=$MOCK_AAD_TOKEN")
62+
63+
$result | Should -Be $expectedHeader
64+
65+
Assert-MockCalled Get-AADToken -Times 1
66+
}
4667
}
4768
}

Diff for: ‎tests/Get-Base64Masterkey.Tests.ps1

+4
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,10 @@ InModuleScope cosmos-db {
2626
}
2727
}
2828

29+
AfterAll {
30+
$env:COSMOS_DB_FLAG_ENABLE_READONLY_KEYS = $null
31+
}
32+
2933
BeforeEach {
3034
# This is defined in the main module
3135
$MASTER_KEY_CACHE = @{}

Diff for: ‎tests/Get-CosmosDbRecordContent.Tests.ps1

+4
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,10 @@ InModuleScope cosmos-db {
99
. $PSScriptRoot\Utils.ps1
1010
}
1111

12+
AfterAll {
13+
$env:COSMOS_DB_FLAG_ENABLE_READONLY_KEYS = $null
14+
}
15+
1216
It "Returns the Content of a successful response" {
1317
$content = @{
1418
Key1 = "Value1";

Diff for: ‎tests/Invoke-CosmosDbApiRequestWithContinuation.Tests.ps1

-1
Original file line numberDiff line numberDiff line change
@@ -387,7 +387,6 @@ InModuleScope cosmos-db {
387387
}
388388

389389
It "Does not refresh auth headers before age threshold" {
390-
# Force each call to refresh
391390
$AUTHORIZATION_HEADER_REFRESH_THRESHOLD = [System.TimeSpan]::FromHours(1)
392391

393392
$continuationTokens = @($null, "100", "200", "300")

0 commit comments

Comments
 (0)