support multiple corpus accounts

src/api-service/__app__/agent_registration/__init__.py

@@ -12,9 +12,9 @@
 from onefuzztypes.requests import AgentRegistrationGet, AgentRegistrationPost
 from onefuzztypes.responses import AgentRegistration
 
-from ..onefuzzlib.azure.containers import StorageType
 from ..onefuzzlib.azure.creds import get_instance_url
 from ..onefuzzlib.azure.queue import get_queue_sas
+from ..onefuzzlib.azure.storage import StorageType
 from ..onefuzzlib.endpoint_authorization import call_if_agent
 from ..onefuzzlib.pools import Node, NodeMessage, NodeTasks, Pool
 from ..onefuzzlib.request import not_ok, ok, parse_uri

src/api-service/__app__/containers/__init__.py

@@ -13,13 +13,13 @@
 from onefuzztypes.responses import BoolResult, ContainerInfo, ContainerInfoBase
 
 from ..onefuzzlib.azure.containers import (
-    StorageType,
     create_container,
     delete_container,
     get_container_metadata,
     get_container_sas_url,
     get_containers,
 )
+from ..onefuzzlib.azure.storage import StorageType
 from ..onefuzzlib.endpoint_authorization import call_if_user
 from ..onefuzzlib.request import not_ok, ok, parse_request
 

src/api-service/__app__/download/__init__.py

@@ -8,11 +8,11 @@
 from onefuzztypes.models import Error, FileEntry
 
 from ..onefuzzlib.azure.containers import (
-    StorageType,
     blob_exists,
     container_exists,
     get_file_sas_url,
 )
+from ..onefuzzlib.azure.storage import StorageType
 from ..onefuzzlib.endpoint_authorization import call_if_user
 from ..onefuzzlib.request import not_ok, parse_uri, redirect
 

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

@@ -4,81 +4,134 @@
 # Licensed under the MIT License.
 
 import datetime
+import logging
 import os
 import urllib.parse
-from enum import Enum
-from typing import Any, Dict, Optional, Union, cast
+from typing import Dict, Optional, Union, cast
 
 from azure.common import AzureHttpError, AzureMissingResourceHttpError
-from azure.storage.blob import BlobPermissions, ContainerPermissions
+from azure.storage.blob import BlobPermissions, BlockBlobService, ContainerPermissions
 from memoization import cached
+from onefuzztypes.primitives import Container
 
-from .creds import get_blob_service, get_func_storage, get_fuzz_storage
+from .storage import (
+    StorageType,
+    choose_account,
+    get_accounts,
+    get_storage_account_name_key,
+)
 
 
-class StorageType(Enum):
-    corpus = "corpus"
-    config = "config"
+@cached
+def get_blob_service(account_id: str) -> BlockBlobService:
+    logging.debug("getting blob container (account_id: %s)", account_id)
+    account_name, account_key = get_storage_account_name_key(account_id)
+    service = BlockBlobService(account_name=account_name, account_key=account_key)
+    return service
 
 
-def get_account_id_by_type(storage_type: StorageType) -> str:
-    if storage_type == StorageType.corpus:
-        account_id = get_fuzz_storage()
-    elif storage_type == StorageType.config:
-        account_id = get_func_storage()
-    else:
-        raise NotImplementedError
-    return account_id
-
-
-@cached(ttl=5)
-def get_blob_service_by_type(storage_type: StorageType) -> Any:
-    account_id = get_account_id_by_type(storage_type)
-    return get_blob_service(account_id)
+def get_service_by_container(
+    container: Container, storage_type: StorageType
+) -> Optional[BlockBlobService]:
+    account = get_account_by_container(container, storage_type)
+    if account is None:
+        return None
+    service = get_blob_service(account)
+    return service
 
 
-@cached(ttl=5)
-def container_exists(name: str, storage_type: StorageType) -> bool:
+def container_exists_on_account(container: Container, account_id: str) -> bool:
     try:
