Fix the batch processing RBAC bug (#579)

* Fix the batch processing RBAC bug * Add some clarifications to the docs * Create storage queue role for keys as well * Initialise queue client differently for rbac and keys * Fix bicep rule * Fix msg policy
Azure-Samples · Apr 4, 2024 · 9b182ab · 9b182ab
1 parent 5a205d2
commit 9b182ab
Show file tree

Hide file tree

Showing 5 changed files with 201 additions and 23 deletions.
diff --git a/code/backend/batch/BatchStartProcessing.py b/code/backend/batch/BatchStartProcessing.py
@@ -2,9 +2,11 @@
 import json
 import azure.functions as func
 import sys
-from azure.storage.queue import QueueClient, BinaryBase64EncodePolicy
 from utilities.helpers.EnvHelper import EnvHelper
-from utilities.helpers.AzureBlobStorageHelper import AzureBlobStorageClient
+from utilities.helpers.AzureBlobStorageHelper import (
+    AzureBlobStorageClient,
+    create_queue_client,
+)
 
 sys.path.append("..")
 bp_batch_start_processing = func.Blueprint()
@@ -25,13 +27,9 @@ def batch_start_processing(req: func.HttpRequest) -> func.HttpResponse:
         else files_data
     )
     files_data = list(map(lambda x: {"filename": x["filename"]}, files_data))
-    # Create the QueueClient object
-    queue_client = QueueClient.from_connection_string(
-        azure_blob_storage_client.connect_str,
-        env_helper.DOCUMENT_PROCESSING_QUEUE_NAME,
-        message_encode_policy=BinaryBase64EncodePolicy(),
-    )
+
     # Send a message to the queue for each file
+    queue_client = create_queue_client()
     for fd in files_data:
         queue_client.send_message(json.dumps(fd).encode("utf-8"))
 

diff --git a/code/backend/batch/utilities/helpers/AzureBlobStorageHelper.py b/code/backend/batch/utilities/helpers/AzureBlobStorageHelper.py
@@ -7,10 +7,35 @@
     ContentSettings,
     UserDelegationKey,
 )
+from azure.storage.queue import QueueClient, BinaryBase64EncodePolicy
 from .EnvHelper import EnvHelper
 from azure.identity import DefaultAzureCredential
 
 
