Use Storage Account types, rather than account_id (#320)

We need to move to supporting data sharding.

One of the steps towards that is stop passing around `account_id`, rather we need to specify the type of storage we need.

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

@@ -13,7 +13,8 @@
 from onefuzztypes.responses import AgentRegistration
 
 from ..onefuzzlib.agent_authorization import call_if_agent
-from ..onefuzzlib.azure.creds import get_fuzz_storage, get_instance_url
+from ..onefuzzlib.azure.containers import StorageType
+from ..onefuzzlib.azure.creds import get_instance_url
 from ..onefuzzlib.azure.queue import get_queue_sas
 from ..onefuzzlib.pools import Node, NodeMessage, NodeTasks, Pool
 from ..onefuzzlib.request import not_ok, ok, parse_uri
@@ -25,7 +26,7 @@ def create_registration_response(machine_id: UUID, pool: Pool) -> func.HttpRespo
     commands_url = "%s/api/agents/commands" % base_address
     work_queue = get_queue_sas(
         pool.get_pool_queue(),
-        account_id=get_fuzz_storage(),
+        StorageType.corpus,
         read=True,
         update=True,
         process=True,

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

@@ -13,6 +13,7 @@
 from onefuzztypes.responses import BoolResult, ContainerInfo, ContainerInfoBase
 
 from ..onefuzzlib.azure.containers import (
+    StorageType,
     create_container,
     delete_container,
     get_container_metadata,
@@ -30,7 +31,7 @@ def get(req: func.HttpRequest) -> func.HttpResponse:
     if isinstance(request, Error):
         return not_ok(request, context="container get")
     if request is not None:
-        metadata = get_container_metadata(request.name)
+        metadata = get_container_metadata(request.name, StorageType.corpus)
         if metadata is None:
             return not_ok(
                 Error(code=ErrorCode.INVALID_REQUEST, errors=["invalid container"]),
@@ -41,6 +42,7 @@ def get(req: func.HttpRequest) -> func.HttpResponse:
             name=request.name,
             sas_url=get_container_sas_url(
                 request.name,
+                StorageType.corpus,
                 read=True,
                 write=True,
                 create=True,
@@ -51,7 +53,7 @@ def get(req: func.HttpRequest) -> func.HttpResponse:
         )
         return ok(info)
 
-    containers = get_containers()
+    containers = get_containers(StorageType.corpus)
 
     container_info = []
     for name in containers:
@@ -66,7 +68,7 @@ def post(req: func.HttpRequest) -> func.HttpResponse:
         return not_ok(request, context="container create")
 
     logging.info("container - creating %s", request.name)
-    sas = create_container(request.name, metadata=request.metadata)
+    sas = create_container(request.name, StorageType.corpus, metadata=request.metadata)
     if sas:
         return ok(
             ContainerInfo(name=request.name, sas_url=sas, metadata=request.metadata)
@@ -83,7 +85,7 @@ def delete(req: func.HttpRequest) -> func.HttpResponse:
         return not_ok(request, context="container delete")
 
     logging.info("container - deleting %s", request.name)
-    return ok(BoolResult(result=delete_container(request.name)))
+    return ok(BoolResult(result=delete_container(request.name, StorageType.corpus)))
 
 
 def main(req: func.HttpRequest) -> func.HttpResponse:

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

@@ -8,6 +8,7 @@
 from onefuzztypes.models import Error, FileEntry
 
 from ..onefuzzlib.azure.containers import (
+    StorageType,
     blob_exists,
     container_exists,
     get_file_sas_url,
@@ -20,21 +21,26 @@ def get(req: func.HttpRequest) -> func.HttpResponse:
     if isinstance(request, Error):
         return not_ok(request, context="download")
 
-    if not container_exists(request.container):
+    if not container_exists(request.container, StorageType.corpus):
         return not_ok(
             Error(code=ErrorCode.INVALID_REQUEST, errors=["invalid container"]),
             context=request.container,
         )
 
-    if not blob_exists(request.container, request.filename):
+    if not blob_exists(request.container, request.filename, StorageType.corpus):
         return not_ok(
             Error(code=ErrorCode.INVALID_REQUEST, errors=["invalid filename"]),
             context=request.filename,
         )
 
     return redirect(
         get_file_sas_url(
-            request.container, request.filename, read=True, days=0, minutes=5
+            request.container,
+            request.filename,
+            StorageType.corpus,
+            read=True,
+            days=0,
+            minutes=5,
         )
     )
 

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

@@ -6,76 +6,108 @@
 import datetime
 import os
 import urllib.parse
-from typing import Dict, Optional, Union, cast
+from enum import Enum
+from typing import Any, Dict, Optional, Union, cast
 
 from azure.common import AzureHttpError, AzureMissingResourceHttpError
 from azure.storage.blob import BlobPermissions, ContainerPermissions
 from memoization import cached
 
-from .creds import get_blob_service
+from .creds import get_blob_service, get_func_storage, get_fuzz_storage
+
+
+class StorageType(Enum):
+    corpus = "corpus"
+    config = "config"
+
+
+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)
 
 
 @cached(ttl=5)
-def container_exists(name: str, account_id: Optional[str] = None) -> bool:
+def container_exists(name: str, storage_type: StorageType) -> bool:
     try:
-        get_blob_service(account_id).get_container_properties(name)
+        get_blob_service_by_type(storage_type).get_container_properties(name)
         return True
     except AzureHttpError:
         return False
 
 
-def get_containers(account_id: Optional[str] = None) -> Dict[str, Dict[str, str]]:
+def get_containers(storage_type: StorageType) -> Dict[str, Dict[str, str]]:
     return {
         x.name: x.metadata
-        for x in get_blob_service(account_id).list_containers(include_metadata=True)
+        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, account_id: Optional[str] = None
+    name: str, storage_type: StorageType
 ) -> Optional[Dict[str, str]]:
     try:
-        result = get_blob_service(account_id).get_container_metadata(name)
+        result = get_blob_service_by_type(storage_type).get_container_metadata(name)
         return cast(Dict[str, str], result)
     except AzureHttpError:
         pass
     return None
 
 
 def create_container(
-    name: str, metadata: Optional[Dict[str, str]], account_id: Optional[str] = None
+    name: str, storage_type: StorageType, metadata: Optional[Dict[str, str]]
 ) -> Optional[str]:
     try:
-        get_blob_service(account_id).create_container(name, metadata=metadata)
+        get_blob_service_by_type(storage_type).create_container(name, metadata=metadata)
     except AzureHttpError:
         # azure storage already logs errors
         return None
 
     return get_container_sas_url(
-        name, read=True, add=True, create=True, write=True, delete=True, list=True
+        name,
+        storage_type,
+        read=True,
+        add=True,
+        create=True,
+        write=True,
+        delete=True,
+        list=True,
     )
 
 
-def delete_container(name: str, account_id: Optional[str] = None) -> bool:
+def delete_container(name: str, storage_type: StorageType) -> bool:
     try:
-        return bool(get_blob_service(account_id).delete_container(name))
+        return bool(get_blob_service_by_type(storage_type).delete_container(name))
     except AzureHttpError:
         # azure storage already logs errors
         return False
 
 
 def get_container_sas_url(
     container: str,
-    account_id: Optional[str] = None,
+    storage_type: StorageType,
+    *,
     read: bool = False,
     add: bool = False,
     create: bool = False,
     write: bool = False,
     delete: bool = False,
     list: bool = False,
 ) -> str:
-    service = get_blob_service(account_id)
+    service = get_blob_service_by_type(storage_type)
     expiry = datetime.datetime.utcnow() + datetime.timedelta(days=30)
     permission = ContainerPermissions(read, add, create, write, delete, list)
 
@@ -91,7 +123,8 @@ def get_container_sas_url(
 def get_file_sas_url(
     container: str,
     name: str,
-    account_id: Optional[str] = None,
+    storage_type: StorageType,
+    *,
     read: bool = False,
     add: bool = False,
     create: bool = False,
@@ -102,7 +135,7 @@ def get_file_sas_url(
     hours: int = 0,
     minutes: int = 0,
 ) -> str:
-    service = get_blob_service(account_id)
+    service = get_blob_service_by_type(storage_type)
     expiry = datetime.datetime.utcnow() + datetime.timedelta(
         days=days, hours=hours, minutes=minutes
     )
@@ -117,38 +150,36 @@ def get_file_sas_url(
 
 
 def save_blob(
-    container: str, name: str, data: Union[str, bytes], account_id: Optional[str] = None
+    container: str, name: str, data: Union[str, bytes], storage_type: StorageType
 ) -> None:
-    service = get_blob_service(account_id)
+    service = get_blob_service_by_type(storage_type)
     service.create_container(container)
     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, account_id: Optional[str] = None
-) -> Optional[bytes]:
-    service = get_blob_service(account_id)
+def get_blob(container: str, name: str, storage_type: StorageType) -> Optional[bytes]:
+    service = get_blob_service_by_type(storage_type)
     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, account_id: Optional[str] = None) -> bool:
-    service = get_blob_service(account_id)
+def blob_exists(container: str, name: str, storage_type: StorageType) -> bool:
+    service = get_blob_service_by_type(storage_type)
     try:
         service.get_blob_properties(container, name)
         return True
     except AzureMissingResourceHttpError:
         return False
 
 
-def delete_blob(container: str, name: str, account_id: Optional[str] = None) -> bool:
-    service = get_blob_service(account_id)
+def delete_blob(container: str, name: str, storage_type: StorageType) -> bool:
+    service = get_blob_service_by_type(storage_type)
     try:
         service.delete_blob(container, name)
         return True

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

@@ -87,12 +87,12 @@ def get_insights_appid() -> str:
     return os.environ["APPINSIGHTS_APPID"]
 
 
-@cached
+# @cached
 def get_fuzz_storage() -> str:
     return os.environ["ONEFUZZ_DATA_STORAGE"]
 
 
-@cached
+# @cached
 def get_func_storage() -> str:
     return os.environ["ONEFUZZ_FUNC_STORAGE"]
 
@@ -109,9 +109,9 @@ def get_instance_url() -> str:
 
 @cached
 def get_instance_id() -> UUID:
-    from .containers import get_blob
+    from .containers import StorageType, get_blob
 
-    blob = get_blob("base-config", "instance_id", account_id=get_func_storage())
+    blob = get_blob("base-config", "instance_id", StorageType.config)
     if blob is None:
         raise Exception("missing instance_id")
     return UUID(blob.decode())