Storing the user assigned managed identity in the scaleset table

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

@@ -55,14 +55,14 @@ 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
+    if len(scalesets) > 0:
+        return True
+
+    # verify object_id against the user assigned managed identity
+    return get_scaleset_principal_id() == token_data.object_id
 
 
 def verify_token(

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

@@ -5,7 +5,7 @@
 
 import datetime
 import logging
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 from uuid import UUID, uuid4
 
 from onefuzztypes.enums import (
@@ -712,14 +712,30 @@ def setup(self) -> None:
                 logging.info("creating scaleset: %s", self.scaleset_id)
         elif vmss.provisioning_state == "Creating":
             logging.info("Waiting on scaleset creation: %s", self.scaleset_id)
-            if vmss.identity and vmss.identity.principal_id:
-                self.client_object_id = vmss.identity.principal_id
+            self.set_identity(vmss)
         else:
             logging.info("scaleset running: %s", self.scaleset_id)
+            self.set_identity(vmss)
             self.state = ScalesetState.running
-            self.client_object_id = vmss.identity.principal_id
         self.save()
 
+    def set_identity(self, vmss: Any) -> None:
+        if (
+            vmss.identity
+            and vmss.user_assigned_identities
+            and (len(vmss.identity.user_assigned_identities) != 1)
+        ):
+            self.error = Error(
+                code=ErrorCode.VM_CREATE_FAILED,
+                errors=[
+                    "The scaleset is expected to have exactly 1 user assigned identity"
+                ],
+            )
+            self.state = ScalesetState.creation_failed
+
+        self.client_object_id = list(vmss.identity.user_assigned_identities.values())[0]
+        return None
+
     # result = 'did I modify the scaleset in azure'
     def cleanup_nodes(self) -> bool:
         if self.state == ScalesetState.halt: