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.
This commit is contained in:
bmc-msft
2020-11-18 09:06:14 -05:00
committed by GitHub
parent 52eca33237
commit e47e89609a
18 changed files with 205 additions and 139 deletions

View File

@ -11,13 +11,13 @@ from uuid import UUID
from onefuzztypes.enums import Compare, ContainerPermission, ContainerType, TaskFeature
from onefuzztypes.models import TaskConfig, TaskDefinition, TaskUnitConfig
from ..azure.containers import blob_exists, container_exists, get_container_sas_url
from ..azure.creds import (
get_func_storage,
get_fuzz_storage,
get_instance_id,
get_instance_url,
from ..azure.containers import (
StorageType,
blob_exists,
container_exists,
get_container_sas_url,
)
from ..azure.creds import get_instance_id, get_instance_url
from ..azure.queue import get_queue_sas
from .defs import TASK_DEFINITIONS
@ -68,7 +68,7 @@ def check_containers(definition: TaskDefinition, config: TaskConfig) -> None:
containers: Dict[ContainerType, List[str]] = {}
for container in config.containers:
if container.name not in checked:
if not container_exists(container.name):
if not container_exists(container.name, StorageType.corpus):
raise TaskConfigError("missing container: %s" % container.name)
checked.add(container.name)
@ -137,7 +137,7 @@ def check_config(config: TaskConfig) -> None:
if TaskFeature.target_exe in definition.features:
container = [x for x in config.containers if x.type == ContainerType.setup][0]
if not blob_exists(container.name, config.task.target_exe):
if not blob_exists(container.name, config.task.target_exe, StorageType.corpus):
err = "target_exe `%s` does not exist in the setup container `%s`" % (
config.task.target_exe,
container.name,
@ -153,7 +153,7 @@ def check_config(config: TaskConfig) -> None:
for tool_path in tools_paths:
if config.task.generator_exe.startswith(tool_path):
generator = config.task.generator_exe.replace(tool_path, "")
if not blob_exists(container.name, generator):
if not blob_exists(container.name, generator, StorageType.corpus):
err = (
"generator_exe `%s` does not exist in the tools container `%s`"
% (
@ -188,7 +188,7 @@ def build_task_config(
telemetry_key=os.environ.get("ONEFUZZ_TELEMETRY"),
heartbeat_queue=get_queue_sas(
"task-heartbeat",
account_id=get_func_storage(),
StorageType.config,
add=True,
),
back_channel_address="https://%s/api/back_channel" % (get_instance_url()),
@ -198,11 +198,11 @@ def build_task_config(
if definition.monitor_queue:
config.input_queue = get_queue_sas(
task_id,
StorageType.corpus,
add=True,
read=True,
update=True,
process=True,
account_id=get_fuzz_storage(),
)
for container_def in definition.containers:
@ -219,6 +219,7 @@ def build_task_config(
"path": "_".join(["task", container_def.type.name, str(i)]),
"url": get_container_sas_url(
container.name,
StorageType.corpus,
read=ContainerPermission.Read in container_def.permissions,
write=ContainerPermission.Write in container_def.permissions,
add=ContainerPermission.Add in container_def.permissions,

View File

@ -18,7 +18,7 @@ from onefuzztypes.webhooks import (
WebhookEventTaskStopped,
)
from ..azure.creds import get_fuzz_storage
from ..azure.containers import StorageType
from ..azure.image import get_os
from ..azure.queue import create_queue, delete_queue
from ..orm import MappingIntStrAny, ORMMixin, QueryFilter
@ -123,7 +123,7 @@ class Task(BASE_TASK, ORMMixin):
}
def init(self) -> None:
create_queue(self.task_id, account_id=get_fuzz_storage())
create_queue(self.task_id, StorageType.corpus)
self.state = TaskState.waiting
self.save()
@ -132,7 +132,7 @@ class Task(BASE_TASK, ORMMixin):
logging.info("stopping task: %s:%s", self.job_id, self.task_id)
ProxyForward.remove_forward(self.task_id)
delete_queue(str(self.task_id), account_id=get_fuzz_storage())
delete_queue(str(self.task_id), StorageType.corpus)
Node.stop_task(self.task_id)
self.state = TaskState.stopped
self.save()

View File

@ -10,8 +10,12 @@ from uuid import UUID
from onefuzztypes.enums import OS, PoolState, TaskState
from onefuzztypes.models import WorkSet, WorkUnit
from ..azure.containers import blob_exists, get_container_sas_url, save_blob
from ..azure.creds import get_func_storage
from ..azure.containers import (
StorageType,
blob_exists,
get_container_sas_url,
save_blob,
)
from ..pools import Pool
from .config import build_task_config, get_setup_container
from .main import Task
@ -60,20 +64,26 @@ def schedule_tasks() -> None:
agent_config = build_task_config(task.job_id, task.task_id, task.config)
setup_container = get_setup_container(task.config)
setup_url = get_container_sas_url(setup_container, read=True, list=True)
setup_url = get_container_sas_url(
setup_container, StorageType.corpus, read=True, list=True
)
setup_script = None
if task.os == OS.windows and blob_exists(setup_container, "setup.ps1"):
if task.os == OS.windows and blob_exists(
setup_container, "setup.ps1", StorageType.corpus
):
setup_script = "setup.ps1"
if task.os == OS.linux and blob_exists(setup_container, "setup.sh"):
if task.os == OS.linux and blob_exists(
setup_container, "setup.sh", StorageType.corpus
):
setup_script = "setup.sh"
save_blob(
"task-configs",
"%s/config.json" % task.task_id,
agent_config.json(exclude_none=True),
account_id=get_func_storage(),
StorageType.config,
)
reboot = False
count = 1