expose the ability manually override node reset (#201)

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

@@ -8,7 +8,13 @@ from typing import Optional, cast
 from uuid import UUID
 
 import azure.functions as func
-from onefuzztypes.enums import ErrorCode, NodeState, NodeTaskState, TaskState
+from onefuzztypes.enums import (
+    ErrorCode,
+    NodeState,
+    NodeTaskState,
+    TaskDebugFlag,
+    TaskState,
+)
 from onefuzztypes.models import (
     Error,
     NodeDoneEventData,
@@ -173,8 +179,21 @@ def on_worker_event(machine_id: UUID, event: WorkerEvent) -> None:
                     ],
                 )
             )
+            if task.config.debug and (
+                TaskDebugFlag.keep_node_on_failure in task.config.debug
+                or TaskDebugFlag.keep_node_on_completion in task.config.debug
+            ):
+                node.debug_keep_node = True
+                node.save()
+
         else:
             task.mark_stopping()
+            if (
+                task.config.debug
+                and TaskDebugFlag.keep_node_on_completion in task.config.debug
+            ):
+                node.debug_keep_node = True
+                node.save()
 
         node.to_reimage(done=True)
     else:

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

@@ -6,7 +6,7 @@
 import azure.functions as func
 from onefuzztypes.enums import ErrorCode
 from onefuzztypes.models import Error
-from onefuzztypes.requests import NodeGet, NodeSearch
+from onefuzztypes.requests import NodeGet, NodeSearch, NodeUpdate
 from onefuzztypes.responses import BoolResult
 
 from ..onefuzzlib.pools import Node, NodeTasks
@@ -42,6 +42,24 @@ def get(req: func.HttpRequest) -> func.HttpResponse:
     return ok(nodes)
 
 
+def post(req: func.HttpRequest) -> func.HttpResponse:
+    request = parse_request(NodeUpdate, req)
+    if isinstance(request, Error):
+        return not_ok(request, context="NodeUpdate")
+
+    node = Node.get_by_machine_id(request.machine_id)
+    if not node:
+        return not_ok(
+            Error(code=ErrorCode.UNABLE_TO_FIND, errors=["unable to find node"]),
+            context=request.machine_id,
+        )
+    if request.debug_keep_node is not None:
+        node.debug_keep_node = request.debug_keep_node
+
+    node.save()
+    return ok(BoolResult(result=True))
+
+
 def delete(req: func.HttpRequest) -> func.HttpResponse:
     request = parse_request(NodeGet, req)
     if isinstance(request, Error):
@@ -55,6 +73,9 @@ def delete(req: func.HttpRequest) -> func.HttpResponse:
         )
 
     node.set_halt()
+    if node.debug_keep_node:
+        node.debug_keep_node = False
+        node.save()
 
     return ok(BoolResult(result=True))
 
@@ -72,6 +93,9 @@ def patch(req: func.HttpRequest) -> func.HttpResponse:
         )
 
     node.stop()
+    if node.debug_keep_node:
+        node.debug_keep_node = False
+        node.save()
     return ok(BoolResult(result=True))
 
 

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

@@ -847,7 +847,16 @@ class Scaleset(BASE_SCALESET, ORMMixin):
             logging.debug("scaleset delete will delete node: %s", self.scaleset_id)
             return
 
-        machine_ids = [x.machine_id for x in nodes]
+        machine_ids = []
+        for node in nodes:
+            if node.debug_keep_node:
+                logging.warning(
+                    "delete manually overridden %s:%s",
+                    self.scaleset_id,
+                    node.machine_id,
+                )
+            else:
+                machine_ids.append(node.machine_id)
 
         logging.info("deleting %s:%s", self.scaleset_id, machine_ids)
         delete_vmss_nodes(self.scaleset_id, machine_ids)
@@ -865,7 +874,16 @@ class Scaleset(BASE_SCALESET, ORMMixin):
             logging.debug("scaleset delete will delete node: %s", self.scaleset_id)
             return
 
-        machine_ids = [x.machine_id for x in nodes]
+        machine_ids = []
+        for node in nodes:
+            if node.debug_keep_node:
+                logging.warning(
+                    "reimage manually overridden %s:%s",
+                    self.scaleset_id,
+                    node.machine_id,
+                )
+            else:
+                machine_ids.append(node.machine_id)
 
         result = reimage_vmss_nodes(self.scaleset_id, machine_ids)
         if isinstance(result, Error):

src/cli/onefuzz/api.py

@@ -666,6 +666,7 @@ class Tasks(Endpoint):
         analyzer_env: Optional[Dict[str, str]] = None,
         tags: Optional[Dict[str, str]] = None,
         prereq_tasks: Optional[List[UUID]] = None,
+        debug: Optional[List[enums.TaskDebugFlag]] = None,
     ) -> models.Task:
         """ Create a task """
         self.logger.debug("creating task: %s", task_type)
@@ -729,6 +730,7 @@ class Tasks(Endpoint):
             pool=models.TaskPool(count=vm_count, pool_name=pool_name),
             containers=containers_submit,
             tags=tags,
