Enable User assigned managed identity for scalesets (#219)

src/api-service/__app__/onefuzzlib/agent_authorization.py

@@ -14,6 +14,7 @@ from onefuzztypes.enums import ErrorCode
 from onefuzztypes.models import Error
 from pydantic import BaseModel
 
+from .azure.creds import get_scaleset_principal_id
 from .pools import Scaleset
 from .request import not_ok
 
@@ -54,6 +55,12 @@ def try_get_token_auth_header(request: func.HttpRequest) -> Union[Error, TokenDa
 
 @cached(ttl=60)
 def is_authorized(token_data: TokenData) -> bool:
+    # verify object_id against the user assigned managed identity
+    if get_scaleset_principal_id() == token_data.object_id:
+        return True
+
+    # backward compatibility case for scalesets deployed before the migration
+    # to user assigned managed id
     scalesets = Scaleset.get_by_object_id(token_data.object_id)
     return len(scalesets) > 0
 

src/api-service/__app__/onefuzzlib/azure/creds.py

@@ -6,6 +6,7 @@
 import logging
 import os
 from typing import Any, List, Optional, Tuple
+from uuid import UUID
 
 from azure.cli.core import CLIError
 from azure.common.client_factory import get_client_from_cli_profile
@@ -123,3 +124,24 @@ def is_member_of(group_id: str, member_id: str) -> bool:
             CheckGroupMembershipParameters(group_id=group_id, member_id=member_id)
         ).value
     )
+
+
+@cached
+def get_scaleset_identity_resource_path() -> str:
+    scaleset_id_name = "%s-scalesetid" % get_instance_name()
+    resource_group_path = "/subscriptions/%s/resourceGroups/%s/providers" % (
+        get_subscription(),
+        get_base_resource_group(),
+    )
+    return "%s/Microsoft.ManagedIdentity/userAssignedIdentities/%s" % (
+        resource_group_path,
+        scaleset_id_name,
+    )
+
+
+@cached
+def get_scaleset_principal_id() -> UUID:
+    api_version = "2018-11-30"  # matches the apiversion in the deplyoment template
+    client = mgmt_client_factory(ResourceManagementClient)
+    uid = client.resources.get_by_id(get_scaleset_identity_resource_path(), api_version)
+    return UUID(uid.properties["principalId"])

src/api-service/__app__/onefuzzlib/azure/vmss.py

@@ -16,7 +16,11 @@ from onefuzztypes.enums import OS, ErrorCode
 from onefuzztypes.models import Error
 from onefuzztypes.primitives import Region
 
-from .creds import get_base_resource_group, mgmt_client_factory
+from .creds import (
+    get_base_resource_group,
+    get_scaleset_identity_resource_path,
+    mgmt_client_factory,
+)
 from .image import get_os
 
 
@@ -234,7 +238,10 @@ def create_vmss(
         "do_not_run_extensions_on_overprovisioned_vms": True,
         "upgrade_policy": {"mode": "Manual"},
         "sku": sku,
-        "identity": {"type": "SystemAssigned"},
+        "identity": {
+            "type": "userAssigned",
+            "userAssignedIdentities": {get_scaleset_identity_resource_path(): {}},
+        },
         "virtual_machine_profile": {
             "priority": "Regular",
             "storage_profile": {"image_reference": image_ref},