build_generation: Add auto discovery agent (#1014)

This PR add a new build_generation agent to auto discovery the target
project in a docker image and create a build script after gathered
enough information.

---------

Signed-off-by: Arthur Chan <[email protected]>

Author: arthurscchan

experimental/build_generator/llm_agent.py

@@ -31,6 +31,7 @@
 
 MAX_PROMPT_LENGTH = 20000
 SAMPLE_HEADERS_COUNT = 30
+MAX_DISCOVERY_ROUND = 100
 
 
 class BuildScriptAgent(BaseAgent):
@@ -51,6 +52,7 @@ def __init__(self,
     self.last_status = False
     self.last_result = ''
     self.target_files = {}
+    self.discovery_stage = False
 
     # Get sample fuzzing harness
     _, _, self.harness_path, self.harness_code = (
@@ -83,31 +85,48 @@ def _parse_tags(self, response: str, tag: str) -> list[str]:
   def _container_handle_bash_commands(self, response: str, tool: BaseTool,
                                       prompt: Prompt) -> Prompt:
     """Handles the command from LLM with container |tool|."""
-    # Update fuzzing harness
-    harness = self._parse_tag(response, 'fuzzer')
-    if harness:
-      self.harness_code = harness
-    if isinstance(tool, ProjectContainerTool):
-      tool.write_to_file(self.harness_code, self.harness_path)
-
-    # Update build script
-    command = '\n'.join(self._parse_tags(response, 'bash'))
+    # Initialise variables
     prompt_text = ''
     success = False
-    if command:
-      # Add set -e to ensure docker image failing is reflected.
-      command = command.replace('#!/bin/bash', '')
-      command = f'#!/bin/bash\nset -e\n{command}'
 
-      # Update build script
+    # Retrieve data from response
+    harness = self._parse_tag(response, 'fuzzer')
+    build_script = '\n'.join(self._parse_tags(response, 'bash'))
+    commands = '; '.join(self._parse_tags(response, 'command'))
+
+    if build_script:
+      self.discovery_stage = False
+
+      # Update fuzzing harness
+      if harness:
+        self.harness_code = harness
       if isinstance(tool, ProjectContainerTool):
-        tool.write_to_file(command, '/src/build.sh')
+        tool.write_to_file(self.harness_code, self.harness_path)
 
-      # Test and parse result
-      result = tool.execute('compile')
-      format_result = self._format_bash_execution_result(result,
-                                                         previous_prompt=prompt)
-      prompt_text = self._parse_tag(format_result, 'stderr') + '\n'
+      # Update build script
+      if build_script:
+        # Add set -e to ensure docker image failing is reflected.
+        build_script = build_script.replace('#!/bin/bash', '')
+        build_script = f'#!/bin/bash\nset -e\n{build_script}'
+
+        # Update build script
+        if isinstance(tool, ProjectContainerTool):
+          tool.write_to_file(build_script, '/src/build.sh')
+
+          # Test and parse result
+          result = tool.execute('compile')
+          format_result = self._format_bash_execution_result(
+              result, previous_prompt=prompt)
+          prompt_text = self._parse_tag(format_result, 'stderr') + '\n'
+          if result.returncode == 0:
+            success = True
+
+    elif commands:
+      # Execute the command directly, then return the formatted result
+      self.discovery_stage = True
+      result = tool.execute(commands)
+      prompt_text = self._format_bash_execution_result(result,
+                                                       previous_prompt=prompt)
       if result.returncode == 0:
         success = True
 
@@ -192,8 +211,6 @@ def _discover_headers(self) -> list[str]:
         if file.endswith((".h", ".hpp")):
           header_path = os.path.join(root, file)
           headers.add(header_path.replace(target_path, ''))
-        if len(headers) > SAMPLE_HEADERS_COUNT:
-          return list(headers)
 
     return list(headers)
 
@@ -205,6 +222,7 @@ def execute(self, result_history: list[Result]) -> BuildResult:
     self.inspect_tool = ProjectContainerTool(benchmark, name='inspect')
     self.inspect_tool.compile(extra_commands=' && rm -rf /out/* > /dev/null')
     cur_round = 1
+    dis_round = 1
     build_result = BuildResult(benchmark=benchmark,
                                trial=last_result.trial,
                                work_dirs=last_result.work_dirs,
@@ -214,7 +232,7 @@ def execute(self, result_history: list[Result]) -> BuildResult:
     prompt = self._initial_prompt(result_history)
     try:
       client = self.llm.get_chat_client(model=self.llm.get_model())
-      while prompt and cur_round < self.max_round:
+      while prompt:
         # Sleep for a minute to avoid over RPM
         time.sleep(60)
 
@@ -224,7 +242,15 @@ def execute(self, result_history: list[Result]) -> BuildResult:
                                  trial=last_result.trial)
         prompt = self._container_tool_reaction(cur_round, response,
                                                build_result)
-        cur_round += 1
+
+        if self.discovery_stage:
+          dis_round += 1
+          if dis_round >= MAX_DISCOVERY_ROUND:
+            break
+        else:
+          cur_round += 1
+          if cur_round >= self.max_round:
+            break
     finally:
       logger.info('Stopping and removing the inspect container %s',
                   self.inspect_tool.container_id,
@@ -291,6 +317,7 @@ def _initial_prompt(self, results: list[Result]) -> Prompt:  # pylint: disable=u
     # Extract template Dockerfile content
     dockerfile_str = templates.CLEAN_OSS_FUZZ_DOCKER
     dockerfile_str = dockerfile_str.replace('{additional_packages}', '')
+    dockerfile_str = dockerfile_str.replace('{fuzzer_dir}', '$SRC/')
     dockerfile_str = dockerfile_str.replace('{repo_url}', self.github_url)
     dockerfile_str = dockerfile_str.replace('{project_repo_dir}',
                                             self.github_url.split('/')[-1])
@@ -304,7 +331,8 @@ def _initial_prompt(self, results: list[Result]) -> Prompt:  # pylint: disable=u
                               self.harness_path.split('/')[-1])
 
     headers = self._discover_headers()
-    problem = problem.replace('{HEADERS}', ','.join(headers))
+    problem = problem.replace('{HEADERS}',
+                              ','.join(headers[:SAMPLE_HEADERS_COUNT]))
 
     prompt.add_priming(templates.LLM_PRIMING)
     prompt.add_problem(problem)
@@ -324,3 +352,66 @@ def execute(self, result_history: list[Result]) -> BuildResult:
                          chat_history={self.name: ''})
 
     return super().execute(result_history)
+
+
+class AutoDiscoveryBuildScriptAgent(BuildScriptAgent):
+  """Generate a working Dockerfile and build script from scratch
+  with LLM auto discovery"""
+
+  def _initial_prompt(self, results: list[Result]) -> Prompt:  # pylint: disable=unused-argument
+    """Constructs initial prompt of the agent."""
+    prompt = self.llm.prompt_type()(None)
+
+    # Extract template Dockerfile content
+    dockerfile_str = templates.CLEAN_OSS_FUZZ_DOCKER
+    dockerfile_str = dockerfile_str.replace('{additional_packages}', '')
+    dockerfile_str = dockerfile_str.replace('{repo_url}', self.github_url)
+    dockerfile_str = dockerfile_str.replace('{project_repo_dir}',
+                                            self.github_url.split('/')[-1])
+
+    # Prepare prompt problem string
+    problem = templates.LLM_AUTO_DISCOVERY
+    problem = problem.replace('{PROJECT_NAME}', self.github_url.split('/')[-1])
+    problem = problem.replace('{DOCKERFILE}', dockerfile_str)
+    problem = problem.replace('{FUZZER}', self.harness_code)
+    problem = problem.replace('{MAX_DISCOVERY_ROUND}', str(MAX_DISCOVERY_ROUND))
+    problem = problem.replace('{FUZZING_FILE}',
+                              self.harness_path.split('/')[-1])
+
+    prompt.add_priming(templates.LLM_PRIMING)
+    prompt.add_problem(problem)
+
+    return prompt
+
+  def _container_tool_reaction(self, cur_round: int, response: str,
+                               build_result: BuildResult) -> Optional[Prompt]:
+    """Validates LLM conclusion or executes its command."""
+    prompt = self.llm.prompt_type()(None)
+
+    if response:
+      prompt = self._container_handle_bash_commands(response, self.inspect_tool,
+                                                    prompt)
+
+      if self.discovery_stage:
+        # Relay the command output back to LLM
+        feedback = templates.LLM_DOCKER_FEEDBACK
+        feedback = feedback.replace('{RESULT}', self.last_result)
+        prompt.add_problem(feedback)
+      else:
+        # Check result and try building with the new builds script
+        prompt = self._container_handle_conclusion(cur_round, response,
+                                                   build_result, prompt)
+
+        if prompt is None:
+          return None
+
+    if not response or not prompt or not prompt.get():
+      prompt = self._container_handle_invalid_tool_usage(
+          self.inspect_tool, cur_round, response, prompt)
+
+    return prompt
+
+  def execute(self, result_history: list[Result]) -> BuildResult:
+    """Executes the agent based on previous result."""
+    self._prepare_repository()
+    return super().execute(result_history)

experimental/build_generator/manager.py

@@ -302,7 +302,8 @@ def create_clean_oss_fuzz_from_empty(github_repo: str, build_worker,
   dockerfile = templates.CLEAN_OSS_FUZZ_DOCKER.format(
       repo_url=github_repo,
       project_repo_dir=project_repo_dir,
-      additional_packages=' '.join(additional_packages))
+      additional_packages=' '.join(additional_packages),
+      fuzzer_dir='$SRC/fuzzers/')
   with open(os.path.join(oss_fuzz_folder, 'Dockerfile'), 'w') as docker_out:
     docker_out.write(dockerfile)
 
@@ -346,7 +347,8 @@ def create_clean_oss_fuzz_from_success(github_repo: str, out_dir: str,
   dockerfile = templates.CLEAN_OSS_FUZZ_DOCKER.format(
       repo_url=github_repo,
       project_repo_dir=project_repo_dir,
-      additional_packages=' '.join(pkgs))
+      additional_packages=' '.join(pkgs),
+      fuzzer_dir='$SRC/fuzzers/')
   with open(os.path.join(oss_fuzz_folder, 'Dockerfile'), 'w') as docker_out:
     docker_out.write(dockerfile)
 
@@ -546,8 +548,8 @@ def auto_generate(github_url, disable_testing_build_scripts=False, outdir=''):
             build_worker.build_suggestion, build_worker.build_script,
             build_worker.build_directory,
             build_worker.executable_files_build.copy())
-        new_worker.build_suggestion.heuristic_id = new_worker.build_suggestion.heuristic_id + '-%d' % (
-            b_idx)
+        new_worker.build_suggestion.heuristic_id = (
+            new_worker.build_suggestion.heuristic_id + f'-{b_idx}')
         new_worker.executable_files_build['refined-static-libs'] = [ref_lib]
         refined_builds.append((test_dir, new_worker))
     refined_builds.append((test_dir, build_worker))

experimental/build_generator/runner.py

@@ -63,6 +63,7 @@ def setup_worker_project(oss_fuzz_base: str,
       file_content = templates.CLEAN_OSS_FUZZ_DOCKER
       file_content = file_content.replace('{additional_packages}', '')
       file_content = file_content.replace('{repo_url}', github_url)
+      file_content = file_content.replace('{fuzzer_dir}', '$SRC/')
       file_content = file_content.replace('{project_repo_dir}',
                                           github_url.split('/')[-1])
     else:
@@ -379,7 +380,10 @@ def run_agent(target_repositories: List[str], args: argparse.Namespace):
   )
 
   # All agents
-  llm_agents = [llm_agent.BuildSystemBuildScriptAgent]
+  llm_agents = [
+      llm_agent.BuildSystemBuildScriptAgent,
+      llm_agent.AutoDiscoveryBuildScriptAgent
+  ]
 
   for target_repository in target_repositories:
     logger.info('Target repository: %s', target_repository)