+            debug=debug,
         )
 
         return self.create_with_config(config)
@@ -983,6 +985,28 @@ class Node(Endpoint):
             data=requests.NodeGet(machine_id=machine_id_expanded),
         )
 
+    def update(
+        self,
+        machine_id: UUID_EXPANSION,
+        *,
+        debug_keep_node: Optional[bool] = None,
+    ) -> responses.BoolResult:
+        self.logger.debug("update node: %s", machine_id)
+        machine_id_expanded = self._disambiguate_uuid(
+            "machine_id",
+            machine_id,
+            lambda: [str(x.machine_id) for x in self.list()],
+        )
+
+        return self._req_model(
+            "POST",
+            responses.BoolResult,
+            data=requests.NodeUpdate(
+                machine_id=machine_id_expanded,
+                debug_keep_node=debug_keep_node,
+            ),
+        )
+
     def list(
         self,
         *,

src/cli/onefuzz/cli.py

@@ -86,6 +86,19 @@ class AsDict(argparse.Action):
         setattr(namespace, self.dest, as_dict)
 
 
+def arg_bool(arg: str) -> bool:
+    acceptable = ["true", "false"]
+    if arg not in acceptable:
+        raise argparse.ArgumentTypeError(
+            "invalid value: %s, must be %s"
+            % (
+                repr(arg),
+                " or ".join(acceptable),
+            )
+        )
+    return arg == "true"
+
+
 def arg_dir(arg: str) -> str:
     if not os.path.isdir(arg):
         raise argparse.ArgumentTypeError("not a directory: %s" % arg)
@@ -241,19 +254,27 @@ class Builder:
             }
             return result
 
-        if issubclass(annotation, bool) and default is False:
-            return {
-                "action": "store_true",
-                "optional": True,
-                "help": "(Default: False.  Sets value to True)",
-            }
-
-        if issubclass(annotation, bool) and default is True:
-            return {
-                "action": "store_false",
-                "optional": True,
-                "help": "(Default: True.  Sets value to False)",
-            }
+        if issubclass(annotation, bool):
+            if default is False:
+                return {
+                    "action": "store_true",
+                    "optional": True,
+                    "help": "(Default: False.  Sets value to True)",
+                }
+            elif default is True:
+                return {
+                    "action": "store_false",
+                    "optional": True,
+                    "help": "(Default: True.  Sets value to False)",
+                }
+            elif default is None:
+                return {
+                    "type": arg_bool,
+                    "optional": True,
+                    "help": "Provide 'true' to set to true and 'false' to set to false",
+                }
+            else:
+                raise Exception("Argument parsing error: %s", repr(default))
 
         if issubclass(annotation, BaseModel):
 

src/cli/onefuzz/templates/afl.py

@@ -5,7 +5,7 @@
 
 from typing import Dict, List, Optional
 
-from onefuzztypes.enums import OS, ContainerType, StatsFormat, TaskType
+from onefuzztypes.enums import OS, ContainerType, StatsFormat, TaskDebugFlag, TaskType
 from onefuzztypes.models import Job, NotificationConfig
 from onefuzztypes.primitives import Container, Directory, File
 
@@ -51,6 +51,7 @@ class AFL(Command):
         existing_inputs: Optional[Container] = None,
         dryrun: bool = False,
         notification_config: Optional[NotificationConfig] = None,