+def connection_string(account_name: str, account_key: str):
+    return f"DefaultEndpointsProtocol=https;AccountName={account_name};AccountKey={account_key};EndpointSuffix=core.windows.net"
+
+
+def create_queue_client():
+    env_helper: EnvHelper = EnvHelper()
+    if env_helper.AZURE_AUTH_TYPE == "rbac":
+        return QueueClient(
+            account_url=f"https://{env_helper.AZURE_BLOB_ACCOUNT_NAME}.queue.core.windows.net/",
+            queue_name=env_helper.DOCUMENT_PROCESSING_QUEUE_NAME,
+            credential=DefaultAzureCredential(),
+            message_encode_policy=BinaryBase64EncodePolicy(),
+        )
+
+    else:
+        return QueueClient.from_connection_string(
+            conn_str=connection_string(
+                env_helper.AZURE_BLOB_ACCOUNT_NAME, env_helper.AZURE_BLOB_ACCOUNT_KEY
+            ),
+            queue_name=env_helper.DOCUMENT_PROCESSING_QUEUE_NAME,
+            message_encode_policy=BinaryBase64EncodePolicy(),
+        )
+
+
 class AzureBlobStorageClient:
     def __init__(
         self,
@@ -25,29 +50,27 @@ def __init__(
             self.account_name = (
                 account_name if account_name else env_helper.AZURE_BLOB_ACCOUNT_NAME
             )
+            self.account_key = None
             self.container_name: str = (
                 container_name
                 if container_name
                 else env_helper.AZURE_BLOB_CONTAINER_NAME
             )
-            credential = DefaultAzureCredential()
-            account_url = f"https://{self.account_name}.blob.core.windows.net/"
             self.blob_service_client = BlobServiceClient(
-                account_url=account_url, credential=credential
+                account_url=f"https://{self.account_name}.blob.core.windows.net/",
+                credential=DefaultAzureCredential(),
             )
             self.user_delegation_key = self.request_user_delegation_key(
                 blob_service_client=self.blob_service_client
             )
-            self.account_key = None
         else:
             self.account_name = (
                 account_name if account_name else env_helper.AZURE_BLOB_ACCOUNT_NAME
             )
             self.account_key = (
                 account_key if account_key else env_helper.AZURE_BLOB_ACCOUNT_KEY
             )
-            self.user_delegation_key = None
-            self.connect_str = f"DefaultEndpointsProtocol=https;AccountName={self.account_name};AccountKey={self.account_key};EndpointSuffix=core.windows.net"
+            self.connect_str = connection_string(self.account_name, self.account_key)
             self.container_name: str = (
                 container_name
                 if container_name
@@ -56,6 +79,7 @@ def __init__(
             self.blob_service_client: BlobServiceClient = (
                 BlobServiceClient.from_connection_string(self.connect_str)
             )
+            self.user_delegation_key = None
 
     def request_user_delegation_key(
         self, blob_service_client: BlobServiceClient

diff --git a/docs/LOCAL_DEPLOYMENT.md b/docs/LOCAL_DEPLOYMENT.md
@@ -19,7 +19,14 @@ The easiest way to run this accelerator is in a VS Code Dev Containers, which wi
 
 > NOTE: It may take up to an hour for the application to be fully deployed. If you see a "Python Developer" welcome screen or an error page, then wait a bit and refresh the page.
 
-**Notes:** the default auth type uses keys, if you want to switch to rbac, please run `azd env set AUTH_TYPE rbac`.
+> NOTE: The default auth type uses keys that are stored in the Azure Keyvault. If you want to use RBAC-based auth (more secure), please run before deploying:
+
+```bash
+azd env set AZURE_AUTH_TYPE rbac
+azd env set USE_KEY_VAULT false
+```
+
+Also please refer to the section on [setting up RBAC auth](#authenticate-using-rbac).
 
 ## Detailed Development Container setup instructions
 
@@ -142,8 +149,9 @@ Or use the [Azure Functions VS Code extension](https://marketplace.visualstudio.
 #### Debugging the batch processing functions locally
 Rename the file `local.settings.json.sample` in the `batch` folder to `local.settings.json` and update the `AzureWebJobsStorage` value with the storage account connection string.
 
-Execute the above [shell command](#L81) to run the function locally. You may need to stop the deployed function on the portal so that all requests are debugged locally.
+Copy the .env file from [previous section](#local-debugging) to the `batch` folder.
 
+Execute the above [shell command](#L81) to run the function locally. You may need to stop the deployed function on the portal so that all requests are debugged locally. To trigger the function, you can click on the corresponding URL that will be printed to the terminal.
 
 ## Environment variables
 

diff --git a/infra/app/function.bicep b/infra/app/function.bicep
@@ -113,15 +113,25 @@ module searchRoleFunction '../core/security/role.bicep' = if (authType == 'rbac'
 }
 
 // Storage Blob Data Contributor
-module storageRoleFunction '../core/security/role.bicep' = if (authType == 'rbac') {
-  name: 'storage-role-function'
+module storageBlobRoleFunction '../core/security/role.bicep' = if (authType == 'rbac') {
+  name: 'storage-blob-role-function'
   params: {
     principalId: function.outputs.identityPrincipalId
     roleDefinitionId: 'ba92f5b4-2d11-453d-a403-e96b0029c9fe'
     principalType: 'ServicePrincipal'
   }
 }
 
+// Storage Queue Data Contributor
+module storageQueueRoleFunction '../core/security/role.bicep' = if (authType == 'rbac') {
+  name: 'storage-queue-role-function'
+  params: {
+    principalId: function.outputs.identityPrincipalId
+    roleDefinitionId: '974c5e8b-45b9-4653-ba55-5f855dd0fb88'
+    principalType: 'ServicePrincipal'
+  }
+}
+
 module functionaccess '../core/security/keyvault-access.bicep' = if (useKeyVault) {
   name: 'function-keyvault-access'
   params: {

diff --git a/infra/main.json b/infra/main.json
@@ -5,7 +5,7 @@
     "_generator": {
       "name": "bicep",
       "version": "0.26.54.24096",
-      "templateHash": "14235243275465943137"
+      "templateHash": "4295550392354123736"
     }
   },
   "parameters": {
@@ -6551,7 +6551,7 @@
             "_generator": {
               "name": "bicep",
               "version": "0.26.54.24096",
-              "templateHash": "12912745910813883723"
+              "templateHash": "17960689149478457729"
             }
           },
           "parameters": {
@@ -7446,7 +7446,7 @@
               "condition": "[equals(parameters('authType'), 'rbac')]",
               "type": "Microsoft.Resources/deployments",
               "apiVersion": "2022-09-01",
-              "name": "storage-role-function",
+              "name": "storage-blob-role-function",
               "properties": {
                 "expressionEvaluationOptions": {
                   "scope": "inner"
@@ -7511,6 +7511,75 @@
                 "[resourceId('Microsoft.Resources/deployments', format('{0}-app-module', parameters('name')))]"
               ]
             },
+            {
+              "condition": "[equals(parameters('authType'), 'rbac')]",
+              "type": "Microsoft.Resources/deployments",
+              "apiVersion": "2022-09-01",
+              "name": "storage-queue-role-function",
+              "properties": {
+                "expressionEvaluationOptions": {
+                  "scope": "inner"
+                },
+                "mode": "Incremental",
+                "parameters": {
+                  "principalId": {
+                    "value": "[reference(resourceId('Microsoft.Resources/deployments', format('{0}-app-module', parameters('name'))), '2022-09-01').outputs.identityPrincipalId.value]"
+                  },
+                  "roleDefinitionId": {
+                    "value": "974c5e8b-45b9-4653-ba55-5f855dd0fb88"
+                  },
+                  "principalType": {
+                    "value": "ServicePrincipal"
+                  }
+                },
+                "template": {
+                  "$schema": "https://schema.management.azure.com/schemas/2019-04-01/deploymentTemplate.json#",
+                  "contentVersion": "1.0.0.0",
+                  "metadata": {
+                    "_generator": {
+                      "name": "bicep",
+                      "version": "0.26.54.24096",
+                      "templateHash": "5795525499710207356"
+                    },
+                    "description": "Creates a role assignment for a service principal."
+                  },
+                  "parameters": {
+                    "principalId": {
+                      "type": "string"
+                    },
+                    "principalType": {
+                      "type": "string",
+                      "defaultValue": "ServicePrincipal",
+                      "allowedValues": [
+                        "Device",
+                        "ForeignGroup",
+                        "Group",
+                        "ServicePrincipal",
+                        "User"
+                      ]
+                    },
+                    "roleDefinitionId": {
+                      "type": "string"
+                    }
+                  },
+                  "resources": [
+                    {
+                      "type": "Microsoft.Authorization/roleAssignments",
+                      "apiVersion": "2022-04-01",
+                      "name": "[guid(subscription().id, resourceGroup().id, parameters('principalId'), parameters('roleDefinitionId'))]",
+                      "properties": {
+                        "principalId": "[parameters('principalId')]",
+                        "principalType": "[parameters('principalType')]",
+                        "roleDefinitionId": "[resourceId('Microsoft.Authorization/roleDefinitions', parameters('roleDefinitionId'))]"
+                      }
+                    }
+                  ]
+                }
+              },
+              "dependsOn": [
+                "[resourceId('Microsoft.Resources/deployments', format('{0}-app-module', parameters('name')))]"
+              ]
+            },
             {
               "condition": "[parameters('useKeyVault')]",
               "type": "Microsoft.Resources/deployments",
@@ -7694,7 +7763,7 @@
             "_generator": {
               "name": "bicep",
               "version": "0.26.54.24096",
-              "templateHash": "12912745910813883723"
+              "templateHash": "17960689149478457729"
             }
           },
           "parameters": {
@@ -8589,7 +8658,7 @@
               "condition": "[equals(parameters('authType'), 'rbac')]",
               "type": "Microsoft.Resources/deployments",
               "apiVersion": "2022-09-01",
-              "name": "storage-role-function",
+              "name": "storage-blob-role-function",
               "properties": {
                 "expressionEvaluationOptions": {
                   "scope": "inner"
@@ -8654,6 +8723,75 @@
                 "[resourceId('Microsoft.Resources/deployments', format('{0}-app-module', parameters('name')))]"
               ]
             },
+            {
+              "condition": "[equals(parameters('authType'), 'rbac')]",
+              "type": "Microsoft.Resources/deployments",
+              "apiVersion": "2022-09-01",
+              "name": "storage-queue-role-function",
+              "properties": {
+                "expressionEvaluationOptions": {
+                  "scope": "inner"
+                },
+                "mode": "Incremental",
+                "parameters": {
+                  "principalId": {
+                    "value": "[reference(resourceId('Microsoft.Resources/deployments', format('{0}-app-module', parameters('name'))), '2022-09-01').outputs.identityPrincipalId.value]"
+                  },
+                  "roleDefinitionId": {
+                    "value": "974c5e8b-45b9-4653-ba55-5f855dd0fb88"
+                  },
+                  "principalType": {
+                    "value": "ServicePrincipal"
+                  }
+                },
+                "template": {
+                  "$schema": "https://schema.management.azure.com/schemas/2019-04-01/deploymentTemplate.json#",
+                  "contentVersion": "1.0.0.0",
+                  "metadata": {
+                    "_generator": {
+                      "name": "bicep",
+                      "version": "0.26.54.24096",
+                      "templateHash": "5795525499710207356"
+                    },
+                    "description": "Creates a role assignment for a service principal."
+                  },
+                  "parameters": {
+                    "principalId": {
+                      "type": "string"
+                    },
+                    "principalType": {
+                      "type": "string",
+                      "defaultValue": "ServicePrincipal",
+                      "allowedValues": [
+                        "Device",
+                        "ForeignGroup",
+                        "Group",
+                        "ServicePrincipal",
+                        "User"
+                      ]
+                    },
+                    "roleDefinitionId": {
+                      "type": "string"
+                    }
+                  },
+                  "resources": [
+                    {
+                      "type": "Microsoft.Authorization/roleAssignments",
+                      "apiVersion": "2022-04-01",
+                      "name": "[guid(subscription().id, resourceGroup().id, parameters('principalId'), parameters('roleDefinitionId'))]",
+                      "properties": {
+                        "principalId": "[parameters('principalId')]",
+                        "principalType": "[parameters('principalType')]",
+                        "roleDefinitionId": "[resourceId('Microsoft.Authorization/roleDefinitions', parameters('roleDefinitionId'))]"
+                      }
+                    }
+                  ]
+                }
+              },
+              "dependsOn": [
+                "[resourceId('Microsoft.Resources/deployments', format('{0}-app-module', parameters('name')))]"
+              ]
+            },
             {
               "condition": "[parameters('useKeyVault')]",
               "type": "Microsoft.Resources/deployments",