-        get_blob_service_by_type(storage_type).get_container_properties(name)
+        get_blob_service(account_id).get_container_properties(container)
         return True
     except AzureHttpError:
         return False
 
 
-def get_containers(storage_type: StorageType) -> Dict[str, Dict[str, str]]:
-    return {
-        x.name: x.metadata
-        for x in get_blob_service_by_type(storage_type).list_containers(
-            include_metadata=True
-        )
-        if not x.name.startswith("$")
-    }
-
-
-def get_container_metadata(
-    name: str, storage_type: StorageType
-) -> Optional[Dict[str, str]]:
+def container_metadata(container: Container, account: str) -> Optional[Dict[str, str]]:
     try:
-        result = get_blob_service_by_type(storage_type).get_container_metadata(name)
+        result = get_blob_service(account).get_container_metadata(container)
         return cast(Dict[str, str], result)
     except AzureHttpError:
         pass
     return None
 
 
-def create_container(
-    name: str, storage_type: StorageType, metadata: Optional[Dict[str, str]]
+def get_account_by_container(
+    container: Container, storage_type: StorageType
 ) -> Optional[str]:
-    try:
-        get_blob_service_by_type(storage_type).create_container(name, metadata=metadata)
-    except AzureHttpError:
-        # azure storage already logs errors
+    accounts = get_accounts(storage_type)
+
+    # check secondary accounts first by searching in reverse.
+    #
+    # By implementation, the primary account is specified first, followed by
+    # any secondary accounts.
+    #
+    # Secondary accounts, if they exist, are preferred for containers and have
+    # increased IOP rates, this should be a slight optimization
+    for account in reversed(accounts):
+        if container_exists_on_account(container, account):
+            return account
+    return None
+
+
+def container_exists(container: Container, storage_type: StorageType) -> bool:
+    return get_account_by_container(container, storage_type) is not None
+
+
+def get_containers(storage_type: StorageType) -> Dict[str, Dict[str, str]]:
+    containers: Dict[str, Dict[str, str]] = {}
+
+    for account_id in get_accounts(storage_type):
+        containers.update(
+            {
+                x.name: x.metadata
+                for x in get_blob_service(account_id).list_containers(
+                    include_metadata=True
+                )
+            }
+        )
+
+    return containers
+
+
+def get_container_metadata(
+    container: Container, storage_type: StorageType
+) -> Optional[Dict[str, str]]:
+    account = get_account_by_container(container, storage_type)
+    if account is None:
         return None
 
-    return get_container_sas_url(
-        name,
-        storage_type,
+    return container_metadata(container, account)
+
+
+def create_container(
+    container: Container,
+    storage_type: StorageType,
+    metadata: Optional[Dict[str, str]],
+) -> Optional[str]:
+    service = get_service_by_container(container, storage_type)
+    if service is None:
+        account = choose_account(storage_type)
+        service = get_blob_service(account)
+        try:
+            service.create_container(container, metadata=metadata)
+        except AzureHttpError as err:
+            logging.error(
+                (
+                    "unable to create container.  account: %s "
+                    "container: %s metadata: %s - %s"
+                ),
+                account,
+                container,
+                metadata,
+                err,
+            )
+            return None
+
+    return get_container_sas_url_service(
+        container,
+        service,
         read=True,
         add=True,
         create=True,
@@ -88,17 +141,19 @@ def create_container(
     )
 
 
-def delete_container(name: str, storage_type: StorageType) -> bool:
-    try:
-        return bool(get_blob_service_by_type(storage_type).delete_container(name))
-    except AzureHttpError:
-        # azure storage already logs errors
-        return False
+def delete_container(container: Container, storage_type: StorageType) -> bool:
+    accounts = get_accounts(storage_type)
+    for account in accounts:
+        service = get_blob_service(account)
+        if bool(service.delete_container(container)):
+            return True
 
+    return False
 
-def get_container_sas_url(
-    container: str,
-    storage_type: StorageType,
+
+def get_container_sas_url_service(
+    container: Container,
+    service: BlockBlobService,
     *,
     read: bool = False,
     add: bool = False,
@@ -107,7 +162,6 @@ def get_container_sas_url(
     delete: bool = False,
     list: bool = False,
 ) -> str:
-    service = get_blob_service_by_type(storage_type)
     expiry = datetime.datetime.utcnow() + datetime.timedelta(days=30)
     permission = ContainerPermissions(read, add, create, write, delete, list)
 
@@ -120,8 +174,35 @@ def get_container_sas_url(
     return str(url)
 
 
+def get_container_sas_url(
+    container: Container,
+    storage_type: StorageType,
+    *,
+    read: bool = False,
+    add: bool = False,
+    create: bool = False,
+    write: bool = False,
+    delete: bool = False,
+    list: bool = False,
+) -> str:
+    service = get_service_by_container(container, storage_type)
+    if not service:
+        raise Exception("unable to create container sas for missing container")
+
+    return get_container_sas_url_service(
+        container,
+        service,
+        read=read,
+        add=add,
+        create=create,
+        write=write,
+        delete=delete,
+        list=list,
+    )
+
+
 def get_file_sas_url(
-    container: str,
+    container: Container,
     name: str,
     storage_type: StorageType,
     *,
@@ -135,7 +216,10 @@ def get_file_sas_url(
     hours: int = 0,
     minutes: int = 0,
 ) -> str:
-    service = get_blob_service_by_type(storage_type)
+    service = get_service_by_container(container, storage_type)
+    if not service:
+        raise Exception("unable to find container: %s - %s" % (container, storage_type))
+
     expiry = datetime.datetime.utcnow() + datetime.timedelta(
         days=days, hours=hours, minutes=minutes
     )
@@ -150,44 +234,60 @@ def get_file_sas_url(
 
 
 def save_blob(
-    container: str, name: str, data: Union[str, bytes], storage_type: StorageType
+    container: Container,
+    name: str,
+    data: Union[str, bytes],
+    storage_type: StorageType,
 ) -> None:
-    service = get_blob_service_by_type(storage_type)
-    service.create_container(container)
+    service = get_service_by_container(container, storage_type)
+    if not service:
+        raise Exception("unable to find container: %s - %s" % (container, storage_type))
+
     if isinstance(data, str):
         service.create_blob_from_text(container, name, data)
     elif isinstance(data, bytes):
         service.create_blob_from_bytes(container, name, data)
 
 
-def get_blob(container: str, name: str, storage_type: StorageType) -> Optional[bytes]:
-    service = get_blob_service_by_type(storage_type)
+def get_blob(
+    container: Container, name: str, storage_type: StorageType
+) -> Optional[bytes]:
+    service = get_service_by_container(container, storage_type)
+    if not service:
+        return None
+
     try:
         blob = service.get_blob_to_bytes(container, name).content
         return cast(bytes, blob)
     except AzureMissingResourceHttpError:
         return None
 
 
-def blob_exists(container: str, name: str, storage_type: StorageType) -> bool:
-    service = get_blob_service_by_type(storage_type)
+def blob_exists(container: Container, name: str, storage_type: StorageType) -> bool:
+    service = get_service_by_container(container, storage_type)
+    if not service:
+        return False
+
     try:
         service.get_blob_properties(container, name)
         return True
     except AzureMissingResourceHttpError:
         return False
 
 
-def delete_blob(container: str, name: str, storage_type: StorageType) -> bool:
-    service = get_blob_service_by_type(storage_type)
+def delete_blob(container: Container, name: str, storage_type: StorageType) -> bool:
+    service = get_service_by_container(container, storage_type)
+    if not service:
+        return False
+
     try:
         service.delete_blob(container, name)
         return True
     except AzureMissingResourceHttpError:
         return False
 
 
-def auth_download_url(container: str, filename: str) -> str:
+def auth_download_url(container: Container, filename: str) -> str:
     instance = os.environ["ONEFUZZ_INSTANCE"]
     return "%s/api/download?%s" % (
         instance,