+        debug: Optional[List[TaskDebugFlag]] = None,
     ) -> Optional[Job]:
         """
         Basic AFL job
@@ -144,6 +145,7 @@ class AFL(Command):
             stats_format=StatsFormat.AFL,
             task_wait_for_files=ContainerType.inputs,
             tags=helper.tags,
+            debug=debug,
         )
 
         report_containers = [
@@ -170,6 +172,7 @@ class AFL(Command):
             check_debugger=True,
             tags=tags,
             prereq_tasks=[fuzzer_task.task_id],
+            debug=debug,
         )
 
         self.logger.info("done creating tasks")

src/cli/onefuzz/templates/libfuzzer.py

@@ -5,7 +5,7 @@
 
 from typing import Dict, List, Optional
 
-from onefuzztypes.enums import ContainerType, TaskType
+from onefuzztypes.enums import ContainerType, TaskDebugFlag, TaskType
 from onefuzztypes.models import Job, NotificationConfig
 from onefuzztypes.primitives import Container, Directory, File
 
@@ -46,6 +46,7 @@ class Libfuzzer(Command):
         tags: Optional[Dict[str, str]] = None,
         check_retry_count: Optional[int] = None,
         crash_report_timeout: Optional[int] = None,
+        debug: Optional[List[TaskDebugFlag]] = None,
     ) -> None:
 
         fuzzer_containers = [
@@ -67,6 +68,7 @@ class Libfuzzer(Command):
             target_env=target_env,
             target_workers=target_workers,
             tags=tags,
+            debug=debug,
         )
 
         coverage_containers = [
@@ -88,6 +90,7 @@ class Libfuzzer(Command):
             target_env=target_env,
             tags=tags,
             prereq_tasks=[fuzzer_task.task_id],
+            debug=debug,
         )
 
         report_containers = [
@@ -114,6 +117,7 @@ class Libfuzzer(Command):
             prereq_tasks=[fuzzer_task.task_id],
             target_timeout=crash_report_timeout,
             check_retry_count=check_retry_count,
+            debug=debug,
         )
 
     def basic(
@@ -141,6 +145,7 @@ class Libfuzzer(Command):
         existing_inputs: Optional[Container] = None,
         dryrun: bool = False,
         notification_config: Optional[NotificationConfig] = None,
+        debug: Optional[List[TaskDebugFlag]] = None,
     ) -> Optional[Job]:
         """ Basic libfuzzer job """
 
@@ -207,6 +212,7 @@ class Libfuzzer(Command):
             tags=helper.tags,
             crash_report_timeout=crash_report_timeout,
             check_retry_count=check_retry_count,
+            debug=debug,
         )
 
         self.logger.info("done creating tasks")

src/cli/onefuzz/templates/ossfuzz.py

@@ -9,7 +9,7 @@ import os
 import subprocess  # nosec
 from typing import Dict, List, Optional, Tuple
 
-from onefuzztypes.enums import OS, ContainerType
+from onefuzztypes.enums import OS, ContainerType, TaskDebugFlag
 from onefuzztypes.models import NotificationConfig
 from onefuzztypes.primitives import File
 
@@ -113,6 +113,7 @@ class OssFuzz(Command):
         max_target_count: int = 20,
         sync_inputs: bool = False,
         notification_config: Optional[NotificationConfig] = None,
+        debug: Optional[List[TaskDebugFlag]] = None,
     ) -> None:
         """ OssFuzz style libfuzzer jobs """
 
@@ -234,6 +235,7 @@ class OssFuzz(Command):
                 target_options=target_options,
                 target_env=target_env,
                 tags=helper.tags,
+                debug=debug,
             )
             helpers.append(helper)
         base_helper.wait()

src/cli/onefuzz/templates/radamsa.py

@@ -5,7 +5,7 @@
 
 from typing import Dict, List, Optional
 
-from onefuzztypes.enums import OS, ContainerType, TaskType
+from onefuzztypes.enums import OS, ContainerType, TaskDebugFlag, TaskType
 from onefuzztypes.models import Job, NotificationConfig
 from onefuzztypes.primitives import Container, Directory, File
 
@@ -47,6 +47,7 @@ class Radamsa(Command):
         disable_check_debugger: bool = False,
         dryrun: bool = False,
         notification_config: Optional[NotificationConfig] = None,
+        debug: Optional[List[TaskDebugFlag]] = None,
     ) -> Optional[Job]:
         """ Basic radamsa job """
 
@@ -160,6 +161,7 @@ class Radamsa(Command):
             check_debugger=not disable_check_debugger,
             tags=helper.tags,
             rename_output=rename_output,
+            debug=debug,
         )
 
         report_containers = [
@@ -190,6 +192,7 @@ class Radamsa(Command):
             check_debugger=not disable_check_debugger,
             check_retry_count=check_retry_count,
             prereq_tasks=[fuzzer_task.task_id],
+            debug=debug,
         )
 
         if helper.platform == OS.windows:
@@ -230,6 +233,7 @@ class Radamsa(Command):
                 analyzer_env=analyzer_env,
                 tags=helper.tags,
                 prereq_tasks=[fuzzer_task.task_id],
+                debug=debug,
             )
 
         self.logger.info("done creating tasks")

src/pytypes/onefuzztypes/enums.py

@@ -351,3 +351,8 @@ class GithubIssueState(Enum):
 class GithubIssueSearchMatch(Enum):
     title = "title"
     body = "body"
+
+
+class TaskDebugFlag(Enum):
+    keep_node_on_failure = "keep_node_on_failure"
+    keep_node_on_completion = "keep_node_on_completion"

src/pytypes/onefuzztypes/models.py

@@ -26,6 +26,7 @@ from .enums import (
     PoolState,
     ScalesetState,
     StatsFormat,
+    TaskDebugFlag,
     TaskFeature,
     TaskState,
     TaskType,
@@ -176,6 +177,7 @@ class TaskConfig(BaseModel):
     pool: Optional[TaskPool]
     containers: List[TaskContainers]
     tags: Dict[str, str]
+    debug: Optional[List[TaskDebugFlag]]
 
 
 class BlobRef(BaseModel):
@@ -444,6 +446,7 @@ class Node(BaseModel):
     version: str = Field(default="1.0.0")
     reimage_requested: bool = Field(default=False)
     delete_requested: bool = Field(default=False)
+    debug_keep_node: bool = Field(default=False)
 
 
 class ScalesetSummary(BaseModel):

src/pytypes/onefuzztypes/requests.py

@@ -141,6 +141,11 @@ class NodeGet(BaseRequest):
     machine_id: UUID
 
 
+class NodeUpdate(BaseRequest):
+    machine_id: UUID
+    debug_keep_node: Optional[bool]
+
+
 class ScalesetSearch(BaseRequest):
     scaleset_id: Optional[UUID]
     state: Optional[List[ScalesetState]]