patcher: improve context retriever agent (#357)

This commit is contained in:
Riccardo Schirone
2025-04-07 12:39:47 +02:00
committed by GitHub
parent 98e736c194
commit 658f9c64c4
22 changed files with 1723 additions and 704 deletions
@@ -0,0 +1,17 @@
crash {
crash {
target {
engine: "libfuzzer"
sanitizer: "address"
task_dir: "/crs_scratch/my-task-id/"
task_id: "my-task-id"
build_type: "fuzzer"
apply_diff: true
}
harness_name: "libpng_read_fuzzer"
crash_input_path: "/crs_scratch/my-task-id-pov.bin"
stacktrace: "+ FUZZER=libpng_read_fuzzer\n+ shift\n+ '[' '!' -v TESTCASE ']'\n+ TESTCASE=/testcase\n+ '[' '!' -f /testcase ']'\n+ export RUN_FUZZER_MODE=interactive\n+ RUN_FUZZER_MODE=interactive\n+ export FUZZING_ENGINE=libfuzzer\n+ FUZZING_ENGINE=libfuzzer\n+ export SKIP_SEED_CORPUS=1\n+ SKIP_SEED_CORPUS=1\n+ run_fuzzer libpng_read_fuzzer -runs=100 /testcase\nvm.mmap_rnd_bits = 28\n/out/libpng_read_fuzzer -rss_limit_mb=2560 -timeout=25 -runs=100 /testcase -dict=png.dict < /dev/null\nDictionary: 28 entries\nINFO: Running with entropic power schedule (0xFF, 100).\nINFO: Seed: 4117878429\nINFO: Loaded 1 modules (5641 inline 8-bit counters): 5641 [0x561561195928, 0x561561196f31), \nINFO: Loaded 1 PC tables (5641 PCs): 5641 [0x561561196f38,0x5615611acfc8), \n/out/libpng_read_fuzzer: Running 1 inputs 100 time(s) each.\nRunning: /testcase\n=================================================================\n==18==ERROR: AddressSanitizer: dynamic-stack-buffer-overflow on address 0x7ffce0754432 at pc 0x5615610dfa9b bp 0x7ffce07543b0 sp 0x7ffce07543a8\nREAD of size 2 at 0x7ffce0754432 thread T0\nSCARINESS: 29 (2-byte-read-dynamic-stack-buffer-overflow)\n #0 0x5615610dfa9a in OSS_FUZZ_png_handle_iCCP /src/libpng/pngrutil.c:1447:10\n #1 0x5615610b3dcd in OSS_FUZZ_png_read_info /src/libpng/pngread.c:229:10\n #2 0x5615610074ae in LLVMFuzzerTestOneInput /src/libpng/contrib/oss-fuzz/libpng_read_fuzzer.cc:156:3\n #3 0x561561025520 in fuzzer::Fuzzer::ExecuteCallback(unsigned char const*, unsigned long) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerLoop.cpp:614:13\n #4 0x561561010795 in fuzzer::RunOneTest(fuzzer::Fuzzer*, char const*, unsigned long) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:327:6\n #5 0x56156101622f in fuzzer::FuzzerDriver(int*, char***, int (*)(unsigned char const*, unsigned long)) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:862:9\n #6 0x5615610414d2 in main /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerMain.cpp:20:10\n #7 0x7fedae19f082 in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x24082) (BuildId: 0323ab4806bee6f846d9ad4bccfc29afdca49a58)\n #8 0x561560f2e83d in _start (/out/libpng_read_fuzzer+0x6c83d)\n\nDEDUP_TOKEN: OSS_FUZZ_png_handle_iCCP--OSS_FUZZ_png_read_info--LLVMFuzzerTestOneInput\nAddress 0x7ffce0754432 is located in stack of thread T0\nSUMMARY: AddressSanitizer: dynamic-stack-buffer-overflow /src/libpng/pngrutil.c:1447:10 in OSS_FUZZ_png_handle_iCCP\nShadow bytes around the buggy address:\n 0x7ffce0754180: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754200: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754280: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754300: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754380: 00 00 00 00 00 00 00 00 ca ca ca ca 00 00 00 00\n=>0x7ffce0754400: 00 00 00 00 00 00[02]cb cb cb cb cb 00 00 00 00\n 0x7ffce0754480: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754500: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754580: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754600: 00 00 00 00 00 00 00 00 f1 f1 f1 f1 00 00 00 f2\n 0x7ffce0754680: f2 f2 f2 f2 00 00 00 00 00 f2 f2 f2 f2 f2 f8 f2\nShadow byte legend (one shadow byte represents 8 application bytes):\n Addressable: 00\n Partially addressable: 01 02 03 04 05 06 07 \n Heap left redzone: fa\n Freed heap region: fd\n Stack left redzone: f1\n Stack mid redzone: f2\n Stack right redzone: f3\n Stack after return: f5\n Stack use after scope: f8\n Global redzone: f9\n Global init order: f6\n Poisoned by user: f7\n Container overflow: fc\n Array cookie: ac\n Intra object redzone: bb\n ASan internal: fe\n Left alloca redzone: ca\n Right alloca redzone: cb\n==18==ABORTING\n"
}
tracer_stacktrace: "+ FUZZER=libpng_read_fuzzer\n+ shift\n+ '[' '!' -v TESTCASE ']'\n+ TESTCASE=/testcase\n+ '[' '!' -f /testcase ']'\n+ export RUN_FUZZER_MODE=interactive\n+ RUN_FUZZER_MODE=interactive\n+ export FUZZING_ENGINE=libfuzzer\n+ FUZZING_ENGINE=libfuzzer\n+ export SKIP_SEED_CORPUS=1\n+ SKIP_SEED_CORPUS=1\n+ run_fuzzer libpng_read_fuzzer -runs=100 /testcase\nvm.mmap_rnd_bits = 28\n/out/libpng_read_fuzzer -rss_limit_mb=2560 -timeout=25 -runs=100 /testcase -dict=png.dict < /dev/null\nDictionary: 28 entries\nINFO: Running with entropic power schedule (0xFF, 100).\nINFO: Seed: 4117878429\nINFO: Loaded 1 modules (5641 inline 8-bit counters): 5641 [0x561561195928, 0x561561196f31), \nINFO: Loaded 1 PC tables (5641 PCs): 5641 [0x561561196f38,0x5615611acfc8), \n/out/libpng_read_fuzzer: Running 1 inputs 100 time(s) each.\nRunning: /testcase\n=================================================================\n==18==ERROR: AddressSanitizer: dynamic-stack-buffer-overflow on address 0x7ffce0754432 at pc 0x5615610dfa9b bp 0x7ffce07543b0 sp 0x7ffce07543a8\nREAD of size 2 at 0x7ffce0754432 thread T0\nSCARINESS: 29 (2-byte-read-dynamic-stack-buffer-overflow)\n #0 0x5615610dfa9a in OSS_FUZZ_png_handle_iCCP /src/libpng/pngrutil.c:1447:10\n #1 0x5615610b3dcd in OSS_FUZZ_png_read_info /src/libpng/pngread.c:229:10\n #2 0x5615610074ae in LLVMFuzzerTestOneInput /src/libpng/contrib/oss-fuzz/libpng_read_fuzzer.cc:156:3\n #3 0x561561025520 in fuzzer::Fuzzer::ExecuteCallback(unsigned char const*, unsigned long) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerLoop.cpp:614:13\n #4 0x561561010795 in fuzzer::RunOneTest(fuzzer::Fuzzer*, char const*, unsigned long) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:327:6\n #5 0x56156101622f in fuzzer::FuzzerDriver(int*, char***, int (*)(unsigned char const*, unsigned long)) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:862:9\n #6 0x5615610414d2 in main /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerMain.cpp:20:10\n #7 0x7fedae19f082 in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x24082) (BuildId: 0323ab4806bee6f846d9ad4bccfc29afdca49a58)\n #8 0x561560f2e83d in _start (/out/libpng_read_fuzzer+0x6c83d)\n\nDEDUP_TOKEN: OSS_FUZZ_png_handle_iCCP--OSS_FUZZ_png_read_info--LLVMFuzzerTestOneInput\nAddress 0x7ffce0754432 is located in stack of thread T0\nSUMMARY: AddressSanitizer: dynamic-stack-buffer-overflow /src/libpng/pngrutil.c:1447:10 in OSS_FUZZ_png_handle_iCCP\nShadow bytes around the buggy address:\n 0x7ffce0754180: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754200: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754280: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754300: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754380: 00 00 00 00 00 00 00 00 ca ca ca ca 00 00 00 00\n=>0x7ffce0754400: 00 00 00 00 00 00[02]cb cb cb cb cb 00 00 00 00\n 0x7ffce0754480: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754500: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754580: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00\n 0x7ffce0754600: 00 00 00 00 00 00 00 00 f1 f1 f1 f1 00 00 00 f2\n 0x7ffce0754680: f2 f2 f2 f2 00 00 00 00 00 f2 f2 f2 f2 f2 f8 f2\nShadow byte legend (one shadow byte represents 8 application bytes):\n Addressable: 00\n Partially addressable: 01 02 03 04 05 06 07 \n Heap left redzone: fa\n Freed heap region: fd\n Stack left redzone: f1\n Stack mid redzone: f2\n Stack right redzone: f3\n Stack after return: f5\n Stack use after scope: f8\n Global redzone: f9\n Global init order: f6\n Poisoned by user: f7\n Container overflow: fc\n Array cookie: ac\n Intra object redzone: bb\n ASan internal: fe\n Left alloca redzone: ca\n Right alloca redzone: cb\n==18==ABORTING\n"
}
vuln_id: "my-vuln-id"
Binary file not shown.

After

Width:  |  Height:  |  Size: 420 B

@@ -0,0 +1,70 @@
+ FUZZER=libpng_read_fuzzer
+ shift
+ '[' '!' -v TESTCASE ']'
+ TESTCASE=/testcase
+ '[' '!' -f /testcase ']'
+ export RUN_FUZZER_MODE=interactive
+ RUN_FUZZER_MODE=interactive
+ export FUZZING_ENGINE=libfuzzer
+ FUZZING_ENGINE=libfuzzer
+ export SKIP_SEED_CORPUS=1
+ SKIP_SEED_CORPUS=1
+ run_fuzzer libpng_read_fuzzer -runs=100 /testcase
vm.mmap_rnd_bits = 28
/out/libpng_read_fuzzer -rss_limit_mb=2560 -timeout=25 -runs=100 /testcase -dict=png.dict < /dev/null
Dictionary: 28 entries
INFO: Running with entropic power schedule (0xFF, 100).
INFO: Seed: 4117878429
INFO: Loaded 1 modules (5641 inline 8-bit counters): 5641 [0x561561195928, 0x561561196f31),
INFO: Loaded 1 PC tables (5641 PCs): 5641 [0x561561196f38,0x5615611acfc8),
/out/libpng_read_fuzzer: Running 1 inputs 100 time(s) each.
Running: /testcase
=================================================================
==18==ERROR: AddressSanitizer: dynamic-stack-buffer-overflow on address 0x7ffce0754432 at pc 0x5615610dfa9b bp 0x7ffce07543b0 sp 0x7ffce07543a8
READ of size 2 at 0x7ffce0754432 thread T0
SCARINESS: 29 (2-byte-read-dynamic-stack-buffer-overflow)
#0 0x5615610dfa9a in OSS_FUZZ_png_handle_iCCP /src/libpng/pngrutil.c:1447:10
#1 0x5615610b3dcd in OSS_FUZZ_png_read_info /src/libpng/pngread.c:229:10
#2 0x5615610074ae in LLVMFuzzerTestOneInput /src/libpng/contrib/oss-fuzz/libpng_read_fuzzer.cc:156:3
#3 0x561561025520 in fuzzer::Fuzzer::ExecuteCallback(unsigned char const*, unsigned long) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerLoop.cpp:614:13
#4 0x561561010795 in fuzzer::RunOneTest(fuzzer::Fuzzer*, char const*, unsigned long) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:327:6
#5 0x56156101622f in fuzzer::FuzzerDriver(int*, char***, int (*)(unsigned char const*, unsigned long)) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:862:9
#6 0x5615610414d2 in main /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerMain.cpp:20:10
#7 0x7fedae19f082 in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x24082) (BuildId: 0323ab4806bee6f846d9ad4bccfc29afdca49a58)
#8 0x561560f2e83d in _start (/out/libpng_read_fuzzer+0x6c83d)
DEDUP_TOKEN: OSS_FUZZ_png_handle_iCCP--OSS_FUZZ_png_read_info--LLVMFuzzerTestOneInput
Address 0x7ffce0754432 is located in stack of thread T0
SUMMARY: AddressSanitizer: dynamic-stack-buffer-overflow /src/libpng/pngrutil.c:1447:10 in OSS_FUZZ_png_handle_iCCP
Shadow bytes around the buggy address:
0x7ffce0754180: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
0x7ffce0754200: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
0x7ffce0754280: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
0x7ffce0754300: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
0x7ffce0754380: 00 00 00 00 00 00 00 00 ca ca ca ca 00 00 00 00
=>0x7ffce0754400: 00 00 00 00 00 00[02]cb cb cb cb cb 00 00 00 00
0x7ffce0754480: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
0x7ffce0754500: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
0x7ffce0754580: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
0x7ffce0754600: 00 00 00 00 00 00 00 00 f1 f1 f1 f1 00 00 00 f2
0x7ffce0754680: f2 f2 f2 f2 00 00 00 00 00 f2 f2 f2 f2 f2 f8 f2
Shadow byte legend (one shadow byte represents 8 application bytes):
Addressable: 00
Partially addressable: 01 02 03 04 05 06 07
Heap left redzone: fa
Freed heap region: fd
Stack left redzone: f1
Stack mid redzone: f2
Stack right redzone: f3
Stack after return: f5
Stack use after scope: f8
Global redzone: f9
Global init order: f6
Poisoned by user: f7
Container overflow: fc
Array cookie: ac
Intra object redzone: bb
ASan internal: fe
Left alloca redzone: ca
Right alloca redzone: cb
==18==ABORTING
@@ -0,0 +1,6 @@
engine: "libfuzzer"
sanitizer: "address"
task_dir: "/crs_scratch/selinux-OSV-2021-902-delta-1/"
task_id: "selinux-OSV-2021-902-delta-1"
build_type: "fuzzer"
apply_diff: true
File diff suppressed because one or more lines are too long
@@ -0,0 +1,83 @@
=================================================================
==14==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x7f013ccb77f0 at pc 0x555dcfad2fef bp 0x7ffd2207d630 sp 0x7ffd2207d628
READ of size 8 at 0x7f013ccb77f0 thread T0
SCARINESS: 33 (8-byte-read-heap-buffer-overflow-far-from-bounds)
#0 0x555dcfad2fee in ebitmap_match_any /src/selinux/libsepol/src/ebitmap.c:235:27
#1 0x555dcfac0756 in check_assertion_avtab_match /src/selinux/libsepol/src/assertion.c:462:7
#2 0x555dcfac7cd7 in avtab_map /src/selinux/libsepol/src/avtab.c:346:10
#3 0x555dcfac0335 in check_assertion /src/selinux/libsepol/src/assertion.c:513:7
#4 0x555dcfa2dacf in cil_check_neverallow /src/selinux/libsepol/src/../cil/src/cil_binary.c:4686:8
#5 0x555dcfa2dacf in cil_check_neverallows /src/selinux/libsepol/src/../cil/src/cil_binary.c:4740:8
#6 0x555dcfa2c204 in cil_binary_create_allocated_pdb /src/selinux/libsepol/src/../cil/src/cil_binary.c:4999:8
#7 0x555dcfa2b7a9 in cil_binary_create /src/selinux/libsepol/src/../cil/src/cil_binary.c:4326:7
#8 0x555dcfa073c1 in cil_build_policydb /src/selinux/libsepol/src/../cil/src/cil.c:672:7
#9 0x555dcf9ed982 in LLVMFuzzerTestOneInput /src/secilc-fuzzer.c:62:6
#10 0x555dcf8a2330 in fuzzer::Fuzzer::ExecuteCallback(unsigned char const*, unsigned long) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerLoop.cpp:614:13
#11 0x555dcf88d5a5 in fuzzer::RunOneTest(fuzzer::Fuzzer*, char const*, unsigned long) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:327:6
#12 0x555dcf89303f in fuzzer::FuzzerDriver(int*, char***, int (*)(unsigned char const*, unsigned long)) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:862:9
#13 0x555dcf8be2e2 in main /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerMain.cpp:20:10
#14 0x7f0141a73082 in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x24082) (BuildId: 0702430aef5fa3dda43986563e9ffcc47efbd75e)
#15 0x555dcf88578d in _start (/out/secilc-fuzzer+0xcd78d)
DEDUP_TOKEN: ebitmap_match_any--check_assertion_avtab_match--avtab_map
0x7f013ccb77f0 is located 16 bytes before 1340512-byte region [0x7f013ccb7800,0x7f013cdfec60)
allocated by thread T0 here:
#0 0x555dcf9ae0ff in malloc /src/llvm-project/compiler-rt/lib/asan/asan_malloc_linux.cpp:68:3
#1 0x555dcfa6e437 in cil_malloc /src/selinux/libsepol/src/../cil/src/cil_mem.c:40:14
#2 0x555dcfa12c4b in __cil_typeattr_bitmap_init /src/selinux/libsepol/src/../cil/src/cil_binary.c:599:23
#3 0x555dcfa130c9 in cil_typeattribute_to_bitmap /src/selinux/libsepol/src/../cil/src/cil_binary.c:636:8
#4 0x555dcfa28f8d in __cil_node_to_policydb /src/selinux/libsepol/src/../cil/src/cil_binary.c:3759:9
#5 0x555dcfab0bfb in cil_tree_walk_core /src/selinux/libsepol/src/../cil/src/cil_tree.c:263:9
#6 0x555dcfab0de1 in cil_tree_walk /src/selinux/libsepol/src/../cil/src/cil_tree.c:307:7
#7 0x555dcfab0de1 in cil_tree_walk_core /src/selinux/libsepol/src/../cil/src/cil_tree.c:275:9
#8 0x555dcfab0de1 in cil_tree_walk /src/selinux/libsepol/src/../cil/src/cil_tree.c:307:7
#9 0x555dcfab0de1 in cil_tree_walk_core /src/selinux/libsepol/src/../cil/src/cil_tree.c:275:9
#10 0x555dcfab0de1 in cil_tree_walk /src/selinux/libsepol/src/../cil/src/cil_tree.c:307:7
#11 0x555dcfab0de1 in cil_tree_walk_core /src/selinux/libsepol/src/../cil/src/cil_tree.c:275:9
#12 0x555dcfab0de1 in cil_tree_walk /src/selinux/libsepol/src/../cil/src/cil_tree.c:307:7
#13 0x555dcfab0de1 in cil_tree_walk_core /src/selinux/libsepol/src/../cil/src/cil_tree.c:275:9
#14 0x555dcfab0fee in cil_tree_walk /src/selinux/libsepol/src/../cil/src/cil_tree.c:307:7
#15 0x555dcfa2bf73 in cil_binary_create_allocated_pdb /src/selinux/libsepol/src/../cil/src/cil_binary.c:4951:8
#16 0x555dcfa2b7a9 in cil_binary_create /src/selinux/libsepol/src/../cil/src/cil_binary.c:4326:7
#17 0x555dcfa073c1 in cil_build_policydb /src/selinux/libsepol/src/../cil/src/cil.c:672:7
#18 0x555dcf9ed982 in LLVMFuzzerTestOneInput /src/secilc-fuzzer.c:62:6
#19 0x555dcf8a2330 in fuzzer::Fuzzer::ExecuteCallback(unsigned char const*, unsigned long) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerLoop.cpp:614:13
#20 0x555dcf88d5a5 in fuzzer::RunOneTest(fuzzer::Fuzzer*, char const*, unsigned long) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:327:6
#21 0x555dcf89303f in fuzzer::FuzzerDriver(int*, char***, int (*)(unsigned char const*, unsigned long)) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:862:9
#22 0x555dcf8be2e2 in main /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerMain.cpp:20:10
#23 0x7f0141a73082 in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x24082) (BuildId: 0702430aef5fa3dda43986563e9ffcc47efbd75e)
DEDUP_TOKEN: __interceptor_malloc--cil_malloc--__cil_typeattr_bitmap_init
SUMMARY: AddressSanitizer: heap-buffer-overflow /src/selinux/libsepol/src/ebitmap.c:235:27 in ebitmap_match_any
Shadow bytes around the buggy address:
0x7f013ccb7500: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
0x7f013ccb7580: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
0x7f013ccb7600: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
0x7f013ccb7680: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
0x7f013ccb7700: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
=>0x7f013ccb7780: fa fa fa fa fa fa fa fa fa fa fa fa fa fa[fa]fa
0x7f013ccb7800: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
0x7f013ccb7880: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
0x7f013ccb7900: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
0x7f013ccb7980: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
0x7f013ccb7a00: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
Shadow byte legend (one shadow byte represents 8 application bytes):
Addressable: 00
Partially addressable: 01 02 03 04 05 06 07
Heap left redzone: fa
Freed heap region: fd
Stack left redzone: f1
Stack mid redzone: f2
Stack right redzone: f3
Stack after return: f5
Stack use after scope: f8
Global redzone: f9
Global init order: f6
Poisoned by user: f7
Container overflow: fc
Array cookie: ac
Intra object redzone: bb
ASan internal: fe
Left alloca redzone: ca
Right alloca redzone: cb
==14==ABORTING
@@ -0,0 +1,4 @@
weight: 1.0
package_name: "selinux"
harness_name: "secilc-fuzzer"
task_id: "selinux-OSV-2021-902-delta-1"
+2 -2
View File
@@ -7,9 +7,9 @@ authors = [{ name = "Trail of Bits", email = "opensource@trailofbits.com" }]
requires-python = ">=3.12,<3.13"
dependencies = [
"common",
"langchain>=0.3.17",
"langchain~=0.3.20",
"langchain-community>=0.3.16",
"langgraph>=0.2.69",
"langgraph==0.3.23",
"openai>=1.60.2",
"program-model",
"pydantic-settings>=2.7.1",
+39 -27
View File
@@ -17,7 +17,6 @@ import re
class PatcherAgentName(Enum):
DIFF_ANALYSIS = "diff_analysis_node"
CONTEXT_RETRIEVER = "context_retriever_node"
ROOT_CAUSE_ANALYSIS = "root_cause_analysis"
CREATE_PATCH = "create_patch"
@@ -34,10 +33,12 @@ class PatcherAgentState(BaseModel):
context: PatchInput
messages: Annotated[list[BaseMessage], add_messages]
relevant_code_snippets: Annotated[list[ContextCodeSnippet], operator.add] = Field(default_factory=list)
relevant_code_snippets: Annotated[set[ContextCodeSnippet], operator.or_] = Field(default_factory=set)
diff_analysis: str | None = None
root_cause: str | None = None
ctx_request_limit: bool = Field(default=False)
patches: list[PatchOutput] = Field(default_factory=list)
patch_tries: int = Field(default=0)
@@ -75,7 +76,7 @@ class PatcherAgentState(BaseModel):
class ContextRetrieverState(BaseModel):
"""State for the Context Retriever Agent."""
code_snippet_requests: list[CodeSnippetKey] = Field(default_factory=list)
code_snippet_requests: list[CodeSnippetRequest] = Field(default_factory=list)
prev_node: str | None = None
@@ -96,6 +97,12 @@ class CodeSnippetKey(BaseModel):
return self.file_path == other.file_path and self.identifier == other.identifier
class CodeSnippetRequest(BaseModel):
"""Code snippet request"""
request: str = Field(description="Detailed explanation of what code snippet is needed")
class ContextCodeSnippet(BaseModel):
"""Code snippet from the Challenge Task. This is the base unit used by the
patcher to build patches. Changes are applied to this units."""
@@ -109,6 +116,29 @@ class ContextCodeSnippet(BaseModel):
code_context: str | None = None
"Additional context around the code snippet, e.g. lines information, etc."
def __str__(self) -> str:
context = (
f"""
<code_context>
{self.code_context}
</code_context>
"""
if self.code_context
else ""
)
return f"""<code_snippet>
<file_path>{self.key.file_path}</file_path>
<identifier>{self.key.identifier}</identifier>
<code>
{self.code}
</code>{context}
</code_snippet>
"""
def __hash__(self) -> int:
"""Hash the code snippet"""
return hash((type(self),) + tuple(self.__dict__.values()))
@dataclass
class PatcherAgentBase:
@@ -142,31 +172,19 @@ class PatcherAgentBase:
return path
def get_code_snippet_requests(
self, response: str, update_state: dict, *, current_node: str, default_goto: str
self, response: str, update_state: dict, ctx_request_limit: bool, *, current_node: str, default_goto: str
) -> tuple[str, dict]:
"""Get the code snippet request from the response."""
CODE_SNIPPET_REQUEST = re.compile(
r"""# CODE SNIPPET REQUESTS:
.*?
(File path: (.*?)
Identifier: (.*?))+
```""",
re.DOTALL | re.IGNORECASE,
)
matches = CODE_SNIPPET_REQUEST.search(response)
if not matches:
if ctx_request_limit:
return default_goto, update_state
# Extract all file path and identifier pairs
pair_pattern = re.compile(r"File path: (.*?)\nIdentifier: (.*?)(?:\n|$)", re.DOTALL)
code_snippet_requests = pair_pattern.findall(matches.group(0))
if len(code_snippet_requests) == 0:
CODE_SNIPPET_REQUEST_RE = re.compile("<code_request>(.*?)</code_request>", re.DOTALL | re.IGNORECASE)
code_snippet_requests_matches = CODE_SNIPPET_REQUEST_RE.findall(response)
if not code_snippet_requests_matches:
return default_goto, update_state
code_snippet_requests = [
CodeSnippetKey(file_path=request[0], identifier=request[1]) for request in code_snippet_requests
CodeSnippetRequest(request=request.strip()) for request in code_snippet_requests_matches
]
return PatcherAgentName.CONTEXT_RETRIEVER.value, {
"code_snippet_requests": code_snippet_requests,
@@ -183,12 +201,6 @@ CONTEXT_DIFF_TMPL = """Diff introducing the vulnerability:
{diff_content}
```
"""
CONTEXT_DIFF_ANALYSIS_TMPL = """Analysis of the diff introducing the \
vulnerability.
```
{diff_analysis}
```
"""
CONTEXT_ROOT_CAUSE_TMPL = """Vulnerability Root cause analysis:
```
{root_cause}
@@ -1,10 +1,27 @@
"""Agent that retrieves code snippets from the project."""
import logging
import re
from dataclasses import dataclass, field
from pathlib import Path
from __future__ import annotations
import logging
import langgraph.errors
import operator
from dataclasses import dataclass, field
from langchain_core.messages import ToolMessage
from typing import Annotated, Sequence, Literal
from langgraph.managed import IsLastStep, RemainingSteps
from pydantic import BaseModel, Field
from pathlib import Path
from enum import Enum
from langchain_core.tools import BaseTool
from langgraph.graph import StateGraph, add_messages
from langgraph.prebuilt import ToolNode
from langgraph.constants import END
from langchain_core.tools import tool
from langgraph.prebuilt import InjectedState
from langchain_core.tools.base import InjectedToolCallId
from buttercup.common.challenge_task import ChallengeTask
from langchain_core.messages import AnyMessage, BaseMessage
from langchain_core.runnables import Runnable
from langchain_core.prompts import ChatPromptTemplate
from buttercup.patcher.agents.common import PatcherAgentBase, ContextRetrieverState, ContextCodeSnippet, CodeSnippetKey
@@ -12,23 +29,236 @@ from buttercup.common.llm import ButtercupLLM, create_default_llm
from langgraph.types import Command
import subprocess
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent, ToolNode
from buttercup.program_model.codequery import CodeQueryPersistent
logger = logging.getLogger(__name__)
SYSTEM_TMPL = """You are a software engineer. Your job is to retrieve the code snippets requested by the user.
You must use the tools provided to you to retrieve the code snippets.
SYSTEM_TMPL = """You are an AI assistant tasked with helping a software engineer find and extract relevant code snippets from a project."""
Do not stop until you have retrieved the code definition of the function/type/class/etc. requested by the user.
Use `get_function_definition`/`get_type_definition` as the last tool in your chain of calls to actually retrieve the code definition.
USER_MSG = """You have access to some tools to navigate the project and search for code.
The project you will be working with is located at:
<project_name>
{PROJECT_NAME}
</project_name>
The software engineer has made the following request:
<engineer_request>
{REQUEST}
</engineer_request>
Throughout this process, maintain a <scratchpad> where you document your thought process, the commands you're using, and the results you're finding. This will help you keep track of your progress and make informed decisions about next steps.
Do not make up any information, only use the provided tools and the information available in the project.
Do not make up any file paths.
Remember to use the provided tools only as defined, and do not attempt to modify or extend their functionality. If you encounter any errors or cannot find the requested information, explain the issue in your answer and suggest potential next steps or alternative approaches.
Try to use `get_function_definition` and `get_type_definition` tools as much as possible and rely on others only if these tools fail or do not work as expected.
Answer with <END> when you have found all the code snippets requested.
"""
MESSAGES = ChatPromptTemplate.from_messages(
[("human", "Please retrieve the code snippet for {file_path} | {function_name}")]
)
RECURSION_LIMIT = 20
MESSAGES = ChatPromptTemplate.from_messages([("human", "Please satisfy this request: {request}")])
class State(BaseModel):
"""State for the context retriever agent."""
messages: Annotated[Sequence[BaseMessage], add_messages]
is_last_step: IsLastStep | None = None
remaining_steps: RemainingSteps | None = None
challenge: ChallengeTask
context_retriever_agent: ContextRetrieverAgent
project_name: str
request: str
code_snippets: Annotated[list[ContextCodeSnippet], operator.add] = Field(default_factory=list)
tmp_code_snippets: TmpCodeSnippets | None = None
class NodeNames(str, Enum):
"""Names of the nodes in the state graph."""
AGENT = "agent"
TOOLS = "tools"
@tool
def ls(
path: str, state: Annotated[BaseModel, InjectedState], tool_call_id: Annotated[str, InjectedToolCallId]
) -> Command:
"""List the files in the given path in the project's source directory."""
path = state.context_retriever_agent.rebase_src_path(path)
if not state.context_retriever_agent.challenge.task_dir.joinpath(path).exists():
raise ValueError(f"Path {path} does not exist in the project's directory")
logger.info("Listing files in %s", path)
args = ["ls", "-l"]
if path:
args.append(str(path))
ls_output = subprocess.check_output(args, cwd=state.challenge.task_dir).decode("utf-8")
return Command(update={"messages": [ToolMessage(ls_output, tool_call_id=tool_call_id)]})
@tool
def grep(
pattern: str,
path: str | None,
state: Annotated[BaseModel, InjectedState],
tool_call_id: Annotated[str, InjectedToolCallId],
) -> Command:
"""Grep for a string and return a 5-line context around the match, together with line numbers. If no path is provided, search the entire project. Prefer using this tool over cat."""
if path:
path = state.context_retriever_agent.rebase_src_path(path)
if not state.challenge.task_dir.joinpath(path).exists():
raise ValueError(f"Path {path} does not exist in the project's directory")
logger.info("Searching for %s in %s", pattern, path)
args = ["grep", "-C", "5", "-nr", pattern]
if path:
args.append(str(path))
grep_output = subprocess.check_output(args, cwd=state.challenge.task_dir).decode("utf-8")
return Command(update={"messages": [ToolMessage(grep_output, tool_call_id=tool_call_id)]})
@tool
def cat(
path: str, state: Annotated[BaseModel, InjectedState], tool_call_id: Annotated[str, InjectedToolCallId]
) -> Command:
"""Read the contents of a file. Use this tool only if grep and get_lines do not work as it might return a large amount of text."""
path = state.context_retriever_agent.rebase_src_path(path)
if not state.challenge.task_dir.joinpath(path).exists():
raise ValueError(f"Path {path} does not exist in the project's directory")
logger.info("Reading contents of %s", path)
cat_output = state.challenge.task_dir.joinpath(path).read_text()
return Command(update={"messages": [ToolMessage(cat_output, tool_call_id=tool_call_id)]})
@tool
def get_lines(
path: str,
start: int,
end: int,
state: Annotated[BaseModel, InjectedState],
tool_call_id: Annotated[str, InjectedToolCallId],
) -> Command:
"""Get a range of lines from a file. Prefer using this tool over cat."""
path = state.context_retriever_agent.rebase_src_path(path)
if not state.challenge.task_dir.joinpath(path).exists():
raise ValueError(f"Path {path} does not exist in the project's directory")
logger.info("Getting lines %d-%d of %s", start, end, path)
get_lines_output = "\n".join(state.challenge.task_dir.joinpath(path).read_text().splitlines()[start:end])
return Command(update={"messages": [ToolMessage(get_lines_output, tool_call_id=tool_call_id)]})
@tool
def get_function_definition(
function_name: str,
file_path: str | None,
state: Annotated[BaseModel, InjectedState],
tool_call_id: Annotated[str, InjectedToolCallId],
) -> Command:
"""Get the definition of a function. If available, pass a file_path, \
otherwise pass None. Use this when you want to get information about a \
function. If not sure about the file path, pass None. Prefer using this \
tool over any other and rely on others only if this tool fails or does \
not work. This tool is just going to return a message whether it found \
the function or not, but it won't provide the code snippet directly. The \
function should be considered as retrieved anyway."""
if file_path:
path = Path(state.context_retriever_agent.rebase_src_path(file_path))
else:
path = None
logger.info("Getting function definition of %s in %s", function_name, path)
functions = state.context_retriever_agent.codequery.get_functions(function_name, path)
if not functions:
functions = state.context_retriever_agent.codequery.get_functions(function_name, path, fuzzy=True)
if not functions:
return Command(
update={"messages": [ToolMessage("No definition found for function", tool_call_id=tool_call_id)]}
)
code_snippets = [
ContextCodeSnippet(
key=CodeSnippetKey(
file_path=function.file_path.as_posix(),
identifier=function.name,
),
code=body.body,
)
for function in functions
for body in function.bodies
]
state.tmp_code_snippets.code_snippets.extend(code_snippets)
return Command(
update={
"messages": [
ToolMessage(
f"Found {len(code_snippets)} code snippets for function {function_name}", tool_call_id=tool_call_id
)
],
"code_snippets": code_snippets,
}
)
@tool
def get_type_definition(
type_name: str,
file_path: str | None,
state: Annotated[BaseModel, InjectedState],
tool_call_id: Annotated[str, InjectedToolCallId],
) -> Command:
"""Get the definition of a type. If available, pass a file_path, \
otherwise pass None. Use this when you want to get information about a type. \
If not sure about the file path, pass None. Prefer using this tool over any \
other and \
rely on others only if this tool fails or does not work. This tool is just \
going to return a message whether it found the type or not, but it won't \
provide the code snippet directly. The type should be considered as \
retrieved anyway."""
if file_path:
path = Path(state.context_retriever_agent.rebase_src_path(file_path))
else:
path = None
logger.info("Getting type definition of %s in %s", type_name, path)
types = state.context_retriever_agent.codequery.get_types(type_name, path)
if not types:
types = state.context_retriever_agent.codequery.get_types(type_name, path, fuzzy=True)
if not types:
return Command(
update={"messages": [ToolMessage("No definition found for type", tool_call_id=tool_call_id)]}
)
code_snippets = [
ContextCodeSnippet(
key=CodeSnippetKey(
file_path=path.as_posix() if path else "<type-path>", # TODO: get this from the type definition
identifier=type_def.name,
),
code=type_def.definition,
)
for type_def in types
]
state.tmp_code_snippets.code_snippets.extend(code_snippets)
return Command(
update={
"messages": [
ToolMessage(f"Found {len(code_snippets)} code snippets for type {type_name}", tool_call_id=tool_call_id)
],
"code_snippets": code_snippets,
}
)
class TmpCodeSnippets(BaseModel):
"""Temporary code snippets."""
code_snippets: list[ContextCodeSnippet] = Field(default_factory=list)
@dataclass
@@ -39,7 +269,10 @@ class ContextRetrieverAgent(PatcherAgentBase):
max_retries: int = 30
recursion_limit: int = 80
llm: Runnable = field(init=False)
max_retries: int = 30
current_retries: int = field(init=False, default=0)
codequery: CodeQueryPersistent = field(init=False)
tools: list[BaseTool] = field(init=False)
def __post_init__(self) -> None:
"""Initialize a few fields"""
@@ -47,114 +280,8 @@ class ContextRetrieverAgent(PatcherAgentBase):
fallback_llms = [
create_default_llm(model_name=ButtercupLLM.CLAUDE_3_5_SONNET.value),
]
self.llm = default_llm.with_fallbacks(fallback_llms)
def _parse_code_snippet_msg(self, msg: str) -> tuple[str, str, str]:
"""Parse the code snippet message."""
# Extract code part from the message using regex
code_pattern = re.compile(r"File path:.*?\nIdentifier:.*?\nCode:\n(.*?)$", re.DOTALL)
code_match = code_pattern.search(msg)
if code_match:
msg = code_match.group(1).strip()
code_block_pattern = re.compile(r"```(?:[a-z]+)?\s*(.*?)\s*```", re.DOTALL)
code_block_match = code_block_pattern.search(msg)
if code_block_match:
# Remove the code block markers
msg = code_block_match.group(1).strip()
return msg
def retrieve_context(self, state: ContextRetrieverState) -> Command:
"""Retrieve the context for the diff analysis."""
if self.current_retries >= self.max_retries:
logger.warning("Reached max context retrieval retries, skipping")
return Command(
goto=state.prev_node,
)
logger.info("Retrieving the context for the diff analysis in Challenge Task %s", self.challenge.name)
logger.debug("Code snippet requests: %s", state.code_snippet_requests)
codequery = CodeQueryPersistent(self.challenge, work_dir=self.work_dir)
@tool
def ls(path: str) -> str:
"""List the files in the given path in the project's source directory."""
path = self.rebase_src_path(path)
if not self.challenge.task_dir.joinpath(path).exists():
raise ValueError(f"File {path} does not exist")
logger.info("Listing files in %s", path)
return subprocess.check_output(["ls", "-l", str(path)], cwd=self.challenge.task_dir).decode("utf-8")
@tool
def grep(path: str, pattern: str) -> str:
"""Grep for a string in a file. Prefer using this tool over cat."""
path = self.rebase_src_path(path)
if not self.challenge.task_dir.joinpath(path).exists():
raise ValueError(f"File {path} does not exist")
logger.info("Searching for %s in %s", pattern, path)
return subprocess.check_output(
["grep", "-C", "10", "-nr", pattern, str(path)], cwd=self.challenge.task_dir
).decode("utf-8")
@tool
def cat(path: str) -> str:
"""Read the contents of a file. Use this tool only if grep and get_lines do not work as it might return a large amount of text."""
path = self.rebase_src_path(path)
if not self.challenge.task_dir.joinpath(path).exists():
raise ValueError(f"File {path} does not exist")
logger.info("Reading contents of %s", path)
return self.challenge.task_dir.joinpath(path).read_text()
@tool
def get_lines(path: str, start: int, end: int) -> str:
"""Get a range of lines from a file. Prefer using this tool over cat."""
path = self.rebase_src_path(path)
logger.info("Getting lines %d-%d of %s", start, end, path)
return "\n".join(self.challenge.task_dir.joinpath(path).read_text().splitlines()[start:end])
@tool
def get_function_definition(function_name: str, path: str | None = None) -> str:
"""Get the definition of a function. This tool, when called, is the last call in your chain of calls."""
if path:
path = self.rebase_src_path(path)
if not self.challenge.task_dir.joinpath(path).exists():
raise ValueError(f"File {path} does not exist")
logger.info("Getting function definition of %s in %s", function_name, path)
functions = codequery.get_functions(function_name, path)
if not functions:
functions = codequery.get_functions(function_name, path, fuzzy=True)
if not functions:
raise ValueError(f"No definition found for function {function_name} in {path}")
# TODO: allow for multiple bodies
return functions[0].bodies[0].body
@tool
def get_type_definition(type_name: str, path: str | None = None) -> str:
"""Get the definition of a type. This tool, when called, is the last call in your chain of calls."""
if path:
path = self.rebase_src_path(path)
if not self.challenge.task_dir.joinpath(path).exists():
raise ValueError(f"File {path} does not exist")
logger.info("Getting type definition of %s in %s", type_name, path)
types = codequery.get_types(type_name, path)
if not types:
types = codequery.get_types(type_name, path, fuzzy=True)
if not types:
raise ValueError(f"No definition found for type {type_name} in {path}")
# TODO: allow for multiple definitions
return types[0].definition
tools = [
self.tools = [
ls,
grep,
get_lines,
@@ -162,40 +289,115 @@ class ContextRetrieverAgent(PatcherAgentBase):
get_function_definition,
get_type_definition,
]
tool_node = ToolNode(tools, handle_tool_errors=True)
agent = create_react_agent(
self.llm,
tool_node,
prompt=SYSTEM_TMPL,
self.llm = default_llm.with_fallbacks(fallback_llms)
self.llm_with_tools = self.llm.bind_tools(self.tools)
self.codequery = CodeQueryPersistent(self.challenge, work_dir=self.work_dir)
def _call_model(self, state: State) -> Command[Literal[NodeNames.TOOLS.value, END]]:
"""Call the model."""
user_msg = USER_MSG.format(
REQUEST=state.request,
PROJECT_NAME=state.project_name,
)
res = self.llm_with_tools.invoke(
[
{"role": "system", "content": SYSTEM_TMPL},
{"role": "user", "content": user_msg},
*state.messages,
]
)
if "<END>" in res.content:
return Command(update={"messages": [res]}, goto=END)
return Command(update={"messages": [res]}, goto=NodeNames.TOOLS.value)
def _prompt(self, state: State) -> list[AnyMessage]:
return [
("system", SYSTEM_TMPL),
(
"user",
USER_MSG.format(
REQUEST=state.request,
PROJECT_NAME=state.project_name,
),
),
*state.messages,
]
def _create_agent(self) -> Runnable:
"""Create the agent."""
workflow = StateGraph(State)
tool_node = ToolNode(self.tools, name=NodeNames.TOOLS.value)
workflow.add_node(NodeNames.AGENT.value, self._call_model)
workflow.add_node(NodeNames.TOOLS.value, tool_node)
workflow.set_entry_point(NodeNames.AGENT.value)
workflow.add_edge(NodeNames.TOOLS.value, NodeNames.AGENT.value)
return workflow.compile()
def retrieve_context(self, state: ContextRetrieverState) -> Command:
"""Retrieve the context for the diff analysis."""
if self.current_retries >= self.max_retries:
logger.warning("Reached max context retrieval retries, skipping")
return Command(
update={
"ctx_request_limit": True,
},
goto=state.prev_node,
)
logger.info("Retrieving the context for the diff analysis in Challenge Task %s", self.challenge.name)
logger.debug("Code snippet requests: %s", state.code_snippet_requests)
agent = self._create_agent()
res = []
for request in state.code_snippet_requests:
logger.info("Retrieving code snippet for %s | %s", request.file_path, request.identifier)
messages = MESSAGES.invoke({"file_path": request.file_path, "function_name": request.identifier})
snippet = agent.invoke({"messages": messages.to_messages()}, {"recursion_limit": self.recursion_limit})
logger.info("Code snippet retrieved for %s | %s", request.file_path, request.identifier)
if not snippet["messages"]:
raise RuntimeError("No messages returned from the agent")
logger.info("Retrieving code snippet for request '%s'", request.request)
tmp_code_snippets = TmpCodeSnippets()
input_state = {
"request": request.request,
"project_name": self.challenge.project_name,
"challenge": self.challenge,
"context_retriever_agent": self,
"tmp_code_snippets": tmp_code_snippets,
}
config = {
"configurable": {
"thread_id": hash(request.request),
},
}
try:
ctx_state_dict: dict = agent.invoke(input_state, {"recursion_limit": RECURSION_LIMIT, **config})
except langgraph.errors.GraphRecursionError:
logger.error("Reached recursion limit for request '%s'", request.request)
ctx_state_dict = {
"messages": [],
"challenge": self.challenge,
"project_name": self.challenge.project_name,
"context_retriever_agent": self,
"request": request.request,
"code_snippets": tmp_code_snippets.code_snippets,
}
msg = snippet["messages"][-1].content
code = self._parse_code_snippet_msg(msg)
if logger.level <= logging.DEBUG:
for message in ctx_state_dict["messages"]:
string_representation = f"{message.type.upper()}: {message.content}"
logger.debug(string_representation)
res.append(
ContextCodeSnippet(
key=CodeSnippetKey(
file_path=request.file_path,
identifier=request.identifier,
),
code=code,
code_context="",
)
)
ctx_state = State(**ctx_state_dict)
if not ctx_state.code_snippets:
# TODO: pass back these errors to the caller so that it can retry in other ways
logger.warning("No code snippet returned from the agent for request '%s'", request.request)
continue
logger.info("Code snippets retrieved for request '%s'", request.request)
res.extend(ctx_state.code_snippets)
self.current_retries += 1
return Command(
update={
"relevant_code_snippets": res,
"relevant_code_snippets": set(res),
},
goto=state.prev_node,
)
@@ -63,7 +63,6 @@ class PatcherLeaderAgent:
)
workflow = StateGraph(PatcherAgentState)
workflow.add_node(PatcherAgentName.DIFF_ANALYSIS.value, rootcause_agent.diff_analysis)
workflow.add_node(PatcherAgentName.CONTEXT_RETRIEVER.value, context_retriever_agent.retrieve_context)
workflow.add_node(PatcherAgentName.ROOT_CAUSE_ANALYSIS.value, rootcause_agent.analyze_vulnerability)
workflow.add_node(PatcherAgentName.CREATE_PATCH.value, swe_agent.create_patch_node)
@@ -73,7 +72,7 @@ class PatcherLeaderAgent:
workflow.add_node(PatcherAgentName.RUN_POV.value, qe_agent.run_pov_node)
workflow.add_node(PatcherAgentName.RUN_TESTS.value, qe_agent.run_tests_node)
workflow.set_entry_point(PatcherAgentName.DIFF_ANALYSIS.value)
workflow.set_entry_point(PatcherAgentName.ROOT_CAUSE_ANALYSIS.value)
return workflow
def run_patch_task(self) -> PatchOutput | None:
+4 -4
View File
@@ -266,7 +266,7 @@ class QEAgent(PatcherAgentBase):
def run_pov_node(
self, state: PatcherAgentState
) -> Command[Literal[PatcherAgentName.RUN_TESTS.value, PatcherAgentName.DIFF_ANALYSIS.value]]:
) -> Command[Literal[PatcherAgentName.RUN_TESTS.value, PatcherAgentName.ROOT_CAUSE_ANALYSIS.value]]:
"""Node in the LangGraph that runs a PoV against a currently built patch"""
logger.info("Testing PoV on Challenge Task %s rebuilt with patch", self.challenge.name)
try:
@@ -288,7 +288,7 @@ class QEAgent(PatcherAgentBase):
"pov_stdout": exc.stdout,
"pov_stderr": exc.stderr,
},
goto=PatcherAgentName.DIFF_ANALYSIS.value,
goto=PatcherAgentName.ROOT_CAUSE_ANALYSIS.value,
)
if not pov_output.command_result.success:
@@ -299,7 +299,7 @@ class QEAgent(PatcherAgentBase):
"pov_stdout": pov_output.command_result.output,
"pov_stderr": pov_output.command_result.error,
},
goto=PatcherAgentName.DIFF_ANALYSIS.value,
goto=PatcherAgentName.ROOT_CAUSE_ANALYSIS.value,
)
logger.info(
@@ -320,7 +320,7 @@ class QEAgent(PatcherAgentBase):
},
goto=PatcherAgentName.RUN_TESTS.value
if not pov_output.did_crash()
else PatcherAgentName.DIFF_ANALYSIS.value,
else PatcherAgentName.ROOT_CAUSE_ANALYSIS.value,
)
def run_tests_node(self, state: PatcherAgentState) -> Command[Literal[PatcherAgentName.CREATE_PATCH.value, END]]:
+124 -253
View File
@@ -1,173 +1,152 @@
"""Software Engineer LLM agent, analyzing the root cause of a vulnerability."""
import logging
import re
from typing import Literal
from dataclasses import dataclass, field
from operator import itemgetter
import tiktoken
from langchain_core.messages import BaseMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import (
ChatPromptTemplate,
MessagesPlaceholder,
)
from langchain_core.runnables import Runnable, RunnableBranch
from langchain_core.messages import BaseMessage
from langchain_core.runnables import Runnable
from buttercup.patcher.agents.common import (
CONTEXT_CODE_SNIPPET_TMPL,
CONTEXT_DIFF_ANALYSIS_TMPL,
CONTEXT_DIFF_TMPL,
CONTEXT_PROJECT_TMPL,
CONTEXT_SANITIZER_TMPL,
PatcherAgentState,
PatcherAgentName,
PatcherAgentBase,
CONTEXT_CODE_SNIPPET_TMPL,
ContextCodeSnippet,
get_code_snippet_request_tmpl,
)
from buttercup.common.llm import ButtercupLLM, create_default_llm
from buttercup.patcher.utils import decode_bytes, get_diff_content
from buttercup.patcher.utils import decode_bytes
from langgraph.types import Command
logger = logging.getLogger(__name__)
ROOT_CAUSE_SYSTEM_TMPL = """You are a security engineer. Your job is to analyze the \
current project's code, a vulnerability description, and determine its root \
cause. The vulnerability is introduced/enabled in the specified diff and it was not \
present before that. Read the project's code and understand the issue.
ROOT_CAUSE_SYSTEM_MSG = """You are an experienced security engineer tasked with determining the root cause of a specific vulnerability. Your analysis will be used to patch the vulnerability."""
ROOT_CAUSE_USER_MSG = """Your goal is to provide a highly technical, detailed, and focused analysis of the vulnerability, going to the root cause of the issue.
You must NOT:
- provide generic recommendations
- provide overly generic root causes
- talk about possible issues that existed before the diff.
- provide code suggestions on how to fix the vulnerability
- make up code, contexts or information
Here is the information you need to analyze:
You MUST:
- provide a detailed root cause analysis of the vulnerability, referencing actual code snippets
- describe a concrete and specific root cause
- focus ONLY on the issue introduced in the diff and ONLY on the issue \
reported in the sanitizer output. If there are multiple vulnerabilities in the \
diff, focus on the one which triggers the sanitizer
- look at the actual code snippets and the code context to understand the root cause, \
do not make up code, contexts or information and do not assume anything
1. Diff introducing the vulnerability (if available):
<diff>
{DIFF}
</diff>
If you need more context to understand the code, ask for more code snippets. Try \
to fully understand the issue by asking for more code snippets instead of \
rushing through an answer. For example, before looking at the code, review the \
fuzzer harness code to understand how the project's code is being used.
2. Project Name:
<project_name>
{PROJECT_NAME}
</project_name>
You can request multiple code snippets at once. Request code snippets in the
following way:
3. Sanitizer used:
<sanitizer>
{SANITIZER}
</sanitizer>
{code_snippet_request_tmpl}
4. Sanitizer output:
<sanitizer_output>
{SANITIZER_OUTPUT}
</sanitizer_output>
You must use the above format to request code snippets.
5. Code snippets from the project:
<code_snippets>
{CODE_SNIPPETS}
</code_snippets>{OLD_ROOT_CAUSE}
Instructions:
1. Review all the provided information carefully.
2. Conduct a thorough analysis of the vulnerability. Focus solely on the issue reported in the sanitizer output (and introduced by the diff if available).
3. Structure your response as follows:
<vulnerability_analysis_process>
[Break down your thought process using this structure:
1. Diff Analysis:
- Quote each specific change in the diff relevant to the vulnerability.
- Explain what each of those changes does to the code.
2. Sanitizer Output Analysis:
- Quote each relevant part of the sanitizer output.
- Interpret what each part means in terms of potential vulnerabilities.
3. Code Snippet Analysis:
- Quote the relevant parts of the provided code snippets.
- Explain how these parts relate to the changes in the diff and the sanitizer output.
4. Integrated Analysis:
- Connect the findings from the diff, sanitizer output, and code snippets.
- Explain how these components work together to introduce the vulnerability.
- Use the reference numbers from previous steps to clearly link your observations.
5. Previous Root Cause Analysis:
- If a previous root cause analysis is available, try to critically analyze it and see if it's still valid, and provide your own extended/improved analysis.
6. Vulnerability Identification:
- Based on the integrated analysis, explain how the changes introduced the vulnerability.
- Use technical details and security concepts to support your conclusion.
7. Determine if additional code snippets need to be requested:
- If you need additional context to understand the code, wrap your code request in <code_request> tags. Be specific about what code you need and why. If there is a previous root cause analysis, try to explore more parts of the code to better understand the vulnerability. Request example:
<code_requests>
<code_request>
Full implementation of the function 'validate_input()' from the file 'input_validation.c', as it's referenced in the diff but not fully visible.
</code_request>
<code_request>
...
</code_request>
</code_requests>
Ensure that you reference and analyze all components (code snippets, diff, and sanitizer output) before drawing any conclusions. Request additional code snippets if necessary for a complete understanding. Remember to quote relevant parts of the code, diff, and sanitizer output throughout your analysis.]
</vulnerability_analysis_process>
<analysis>
[Provide a detailed, technical explanation of the vulnerability. Include:
- The exact nature of the vulnerability
- How it was introduced in the diff
- Relevant technical details about how the vulnerability could be exploited
Ground your explanation in the provided code snippets, diff, and sanitizer output.]
</analysis>
<summary>
[Provide a clear, concise description of the vulnerability introduced by the diff. Highlight the most critical aspects of the security issue.]
</summary>
4. In your analysis:
- Provide a detailed, technical explanation of the vulnerability.
- Describe a concrete and specific problem.
- If there are multiple vulnerabilities in the diff, focus on the one which triggers the sanitizer.
- Base your analysis on the provided code snippets, diff, and sanitizer output.
- Request and analyze additional code snippets if necessary for a complete understanding.
5. Do NOT:
- Provide generic recommendations.
- Offer overly generic analyses.
- Discuss possible issues that existed before the diff.
- Suggest code fixes for the vulnerability.
- Make up any information not present in the provided data.
Maintain a highly technical approach throughout your analysis, focusing solely on the vulnerability introduced by the diff and avoiding any generic advice or recommendations for fixing the code.
"""
ROOT_CAUSE_SYSTEM_TMPL = ROOT_CAUSE_SYSTEM_TMPL.format(
code_snippet_request_tmpl=get_code_snippet_request_tmpl(2),
)
DIFF_ANALYSIS_SYSTEM_TMPL = """You are a security engineer. Your job is to \
analyze a diff of a project and determine the vulnerability it introduces. The \
vulnerability is introduced in the specified diff and it was not present \
before that.
You MUST:
- provide a detailed analysis of the vulnerability introduced in the diff
- describe a concrete and specific problem
- focus ONLY on the issue introduced in the diff and ONLY on the issue \
reported in the sanitizer output. If there are multiple vulnerabilities in the \
diff, focus on the one which triggers the sanitizer
You must NOT:
- provide generic recommendations
- provide overly generic analyses
- talk about possible issues that existed before the diff.
- provide code suggestions on how to fix the vulnerability
If you need more context to understand the code, you can ask for more code
snippets. Try to fully understand the issue by asking for more code snippets
instead of rushing through an answer.
You can request multiple code snippets at once. Request code snippets in the
following way:
{code_snippet_request_tmpl}
"""
DIFF_ANALYSIS_SYSTEM_TMPL = DIFF_ANALYSIS_SYSTEM_TMPL.format(
code_snippet_request_tmpl=get_code_snippet_request_tmpl(2),
)
ROOT_CAUSE_PROMPT = ChatPromptTemplate.from_messages(
[
("system", ROOT_CAUSE_SYSTEM_TMPL),
MessagesPlaceholder(variable_name="context"),
MessagesPlaceholder(variable_name="messages", optional=True),
("user", "Provide a detailed and complete analysis"),
("system", ROOT_CAUSE_SYSTEM_MSG),
("user", ROOT_CAUSE_USER_MSG),
("ai", "<vulnerability_analysis_process>"),
]
)
DIFF_ANALYSIS_PROMPT = ChatPromptTemplate.from_messages(
[
("system", DIFF_ANALYSIS_SYSTEM_TMPL),
MessagesPlaceholder(variable_name="context"),
MessagesPlaceholder(variable_name="messages", optional=True),
("user", "Provide a detailed and complete analysis"),
]
)
OLD_ROOT_CAUSE_TMPL = """
OLD_DIFF_ANALYSIS_PROMPT = ChatPromptTemplate.from_messages(
[
(
"ai",
"""Diff analysis:
```
{diff_analysis}
```""",
),
(
"user",
"I tried to fix the vulnerability according to the diff analysis you \
provided, but the Proof of Vulnerability (PoV) was still triggered. Please \
review your previous analysis and provide a new or an improved analysis to \
help me correctly fix the vulnerability.",
),
]
)
6. Existing root cause analysis (potentially outdated, wrong, or incomplete):
OLD_ROOT_CAUSE_PROMPT = ChatPromptTemplate.from_messages(
[
(
"ai",
"""Root cause analysis:
```
<old_root_cause>
{root_cause}
```""",
),
(
"user",
"I tried to fix the vulnerability according to the root cause you \
provided, but the Proof of Vulnerability (PoV) was still triggered. Please \
review your previous analysis and provide a new or an improved analysis to \
help me correctly fix the vulnerability.",
),
]
)
REDUCE_ANALYSES = ChatPromptTemplate.from_messages(
[
(
"system",
"Please reduce the analyses to a single, detailed and \
complete, analysis that can be easily understood by another software \
engineer without other context.",
),
("user", "Previous analysis:\n```\n{previous_analysis}\n```"),
("user", "New analysis:\n```\n{new_analysis}\n```"),
]
)
</old_root_cause>"""
BUILD_FAILURE_SYSTEM_TMPL = """\
You are a software engineer helping fixing a bug in a \
@@ -200,10 +179,7 @@ class RootCauseAgent(PatcherAgentBase):
llm: Runnable = field(init=False)
root_cause_chain: Runnable = field(init=False)
diff_analysis_chain: Runnable = field(init=False)
diff_analysis_one_chain: Runnable = field(init=False)
build_failure_analysis_chain: Runnable = field(init=False)
encoding: tiktoken.Encoding = field(init=False)
def __post_init__(self) -> None:
"""Initialize a few fields"""
@@ -214,119 +190,10 @@ class RootCauseAgent(PatcherAgentBase):
self.llm = default_llm.with_fallbacks(fallback_llms)
self.root_cause_chain = ROOT_CAUSE_PROMPT | self.llm | StrOutputParser()
self.diff_analysis_chain = DIFF_ANALYSIS_PROMPT | self.llm | StrOutputParser()
self.diff_analysis_one_chain = RunnableBranch(
(
lambda x: x["diff_analysis"],
{
"previous_analysis": itemgetter("diff_analysis"),
"new_analysis": self.diff_analysis_chain,
}
| REDUCE_ANALYSES
| self.llm
| StrOutputParser(),
),
self.diff_analysis_chain,
)
self.build_failure_analysis_chain = BUILD_FAILURE_PROMPT | self.llm | StrOutputParser()
self.encoding = tiktoken.encoding_for_model("gpt-4o")
def get_diff_analysis_context(self, state: PatcherAgentState) -> list[BaseMessage | str]:
"""Get the messages for the diff analysis context."""
messages: list[BaseMessage | str] = []
messages += [CONTEXT_PROJECT_TMPL.format(project_name=self.challenge.name)]
diff_content = get_diff_content(self.challenge)
messages += [CONTEXT_DIFF_TMPL.format(diff_content=diff_content)]
if state.context.sanitizer and state.context.sanitizer_output:
messages += [
CONTEXT_SANITIZER_TMPL.format(
sanitizer=state.context.sanitizer,
sanitizer_output=state.context.sanitizer_output,
)
]
return messages
def get_root_cause_context(self, state: PatcherAgentState) -> list[BaseMessage | str]:
"""Get the messages for the root cause analysis context."""
messages: list[BaseMessage | str] = []
messages += [CONTEXT_PROJECT_TMPL.format(project_name=self.challenge.name)]
if state.diff_analysis:
messages += [CONTEXT_DIFF_ANALYSIS_TMPL.format(diff_analysis=state.diff_analysis)]
if state.context.sanitizer and state.context.sanitizer_output:
messages += [
CONTEXT_SANITIZER_TMPL.format(
sanitizer=state.context.sanitizer,
sanitizer_output=state.context.sanitizer_output,
)
]
messages += self._get_relevant_code_snippets_msgs(state.relevant_code_snippets)
return messages
def diff_analysis(
self, state: PatcherAgentState
) -> Command[Literal[PatcherAgentName.ROOT_CAUSE_ANALYSIS.value, PatcherAgentName.CONTEXT_RETRIEVER.value]]:
"""Analyze the diff to understand the vulnerability."""
logger.info("Analyzing the diff in Challenge Task %s", self.challenge.name)
messages = []
if state.diff_analysis:
messages += OLD_DIFF_ANALYSIS_PROMPT.format_messages(diff_analysis=state.diff_analysis)
diff_analysis = self.chain_call(
lambda x, y: x + y,
self.diff_analysis_one_chain,
{
"context": self.get_diff_analysis_context(state),
"messages": messages,
"diff_analysis": state.diff_analysis,
},
default="",
)
if not diff_analysis:
logger.error("Could not analyze the diff")
raise ValueError("Could not analyze the diff")
update_state = {
"diff_analysis": diff_analysis,
"root_cause": None,
}
goto, update_state = self.get_code_snippet_requests(
diff_analysis,
update_state,
current_node=PatcherAgentName.DIFF_ANALYSIS.value,
default_goto=PatcherAgentName.ROOT_CAUSE_ANALYSIS.value,
)
return Command(
update=update_state,
goto=goto,
)
def _parse_code_snippet_msg(self, msg: str) -> tuple[str, str, str]:
"""Parse the code snippet message."""
# Extract code part from the message using regex
code_pattern = re.compile(r"File path:.*?\nIdentifier:.*?\nCode:\n(.*?)$", re.DOTALL)
code_match = code_pattern.search(msg)
if code_match:
code_block_pattern = re.compile(r"```(?:[a-z]+)?\s*(.*?)\s*```", re.DOTALL)
code_block_match = code_block_pattern.search(code_match.group(1))
if code_block_match:
# Remove the code block markers
msg = code_block_match.group(1).strip()
else:
# If we can't find the code block, just return the whole part after "Code:"
msg = code_match.group(1).strip()
return msg
def _get_relevant_code_snippets_msgs(
self, relevant_code_snippets: list[ContextCodeSnippet]
self, relevant_code_snippets: set[ContextCodeSnippet]
) -> list[BaseMessage | str]:
messages: list[BaseMessage | str] = []
for code_snippet in relevant_code_snippets:
@@ -341,20 +208,24 @@ class RootCauseAgent(PatcherAgentBase):
return messages
def analyze_vulnerability(self, state: PatcherAgentState) -> Command[Literal[PatcherAgentName.CREATE_PATCH.value]]:
def analyze_vulnerability(
self, state: PatcherAgentState
) -> Command[Literal[PatcherAgentName.CONTEXT_RETRIEVER.value, PatcherAgentName.CREATE_PATCH.value]]:
"""Analyze the diff analysis and the code to understand the
vulnerability in the current code."""
logger.info("Analyzing the vulnerability in Challenge Task %s", self.challenge.name)
messages = []
if state.root_cause:
messages += OLD_ROOT_CAUSE_PROMPT.format_messages(root_cause=state.root_cause)
diff_content = "\n".join(diff.read_text() for diff in self.challenge.get_diffs())
root_cause = self.chain_call(
lambda x, y: x + y,
self.root_cause_chain,
{
"context": self.get_root_cause_context(state),
"messages": messages,
"DIFF": diff_content,
"PROJECT_NAME": self.challenge.project_name,
"SANITIZER": state.context.sanitizer,
"SANITIZER_OUTPUT": state.context.sanitizer_output,
"CODE_SNIPPETS": "\n".join(map(str, state.relevant_code_snippets)),
"OLD_ROOT_CAUSE": OLD_ROOT_CAUSE_TMPL.format(root_cause=state.root_cause) if state.root_cause else "",
},
default="",
)
@@ -364,7 +235,6 @@ class RootCauseAgent(PatcherAgentBase):
update_state = {
"root_cause": root_cause,
"patches": [],
"build_succeeded": None,
"pov_fixed": None,
"tests_passed": None,
@@ -372,6 +242,7 @@ class RootCauseAgent(PatcherAgentBase):
goto, update_state = self.get_code_snippet_requests(
root_cause,
update_state,
state.ctx_request_limit,
current_node=PatcherAgentName.ROOT_CAUSE_ANALYSIS.value,
default_goto=PatcherAgentName.CREATE_PATCH.value,
)
+219 -267
View File
@@ -15,235 +15,173 @@ from langchain_core.messages import BaseMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import (
ChatPromptTemplate,
FewShotChatMessagePromptTemplate,
MessagesPlaceholder,
)
from pydantic import BaseModel, Field
from langchain_core.runnables import (
ConfigurableField,
Runnable,
RunnableConfig,
)
from buttercup.patcher.agents.common import (
CONTEXT_CODE_SNIPPET_TMPL,
CONTEXT_DIFF_TMPL,
CONTEXT_PROJECT_TMPL,
CONTEXT_ROOT_CAUSE_TMPL,
ContextRetrieverState,
PatcherAgentState,
PatcherAgentName,
PatcherAgentBase,
CodeSnippetKey,
get_code_snippet_request_tmpl,
CodeSnippetRequest,
)
from buttercup.common.llm import ButtercupLLM, create_default_llm
from buttercup.patcher.utils import decode_bytes, PatchOutput, get_diff_content
logger = logging.getLogger(__name__)
SYSTEM_TMPL = """You are a security engineer tasked with fixing a bug in a \
software project. Fix the bug. Modify only the "Code" section(s) and perform as \
less changes as possible to *only* fix the bug explained in the root cause \
analysis (not other bugs).
You can modify one or more code snippets.
You don't have to modify all the code snippets.
You don't have to output a code snippet if you don't modify it.
If you modify a code snippet, you must provide the "File path" and the "Function
name", together with the old and new code parts. You don't need to rewrite the
whole code snippet, but make sure to write both the old part and the new part.
Moreover, if you don't modify the whole code snippet, include AT LEAST 5 lines
before and after the old part.
If you need more context to create a patch, you can ask for more code snippets.
You can request multiple code snippets at once. Request code snippets in the
following way:
{code_snippet_request_tmpl}
Always follow the scheme:
File path: <path-to-file>
Identifier: <identifier>
Old code:
```<language>
<old-code>
```
Code:
```<language>
<new-code>
```
For example, the output should look like this:
File path: path/to/file.c
Identifier: my_function
Old code:
```c
int my_function(int a) {{{{
return a;
}}}}
```
Code:
```c
int my_function(int a) {{{{
return a + 1;
}}}}
```
File path: path/to/file2.c
Identifier: my_function2
Old code:
```c
buf2 = get_buf_from_input();
printf(buf2);
// Existing comment
// Another existing comment
printf("Hello, World!");
```
Code:
```c
buf2 = get_buf_from_input();
printf("%%s", buf2);
// Existing comment
// Another existing comment
printf("Hello, World!");
```
""" # noqa: W293
SYSTEM_TMPL = SYSTEM_TMPL.format(code_snippet_request_tmpl=get_code_snippet_request_tmpl(2))
CODE_SNIPPET_REGEX = re.compile(
r"""File path: (.*?)
Identifier: (.*?)
Old code:
```.*?
(.*?)
```
Code:
```.*?
(.*?)
```""",
re.DOTALL | re.IGNORECASE,
SYSTEM_MSG = (
"""You are a skilled software engineer tasked with generating a patch for a specific vulnerability in a project."""
)
USER_MSG = """Your primary goal is to fix only the described vulnerability.
CREATE_PATCH_STR_PROMPT = ChatPromptTemplate.from_messages(
First, review the following project information and vulnerability analysis:
Project Name:
<project_name>
{PROJECT_NAME}
</project_name>
Root Cause Analysis:
<root_cause_analysis>
{ROOT_CAUSE_ANALYSIS}
</root_cause_analysis>
Code Snippets that may need modification:
<code_snippets>
{CODE_SNIPPETS}
</code_snippets>
{PREVIOUS_PATCH_PROMPT}{REVIEW_PROMPT}{BUILD_ANALYSIS_PROMPT}{POV_FAILED_PROMPT}{TESTS_FAILED_PROMPT}
Instructions:
1. Review the provided information carefully.
2. Plan your patch approach. Wrap your analysis and solution planning inside <vulnerability_analysis_and_solution> tags. In this section:
a) List all vulnerable parts of the code
b) If available
b.1) consider the Quality Engineer review comments and focus on addressing them
b.2) consider the build failure analysis and focus on addressing it
b.3) consider the POV failure analysis and focus on addressing it
b.4) consider the tests failure analysis and focus on addressing it
c) Propose multiple potential solutions for the vulnerability
d) Evaluate each solution's pros and cons
e) If you need additional context to patch the project, request it using <code_requests> tags. Be specific about what code you need and why. For example:
<code_requests>
<code_request>
Full implementation of the function 'validate_input()' from the file 'input_validation.c', as it's referenced in the analysis but not fully visible.
</code_request>
<code_request>
...
</code_request>
</code_requests>
If you make a code request, wait for a response before proceeding.
f) Choose the best solution:
Consider:
- The specific part(s) of the code that need modification
- Potential solutions and their pros/cons
3. Explain the changes you intend to make and how they fix the vulnerability. Use <explanation> tags for this section.
- If you need other code snippets to fix the vulnerability, request them using <code_requests> tags.
4. Generate the patch based on your planning and explanation. Remember:
- Only fix the described vulnerability
- You can modify one or more code snippets
- You don't have to modify all code snippets
- You don't need to output snippets you haven't modified
5. Format your output as follows for each modified code snippet:
<patch>
<identifier>[Identifier of the code snippet]</identifier>
<file_path>[File path of the code snippet]</file_path>
<old_code>
[Include at least 5 lines before the modified part, if available]
[Old code that needs to be replaced]
[Include at least 5 lines after the modified part, if available]
</old_code>
<new_code>
[Include at least 5 lines before the modified part, if available]
[New code that fixes the vulnerability]
[Include at least 5 lines after the modified part, if available]
</new_code>
</patch>
Remember to focus solely on fixing the described vulnerability. Do not make any unrelated changes or improvements to the code.
Begin your vulnerability analysis and solution planning now. If you need to request additional code, do so before providing any patches.
"""
PROMPT = ChatPromptTemplate.from_messages(
[
("system", SYSTEM_TMPL),
MessagesPlaceholder(variable_name="context"),
MessagesPlaceholder(variable_name="messages", optional=True),
(
"user",
"As output, write the modified code snippets, as instructed, but "
"with the bug fixed. Leave the extra context intact and do not "
"output it. Do not output a code snippet that does not require a change. "
"If necessary, request more code snippets.",
),
("system", SYSTEM_MSG),
("user", USER_MSG),
("ai", "<vulnerability_analysis_and_solution>"),
]
)
ADDRESS_REVIEW_PATCH_PROMPT = ChatPromptTemplate.from_messages(
[
(
"human",
"""The last patch you have created does not follow the quality \
assurance guidelines.
ADDRESS_REVIEW_PATCH_PROMPT = """
The last patch you have created does not pass the review process performed by the Quality Engineer. The Quality Engineer has provided the following review:
The Quality Engineer review:
```
{review}
```
<quality_engineer_review>
{PATCH_REVIEW}
</quality_engineer_review>
Please generate a patch that fixes the vulnerability and addresses the Quality \
Engineer review. Do not generate an already existing patch.
""",
),
]
)
Please generate a new patch that fixes the vulnerability and addresses the Quality Engineer review. Do not generate an already existing patch.
When considering possible patch solutions, consider the review comments.
"""
FIX_BUILD_PROMPT = ChatPromptTemplate.from_messages(
[
(
"human",
"""The last patch you have created does not build correctly.
Build stdout:
```
{build_stdout}
```
Build stderr:
```
{build_stderr}
```
Please generate a patch that builds correctly and fixes the vulnerability. Do \
not generate an already existing patch.""",
),
]
)
BUILD_ANALYSIS_PROMPT = ChatPromptTemplate.from_messages(
[
(
"human",
"""The last patch you have created does not build correctly. Another \
BUILD_ANALYSIS_PROMPT = """
The last patch you have created does not build correctly. Another \
software engineer has provided the following analysis about the build failures.
Build failure analysis:
```
<build_failure_analysis>
{build_failure_analysis}
```
</build_failure_analysis>
Please generate a patch that builds correctly and fixes the vulnerability. Do \
not generate an already existing patch.""",
)
]
)
not generate an already existing patch.
"""
PATCH_PROMPT = ChatPromptTemplate.from_messages(
[
(
"ai",
"""Patch:
```
{patch}
```""",
),
("user", "This patch does not work (build, fix vulnerability, tests, etc.)."),
]
)
PATCH_PROMPT = """
You previously tried the following patch, but it was not good enough (it failed the review process, failed to build, failed to fix the vulnerability, failed the tests, etc.).
POV_FAILED_PROMPT = ChatPromptTemplate.from_messages(
[
(
"human",
"The last patch you have created does not fix the vulnerability \
<previous_patch>
{PREVIOUS_PATCH}
</previous_patch>
"""
POV_FAILED_PROMPT = """
The last patch you have created does not fix the vulnerability \
correctly. Please generate a patch that fixes the vulnerability. Do not generate \
an already existing patch.",
),
]
)
an already existing patch.
"""
TESTS_FAILED_PROMPT = ChatPromptTemplate.from_messages(
[
(
"human",
"""The last patch you have created does not pass some tests.
TESTS_FAILED_PROMPT = """
The last patch you have created does not pass some tests.
Tests stdout:
```
<tests_stdout>
{tests_stdout}
```
</tests_stdout>
Tests stderr:
```
<tests_stderr>
{tests_stderr}
```
</tests_stderr>
Please generate a new patch that fixes the vulnerability but also makes the \
tests work. Do not generate an already existing patch.""",
),
]
)
tests work. Do not generate an already existing patch.
"""
class CodeSnippetChange(BaseModel):
@@ -287,9 +225,7 @@ class SWEAgent(PatcherAgentBase):
def __post_init__(self) -> None:
"""Initialize a few fields"""
default_llm = create_default_llm(
model_name=ButtercupLLM.OPENAI_GPT_4O.value, temperature=0.1
).configurable_fields(
default_llm = create_default_llm(model_name=ButtercupLLM.OPENAI_GPT_4O.value).configurable_fields(
temperature=ConfigurableField(
id="llm_temperature",
name="LLM temperature",
@@ -301,7 +237,7 @@ class SWEAgent(PatcherAgentBase):
ButtercupLLM.CLAUDE_3_5_SONNET,
]:
fallback_llms.append(
create_default_llm(model_name=fb_model.value, temperature=0.1).configurable_fields(
create_default_llm(model_name=fb_model.value).configurable_fields(
temperature=ConfigurableField(
id="llm_temperature",
name="LLM temperature",
@@ -311,19 +247,23 @@ class SWEAgent(PatcherAgentBase):
)
self.llm = default_llm.with_fallbacks(fallback_llms)
self.code_snippets_chain = CREATE_PATCH_STR_PROMPT | self.llm | StrOutputParser()
self.code_snippets_chain = PROMPT | self.llm | StrOutputParser()
def parse_code_snippets(self, msg: str) -> CodeSnippetChanges:
"""Parse the code snippets from the string."""
matches = CODE_SNIPPET_REGEX.findall(msg)
code_snippets_re = re.compile(
r"<patch>.*?<identifier>(.*?)</identifier>.*?<file_path>(.*?)</file_path>.*?<old_code>(.*?)</old_code>.*?<new_code>(.*?)</new_code>.*?</patch>",
re.DOTALL | re.IGNORECASE,
)
matches = code_snippets_re.findall(msg)
items: list[CodeSnippetChange] = []
for match in matches:
file_path, function_name, old_code, code = match
identifier, file_path, old_code, code = match
items.append(
CodeSnippetChange(
key=CodeSnippetKey(file_path=file_path.strip(), identifier=function_name.strip()),
old_code=old_code,
code=code,
key=CodeSnippetKey(file_path=file_path.strip(), identifier=identifier.strip()),
old_code=old_code.strip("\n"),
code=code.strip("\n"),
)
)
@@ -442,17 +382,15 @@ class SWEAgent(PatcherAgentBase):
)
code_snippet_key = closest_match
else:
logger.warning(
"Code snippet not found in the original context, trying to continue anyway: %s",
code_snippet_key,
)
return None
return code_snippet_key
def _get_snippets_patches(
self, code_snippets: CodeSnippetChanges, orig_code_snippets: dict[CodeSnippetKey, str]
) -> list[PatchOutput]:
) -> tuple[list[PatchOutput], list[CodeSnippetRequest]]:
patches: list[PatchOutput] = []
code_snippet_requests: list[CodeSnippetRequest] = []
for code_snippet_idx, code_snippet in enumerate(code_snippets.items or []):
if not code_snippet.is_valid():
logger.warning("Invalid code snippet: %s (%d)", code_snippet.key, code_snippet_idx)
@@ -463,9 +401,13 @@ class SWEAgent(PatcherAgentBase):
logger.warning(
"Could not find a valid code snippet key for %s (%d)", code_snippet.key, code_snippet_idx
)
code_snippet_key = CodeSnippetKey(
file_path=code_snippet.key.file_path, identifier=code_snippet.key.identifier
code_snippet_requests.append(
CodeSnippetRequest(
request="Provide code snippet %s / %s"
% (code_snippet.key.file_path, code_snippet.key.identifier)
)
)
continue
get_file_content_result = self._get_file_content(code_snippet_key.file_path)
if get_file_content_result is None:
@@ -521,24 +463,24 @@ class SWEAgent(PatcherAgentBase):
)
patches.append(patch)
return patches
return patches, code_snippet_requests
def create_upatch(self, inp: CreateUPatchInput | dict) -> PatchOutput:
def create_upatch(self, inp: CreateUPatchInput | dict) -> PatchOutput | list[CodeSnippetRequest]:
"""Extract the patch the new vulnerable code."""
inp = inp if isinstance(inp, CreateUPatchInput) else CreateUPatchInput(**inp)
code_snippets, state = inp.code_snippets, inp.state
orig_code_snippets = {
CodeSnippetKey(file_path=cs.key.file_path, identifier=cs.key.identifier): cs.code
for cs in (state.relevant_code_snippets)
for cs in state.relevant_code_snippets
if cs.code and cs.key.identifier
}
logger.debug("Creating patches for %d code snippets", len(code_snippets.items or []))
patches: list[PatchOutput] = self._get_snippets_patches(code_snippets, orig_code_snippets)
if not patches:
patches, code_snippet_requests = self._get_snippets_patches(code_snippets, orig_code_snippets)
if not patches or code_snippet_requests:
logger.warning("No valid patches generated")
return PatchOutput(task_id="", vulnerability_id="", patch="")
return code_snippet_requests
# Concatenate all patches in one
logger.debug("Concatenating %d patches", len(patches))
@@ -552,7 +494,14 @@ class SWEAgent(PatcherAgentBase):
def create_patch_node(
self, state: PatcherAgentState
) -> Command[Literal[PatcherAgentName.REVIEW_PATCH.value, PatcherAgentName.CONTEXT_RETRIEVER.value, END]]:
) -> Command[
Literal[
PatcherAgentName.REVIEW_PATCH.value,
PatcherAgentName.ROOT_CAUSE_ANALYSIS.value,
PatcherAgentName.CONTEXT_RETRIEVER.value,
END,
]
]:
"""Node in the LangGraph that generates a patch (in diff format)"""
if state.patch_tries >= self.max_patch_retries:
logger.warning("Reached max patch tries, terminating the patching process")
@@ -571,81 +520,77 @@ class SWEAgent(PatcherAgentBase):
"tests_passed": None,
}
messages = []
previous_patch_prompt = ""
review_prompt = ""
build_analysis_prompt = ""
pov_failed_prompt = ""
tests_failed_prompt = ""
last_patch = state.get_last_patch()
if last_patch:
old_patches = FewShotChatMessagePromptTemplate(
input_variables=[],
examples=[{"patch": last_patch.patch}],
example_prompt=PATCH_PROMPT,
)
messages += old_patches.format_messages()
previous_patch_prompt = PATCH_PROMPT.format(PREVIOUS_PATCH=last_patch.patch)
if state.patch_review is not None:
messages += ADDRESS_REVIEW_PATCH_PROMPT.format_messages(review=state.patch_review)
review_prompt = ADDRESS_REVIEW_PATCH_PROMPT.format(PATCH_REVIEW=state.patch_review)
elif state.build_succeeded is False:
messages += BUILD_ANALYSIS_PROMPT.format_messages(
build_failure_analysis=state.build_analysis,
)
build_analysis_prompt = BUILD_ANALYSIS_PROMPT.format(build_failure_analysis=state.build_analysis)
elif state.pov_fixed is False:
messages += POV_FAILED_PROMPT.format_messages()
pov_failed_prompt = POV_FAILED_PROMPT.format()
elif state.tests_passed is False:
messages += TESTS_FAILED_PROMPT.format_messages(
tests_failed_prompt = TESTS_FAILED_PROMPT.format(
tests_stdout=decode_bytes(state.tests_stdout),
tests_stderr=decode_bytes(state.tests_stderr),
)
temperature = 0.1
is_patch_generated = False
patch_str: str = self.chain_call(
lambda x, y: x + y,
self.code_snippets_chain,
{
"PROJECT_NAME": self.challenge.name,
"ROOT_CAUSE_ANALYSIS": state.root_cause,
"CODE_SNIPPETS": "\n".join(map(str, state.relevant_code_snippets)),
"PREVIOUS_PATCH_PROMPT": previous_patch_prompt,
"REVIEW_PROMPT": review_prompt,
"BUILD_ANALYSIS_PROMPT": build_analysis_prompt,
"POV_FAILED_PROMPT": pov_failed_prompt,
"TESTS_FAILED_PROMPT": tests_failed_prompt,
},
default="",
)
goto, update_state = self.get_code_snippet_requests(
patch_str,
update_state,
state.ctx_request_limit,
current_node=PatcherAgentName.CREATE_PATCH.value,
default_goto=PatcherAgentName.REVIEW_PATCH.value,
)
if goto == PatcherAgentName.CONTEXT_RETRIEVER.value:
return Command(
update=update_state,
goto=goto,
)
for _ in range(3):
patch_str: str = self.chain_call(
lambda x, y: x + y,
self.code_snippets_chain,
{
"context": self.get_context(state),
"state": state,
"messages": messages,
},
config=RunnableConfig(
configurable={
"llm_temperature": temperature,
},
create_upatch_input = CreateUPatchInput(
code_snippets=self.parse_code_snippets(patch_str),
state=state,
)
patch = self.create_upatch(create_upatch_input)
if isinstance(patch, list):
logger.info("Requesting new code snippets for patch generation")
return Command(
update=ContextRetrieverState(
code_snippet_requests=patch,
prev_node=PatcherAgentName.CREATE_PATCH.value,
),
default="",
goto=PatcherAgentName.CONTEXT_RETRIEVER.value,
)
goto, update_state = self.get_code_snippet_requests(
patch_str,
update_state,
current_node=PatcherAgentName.CREATE_PATCH.value,
default_goto=PatcherAgentName.REVIEW_PATCH.value,
if not patch or not patch.patch:
logger.error("Could not generate a patch")
return Command(
update=update_state,
goto=PatcherAgentName.ROOT_CAUSE_ANALYSIS.value,
)
if goto == PatcherAgentName.CONTEXT_RETRIEVER.value:
return Command(
update=update_state,
goto=goto,
)
create_upatch_input = CreateUPatchInput(
code_snippets=self.parse_code_snippets(patch_str),
state=state,
)
patch = self.create_upatch(create_upatch_input)
if not patch or not patch.patch:
logger.error("Could not generate a patch")
temperature += 0.1
continue
if patch.patch in [p.patch for p in (state.patches)]:
logger.error("Generated patch already exists, try again with higher temperature...")
temperature += 0.2
continue
is_patch_generated = True
break
if not is_patch_generated:
raise ValueError("Could not generate a new different patch")
logger.info("Generated a patch for Challenge Task %s", self.challenge.name)
logger.debug("Patch: %s", patch.patch)
@@ -657,6 +602,13 @@ class SWEAgent(PatcherAgentBase):
"patch_tries": patch_tries,
}
)
if patch.patch in [p.patch for p in (state.patches)]:
logger.warning("Generated patch already exists, going back to root cause analysis")
return Command(
update=update_state,
goto=PatcherAgentName.ROOT_CAUSE_ANALYSIS.value,
)
return Command(
update=update_state,
goto=PatcherAgentName.REVIEW_PATCH.value,
+19
View File
@@ -0,0 +1,19 @@
import pytest
def pytest_addoption(parser):
parser.addoption("--runintegration", action="store_true", default=False, help="run integration tests")
def pytest_configure(config):
config.addinivalue_line("markers", "integration: mark test as an integration test")
def pytest_collection_modifyitems(config, items):
if config.getoption("--runintegration"):
# --runintegration given in cli: do not skip integration tests
return
skip_integration = pytest.mark.skip(reason="need --runintegration option to run")
for item in items:
if "integration" in item.keywords:
item.add_marker(skip_integration)
+520
View File
@@ -0,0 +1,520 @@
"""Tests for the ContextRetrieverAgent."""
import pytest
from pathlib import Path
import subprocess
from unittest.mock import MagicMock, patch
import os
from buttercup.patcher.agents.context_retriever import ContextRetrieverAgent
from buttercup.patcher.agents.common import ContextRetrieverState, CodeSnippetRequest
from buttercup.patcher.patcher import PatchInput
from langgraph.types import Command
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.task_meta import TaskMeta
from langchain_core.messages import AIMessage
from langchain_core.messages.tool import ToolCall
@pytest.fixture
def mock_llm():
return MagicMock()
@pytest.fixture
def mock_chain(mock_llm: MagicMock):
res = MagicMock()
res.with_fallbacks.return_value = mock_llm
res.with_fallbacks.return_value.bind_tools.return_value = mock_llm
return res
@pytest.fixture(autouse=True)
def mock_llm_functions(mock_chain: MagicMock):
"""Mock LLM creation functions and environment variables."""
with (
patch.dict(os.environ, {"BUTTERCUP_LITELLM_HOSTNAME": "http://test-host", "BUTTERCUP_LITELLM_KEY": "test-key"}),
patch("buttercup.common.llm.create_default_llm", return_value=mock_chain),
patch("buttercup.common.llm.create_llm", return_value=mock_chain),
):
yield
@pytest.fixture
def task_dir(tmp_path: Path) -> Path:
"""Create a mock challenge task directory structure."""
# Create the main directories
tmp_path = tmp_path / "test-challenge-task"
oss_fuzz = tmp_path / "fuzz-tooling" / "my-oss-fuzz"
source = tmp_path / "src" / "my-source"
diffs = tmp_path / "diff" / "my-diff"
oss_fuzz.mkdir(parents=True, exist_ok=True)
source.mkdir(parents=True, exist_ok=True)
diffs.mkdir(parents=True, exist_ok=True)
# Create project.yaml file
project_yaml_path = oss_fuzz / "projects" / "example_project" / "project.yaml"
project_yaml_path.parent.mkdir(parents=True, exist_ok=True)
project_yaml_path.write_text("""name: example_project
language: c
""")
# Create some mock patch files
(diffs / "patch1.diff").write_text("mock patch 1")
(diffs / "patch2.diff").write_text("mock patch 2")
# Create a mock helper.py file
helper_path = oss_fuzz / "infra/helper.py"
helper_path.parent.mkdir(parents=True, exist_ok=True)
helper_path.write_text("import sys;\nsys.exit(0)\n")
# Create a mock test.c file
(source / "test.c").write_text("int main() { return 0; }")
(source / "test.h").write_text("struct ebitmap_t { int a; };")
TaskMeta(project_name="example_project", focus="my-source", task_id="task-id-challenge-task").save(tmp_path)
return tmp_path
@pytest.fixture
def mock_challenge(task_dir: Path) -> ChallengeTask:
"""Create a mock challenge task for testing."""
return ChallengeTask(
read_only_task_dir=task_dir,
local_task_dir=task_dir,
)
@pytest.fixture
def selinux_oss_fuzz_task(tmp_path: Path) -> ChallengeTask:
"""Create a challenge task using a real OSS-Fuzz repository."""
tmp_path = tmp_path / "selinux"
tmp_path.mkdir(parents=True)
oss_fuzz_dir = tmp_path / "fuzz-tooling"
oss_fuzz_dir.mkdir(parents=True)
source_dir = tmp_path / "src"
source_dir.mkdir(parents=True)
subprocess.run(
[
"git",
"-C",
str(oss_fuzz_dir),
"clone",
"https://github.com/google/oss-fuzz.git",
],
check=True,
)
# Restore libjpeg-turbo project directory to specific commit
subprocess.run(
[
"git",
"-C",
str(oss_fuzz_dir / "oss-fuzz"),
"checkout",
"ef2f42b3b10af381d3d55cc901fde0729e54573b",
"--",
"projects/selinux",
],
check=True,
)
# Download selinux source code
url = "https://github.com/SELinuxProject/selinux"
subprocess.run(["git", "-C", str(source_dir), "clone", url], check=True)
subprocess.run(
[
"git",
"-C",
str(source_dir / "selinux"),
"checkout",
"c35919a703302bd571476f245d856174a1fe1926",
],
check=True,
)
# Create task metadata
TaskMeta(project_name="selinux", focus="selinux", task_id="task-id-selinux").save(tmp_path)
return ChallengeTask(
read_only_task_dir=tmp_path,
local_task_dir=tmp_path,
)
@pytest.fixture
def selinux_agent(selinux_oss_fuzz_task: ChallengeTask, tmp_path: Path) -> ContextRetrieverAgent:
"""Create a ContextRetrieverAgent instance."""
patch_input = PatchInput(
challenge_task_dir=selinux_oss_fuzz_task.task_dir,
task_id=selinux_oss_fuzz_task.task_meta.task_id,
vulnerability_id="vuln-id-selinux",
harness_name="secilc-fuzzer",
# not used by the context retriever
pov=Path("pov-path-selinux"),
sanitizer_output="sanitizer-output-selinux",
engine="libfuzzer",
sanitizer="address",
)
wdir = tmp_path / "work_dir"
wdir.mkdir(parents=True)
return ContextRetrieverAgent(
input=patch_input,
chain_call=lambda _, runnable, args, config, default: runnable.invoke(args, config=config),
challenge=selinux_oss_fuzz_task,
work_dir=wdir,
)
@pytest.fixture
def mock_agent(mock_challenge: ChallengeTask, tmp_path: Path) -> ContextRetrieverAgent:
"""Create a ContextRetrieverAgent instance."""
patch_input = PatchInput(
challenge_task_dir=mock_challenge.task_dir,
task_id=mock_challenge.task_meta.task_id,
vulnerability_id="vuln-id-mock",
harness_name="mock-harness",
# not used by the context retriever
pov=Path("pov-path-mock"),
sanitizer_output="sanitizer-output-mock",
engine="libfuzzer",
sanitizer="address",
)
wdir = tmp_path / "work_dir"
wdir.mkdir(parents=True)
return ContextRetrieverAgent(
input=patch_input,
chain_call=lambda _, runnable, args, config, default: runnable.invoke(args, config=config),
challenge=mock_challenge,
work_dir=wdir,
)
@pytest.mark.integration
def test_retrieve_context_basic(selinux_agent: ContextRetrieverAgent, mock_llm: MagicMock) -> None:
"""Test basic context retrieval functionality."""
# Create a test state with a simple request
state = ContextRetrieverState(
code_snippet_requests=[
CodeSnippetRequest(
request="Find function ebitmap_match_any",
)
],
prev_node="test_node",
)
# Execute the retrieve_context method
mock_llm.invoke.side_effect = [
# First response: Agent decides to use grep to search for the function
AIMessage(
content="I'll search for the function using grep.",
tool_calls=[
ToolCall(
id="grep_call_1",
name="grep",
args={
"pattern": "ebitmap_match_any",
},
)
],
),
# Agent decides to get the function definition
AIMessage(
content="I found the function. Let me get its definition.",
tool_calls=[
ToolCall(
id="get_function_definition_call_1",
name="get_function_definition",
args={"function_name": "ebitmap_match_any", "file_path": "libsepol/src/ebitmap.c"},
)
],
),
AIMessage(
content="I'm done <END>",
),
]
result = selinux_agent.retrieve_context(state)
# Verify the result
assert isinstance(result, Command)
assert result.goto == "test_node"
assert "relevant_code_snippets" in result.update
assert len(result.update["relevant_code_snippets"]) == 1
# Verify the code snippet content
snippet = next(iter(result.update["relevant_code_snippets"]))
assert "ebitmap_match_any" in snippet.code
assert "const ebitmap_t *e1" in snippet.code
def test_missing_arg_tool_call(mock_agent: ContextRetrieverAgent, mock_llm: MagicMock) -> None:
"""Test basic context retrieval functionality."""
# Create a test state with a simple request
state = ContextRetrieverState(
code_snippet_requests=[
CodeSnippetRequest(
request="Find function main",
)
],
prev_node="test_node",
)
# Execute the retrieve_context method
mock_llm.invoke.side_effect = [
AIMessage(
content="I'll list the files in the project.",
tool_calls=[
ToolCall(
id="list_files_call_1",
name="ls",
args={}, # no args for list_files
)
],
),
AIMessage(
content="Let me pass the path to ls",
tool_calls=[
ToolCall(
id="ls_call_1",
name="ls",
args={
"path": ".",
},
)
],
),
AIMessage(
content="Let's get the definition",
tool_calls=[
ToolCall(
id="get_function_definition_call_1",
name="get_function_definition",
args={
"function_name": "main",
"file_path": "test.c",
},
)
],
),
AIMessage(
content="I'm done <END>",
),
]
result = mock_agent.retrieve_context(state)
# Verify the result
assert isinstance(result, Command)
assert result.goto == "test_node"
assert "relevant_code_snippets" in result.update
assert len(result.update["relevant_code_snippets"]) == 1
def test_recursion_limit(mock_agent: ContextRetrieverAgent, mock_llm: MagicMock) -> None:
"""Test hitting the context request limit."""
state = ContextRetrieverState(
code_snippet_requests=[
CodeSnippetRequest(
request="Find function main",
)
],
prev_node="test_node",
)
mock_llm.invoke.side_effect = [
AIMessage(
content="I'll list the files in the dir %s." % (i,),
tool_calls=[
ToolCall(
id="list_files_call_%s" % (i,),
name="ls",
args={
"path": str(i),
},
)
],
)
for i in range(1, 1000)
]
result = mock_agent.retrieve_context(state)
# Verify the result
assert isinstance(result, Command)
assert result.goto == "test_node"
assert "relevant_code_snippets" in result.update
assert len(result.update["relevant_code_snippets"]) == 0, (
"Should not have any code snippets, we hit the recursion limit and never called the get_function_definition/get_type_definition tool"
)
def test_recursion_limit_tmp_code_snippets(mock_agent: ContextRetrieverAgent, mock_llm: MagicMock) -> None:
"""Test hitting the context request limit but getting some partial results."""
state = ContextRetrieverState(
code_snippet_requests=[
CodeSnippetRequest(request="Find function main"),
],
prev_node="test_node",
)
llm_invoke_side_effect = [
AIMessage(
content="I'll get the function definition.",
tool_calls=[
ToolCall(
id="get_function_definition_call",
name="get_function_definition",
args={
"function_name": "main",
"file_path": "test.c",
},
)
],
)
]
llm_invoke_side_effect += [
AIMessage(
content="I'll list the files in the dir %s." % (i,),
tool_calls=[
ToolCall(
id="list_files_call_%s" % (i,),
name="ls",
args={
"path": str(i),
},
)
],
)
for i in range(1, 1000)
]
mock_llm.invoke.side_effect = llm_invoke_side_effect
result = mock_agent.retrieve_context(state)
assert isinstance(result, Command)
assert result.goto == "test_node"
assert "relevant_code_snippets" in result.update
assert len(result.update["relevant_code_snippets"]) == 1
code_snippet = next(iter(result.update["relevant_code_snippets"]))
assert code_snippet.code == "int main() { return 0; }"
def test_dupped_code_snippets(mock_agent: ContextRetrieverAgent, mock_llm: MagicMock) -> None:
"""Test that we don't return duplicate code snippets."""
state = ContextRetrieverState(
code_snippet_requests=[
CodeSnippetRequest(request="Find function main"),
],
prev_node="test_node",
)
mock_llm.invoke.side_effect = [
AIMessage(
content="I'll get the function definition.",
tool_calls=[
ToolCall(
id="get_function_definition_call",
name="get_function_definition",
args={"function_name": "main", "file_path": "test.c"},
)
],
),
AIMessage(
content="I'll get the function definition.",
tool_calls=[
ToolCall(
id="get_function_definition_call",
name="get_function_definition",
args={"function_name": "main", "file_path": "test.c"},
)
],
),
AIMessage(
content="I'm done <END>",
),
]
result = mock_agent.retrieve_context(state)
assert isinstance(result, Command)
assert result.goto == "test_node"
assert "relevant_code_snippets" in result.update
assert len(result.update["relevant_code_snippets"]) == 1
code_snippet = next(iter(result.update["relevant_code_snippets"]))
assert code_snippet.code == "int main() { return 0; }"
def test_get_type_definition(mock_agent: ContextRetrieverAgent, mock_llm: MagicMock) -> None:
"""Test that we can get the type definition."""
state = ContextRetrieverState(
code_snippet_requests=[
CodeSnippetRequest(request="Find type ebitmap_t"),
],
prev_node="test_node",
)
mock_llm.invoke.side_effect = [
AIMessage(
content="I'll get the type definition.",
tool_calls=[
ToolCall(
id="get_type_definition_call",
name="get_type_definition",
args={"type_name": "ebitmap_t", "file_path": "test.h"},
)
],
),
AIMessage(
content="I'm done <END>",
),
]
result = mock_agent.retrieve_context(state)
assert isinstance(result, Command)
assert result.goto == "test_node"
assert "relevant_code_snippets" in result.update
assert len(result.update["relevant_code_snippets"]) == 1
code_snippet = next(iter(result.update["relevant_code_snippets"]))
assert code_snippet.code == "struct ebitmap_t { int a; }"
def test_get_definitions_no_paths(mock_agent: ContextRetrieverAgent, mock_llm: MagicMock) -> None:
"""Test that we can get the type definition even if the file path is not provided."""
state = ContextRetrieverState(
code_snippet_requests=[
CodeSnippetRequest(request="Find type ebitmap_t"),
],
prev_node="test_node",
)
mock_llm.invoke.side_effect = [
AIMessage(
content="I'll get the function definition.",
tool_calls=[
ToolCall(
id="get_function_definition_call",
name="get_function_definition",
args={"function_name": "main", "file_path": None},
)
],
),
AIMessage(
content="I'll get the type definition.",
tool_calls=[
ToolCall(
id="get_type_definition_call",
name="get_type_definition",
args={"type_name": "ebitmap_t", "file_path": None},
)
],
),
AIMessage(
content="I'm done <END>",
),
]
result = mock_agent.retrieve_context(state)
assert isinstance(result, Command)
assert result.goto == "test_node"
assert "relevant_code_snippets" in result.update
assert len(result.update["relevant_code_snippets"]) == 2
code_snippets = result.update["relevant_code_snippets"]
assert any(
snippet.code == "int main() { return 0; }" and snippet.key.file_path == "src/my-source/test.c"
for snippet in code_snippets
)
# FIXME: this is not working as expected
# assert any(snippet.code == "struct ebitmap_t { int a; }" and snippet.key.file_path == "src/my-source/test.h" for snippet in code_snippets)
+207
View File
@@ -0,0 +1,207 @@
"""Tests for the RootCauseAgent."""
import pytest
from pathlib import Path
from unittest.mock import MagicMock, patch
import os
from typing import Any, Callable
from langchain_core.runnables import Runnable, RunnableConfig
from langgraph.types import Command
from buttercup.common.challenge_task import ChallengeTask
from buttercup.common.task_meta import TaskMeta
from buttercup.patcher.agents.rootcause import RootCauseAgent
from buttercup.patcher.agents.common import PatcherAgentState, PatcherAgentName
from buttercup.patcher.patcher import PatchInput
from buttercup.patcher.agents.common import ContextCodeSnippet, CodeSnippetKey
@pytest.fixture
def mock_llm():
return MagicMock()
@pytest.fixture
def mock_chain(mock_llm: MagicMock):
res = MagicMock()
res.with_fallbacks.return_value = mock_llm
res.with_fallbacks.return_value.bind_tools.return_value = mock_llm
return res
@pytest.fixture(autouse=True)
def mock_llm_functions(mock_chain: MagicMock):
"""Mock LLM creation functions and environment variables."""
with (
patch.dict(os.environ, {"BUTTERCUP_LITELLM_HOSTNAME": "http://test-host", "BUTTERCUP_LITELLM_KEY": "test-key"}),
patch("buttercup.common.llm.create_default_llm", return_value=mock_chain),
patch("buttercup.common.llm.create_llm", return_value=mock_chain),
):
yield
@pytest.fixture
def task_dir(tmp_path: Path) -> Path:
"""Create a mock challenge task directory structure."""
# Create the main directories
tmp_path = tmp_path / "test-challenge-task"
oss_fuzz = tmp_path / "fuzz-tooling" / "my-oss-fuzz"
source = tmp_path / "src" / "my-source"
diffs = tmp_path / "diff" / "my-diff"
oss_fuzz.mkdir(parents=True, exist_ok=True)
source.mkdir(parents=True, exist_ok=True)
diffs.mkdir(parents=True, exist_ok=True)
# Create project.yaml file
project_yaml_path = oss_fuzz / "projects" / "example_project" / "project.yaml"
project_yaml_path.parent.mkdir(parents=True, exist_ok=True)
project_yaml_path.write_text("""name: example_project
language: c
""")
# Create some mock patch files
(diffs / "patch1.diff").write_text("mock patch 1")
(diffs / "patch2.diff").write_text("mock patch 2")
# Create a mock helper.py file
helper_path = oss_fuzz / "infra/helper.py"
helper_path.parent.mkdir(parents=True, exist_ok=True)
helper_path.write_text("import sys;\nsys.exit(0)\n")
# Create a mock test.c file
(source / "test.c").write_text("int main() { return 0; }")
(source / "test.h").write_text("struct ebitmap_t { int a; };")
TaskMeta(project_name="example_project", focus="my-source", task_id="task-id-challenge-task").save(tmp_path)
return tmp_path
@pytest.fixture
def mock_challenge(task_dir: Path) -> ChallengeTask:
"""Create a mock challenge task for testing."""
return ChallengeTask(
read_only_task_dir=task_dir,
local_task_dir=task_dir,
)
def _chain_call(
reduce_function: Callable,
runnable: Runnable,
args: dict[str, Any],
config: RunnableConfig | None = None,
default: Any = None,
) -> Any:
return runnable.invoke(args, config=config)
@pytest.fixture
def mock_agent(mock_challenge: ChallengeTask, tmp_path: Path) -> RootCauseAgent:
"""Create a RootCauseAgent instance."""
patch_input = PatchInput(
challenge_task_dir=mock_challenge.task_dir,
task_id=mock_challenge.task_meta.task_id,
vulnerability_id="vuln-id-mock",
harness_name="mock-harness",
# not used by the context retriever
pov=Path("pov-path-mock"),
sanitizer_output="sanitizer-output-mock",
engine="libfuzzer",
sanitizer="address",
)
wdir = tmp_path / "work_dir"
wdir.mkdir(parents=True)
return RootCauseAgent(
input=patch_input,
chain_call=_chain_call,
challenge=mock_challenge,
)
@pytest.fixture
def patcher_agent_state(mock_agent: RootCauseAgent) -> PatcherAgentState:
"""Create a PatcherAgentState instance."""
return PatcherAgentState(
context=mock_agent.input,
messages=[],
relevant_code_snippets=[
ContextCodeSnippet(
key=CodeSnippetKey(file_path="src/test.c", identifier="main"),
code="int main() { return 0; }",
),
],
)
def test_rootcause_simple_request(mock_agent: RootCauseAgent, patcher_agent_state: PatcherAgentState) -> None:
"""Test that the agent requests a code snippet."""
mock_agent.root_cause_chain = MagicMock()
mock_agent.root_cause_chain.invoke.side_effect = [
"""
<code_request>
Find the function definition for main
</code_request>
""",
]
result = mock_agent.analyze_vulnerability(patcher_agent_state)
assert isinstance(result, Command)
assert result.goto == PatcherAgentName.CONTEXT_RETRIEVER.value
assert len(result.update["code_snippet_requests"]) == 1
assert result.update["code_snippet_requests"][0].request == "Find the function definition for main"
def test_rootcause_multiple_requests(mock_agent: RootCauseAgent, patcher_agent_state: PatcherAgentState) -> None:
"""Test that the agent requests multiple code snippets."""
mock_agent.root_cause_chain = MagicMock()
mock_agent.root_cause_chain.invoke.side_effect = [
"""
<code_request>
Find the function definition for main
</code_request>
<code_request> Find the function definition for main2</code_request>
<code_request>Find the function definition for main3 </code_request>
<code_request>Find the function definition for
main4</code_request>
<code_request>Find the function definition for main5</code_request>
<code_request>
Find the function definition for main6
</code_request>
""",
]
result = mock_agent.analyze_vulnerability(patcher_agent_state)
assert isinstance(result, Command)
assert result.goto == PatcherAgentName.CONTEXT_RETRIEVER.value
assert len(result.update["code_snippet_requests"]) == 6
assert result.update["code_snippet_requests"][0].request == "Find the function definition for main"
assert result.update["code_snippet_requests"][1].request == "Find the function definition for main2"
assert result.update["code_snippet_requests"][2].request == "Find the function definition for main3"
assert result.update["code_snippet_requests"][3].request == "Find the function definition for \n main4"
assert result.update["code_snippet_requests"][4].request == "Find the function definition for main5"
assert result.update["code_snippet_requests"][5].request == "Find the function definition for main6"
def test_no_code_snippet_requests(mock_agent: RootCauseAgent, patcher_agent_state: PatcherAgentState) -> None:
"""Test that no code snippets are requested if not asked by the LLM when there are some snippets already."""
mock_agent.root_cause_chain = MagicMock()
mock_agent.root_cause_chain.invoke.side_effect = ["This is my analysis."]
result = mock_agent.analyze_vulnerability(patcher_agent_state)
assert isinstance(result, Command)
assert result.goto == PatcherAgentName.CREATE_PATCH.value
assert "code_snippet_requests" not in result.update
@pytest.mark.skip("TODO: Implement this")
def test_cant_call_with_no_snippets(mock_agent: RootCauseAgent, patcher_agent_state: PatcherAgentState) -> None:
"""Test that the agent can't call with no snippets."""
patcher_agent_state.relevant_code_snippets = []
with pytest.raises(ValueError):
mock_agent.analyze_vulnerability(patcher_agent_state)
+45 -5
View File
@@ -1,4 +1,5 @@
version = 1
revision = 1
requires-python = "==3.12.*"
resolution-markers = [
"python_version < '0'",
@@ -715,16 +716,18 @@ wheels = [
[[package]]
name = "langgraph"
version = "0.2.69"
version = "0.3.23"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "langchain-core" },
{ name = "langgraph-checkpoint" },
{ name = "langgraph-prebuilt" },
{ name = "langgraph-sdk" },
{ name = "xxhash" },
]
sdist = { url = "https://files.pythonhosted.org/packages/9b/36/d75836c6c1b4a817bb22f591137dc533be02fdba171df4d80eac49e22043/langgraph-0.2.69.tar.gz", hash = "sha256:77bd6efd967b4f092ec31d2148b3e6ba3c31e202b4f3a975dbb082b19b5bb057", size = 128593 }
sdist = { url = "https://files.pythonhosted.org/packages/2e/73/d4700fa35229618dfbb9bf1313eff8f31a43fc253bf1a989bb14e8e35dff/langgraph-0.3.23.tar.gz", hash = "sha256:054e64636be453626ea6e97be9172b2bae10d643253d0db313a0233f6a34f9dd", size = 117715 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/4d/be/f3489ee7a67bb5bb3dc256b9e95d19a44c31e337f0345a38b443018355e3/langgraph-0.2.69-py3-none-any.whl", hash = "sha256:b64a5755fa2c7f2f67608ff4ce0ef8c168b30a0fb551a6c1d2e19bf1d2268ce4", size = 148716 },
{ url = "https://files.pythonhosted.org/packages/99/3b/86c67f1df41e321c545d19ace9eb1272ea461c091aead5cf6a6024a86278/langgraph-0.3.23-py3-none-any.whl", hash = "sha256:b089f360c8a7d7891689c04fd2918e1c1bd288b231d1e33f205d6ad2c9c5d727", size = 139736 },
]
[[package]]
@@ -740,6 +743,19 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/4d/ef/c320b52035e29081f2693377602289a00545016b4adcc963d5e202ac0c92/langgraph_checkpoint-2.0.10-py3-none-any.whl", hash = "sha256:0d592cfda2df93844c6ea44d142170a8f7e5ba5320274e0e5e60e27f2749392c", size = 37476 },
]
[[package]]
name = "langgraph-prebuilt"
version = "0.1.7"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "langchain-core" },
{ name = "langgraph-checkpoint" },
]
sdist = { url = "https://files.pythonhosted.org/packages/fb/ed/e05dc6561aae82b2ec14b0a663f9bfd67a1da169d8895acecbb149cf9898/langgraph_prebuilt-0.1.7.tar.gz", hash = "sha256:5b086fad2ebfabe743bfbfa1e0248ce0d7d10eb0cc939b8bc7a9b41360b9d518", size = 23642 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/52/c2/e88798e0c698ae92f30b41966a781e1b28fa48ae4825462062722d1942bb/langgraph_prebuilt-0.1.7-py3-none-any.whl", hash = "sha256:35eff90fb86edd0b026a6744942bc59bad4a2d82a482aa860ff8d081af423956", size = 25021 },
]
[[package]]
name = "langgraph-sdk"
version = "0.1.51"
@@ -1120,9 +1136,9 @@ tests = [
requires-dist = [
{ name = "common", editable = "../common" },
{ name = "flake8", marker = "extra == 'tests'", specifier = ">=7.1.1" },
{ name = "langchain", specifier = ">=0.3.17" },
{ name = "langchain", specifier = "~=0.3.20" },
{ name = "langchain-community", specifier = ">=0.3.16" },
{ name = "langgraph", specifier = ">=0.2.69" },
{ name = "langgraph", specifier = "==0.3.23" },
{ name = "mypy", marker = "extra == 'tests'", specifier = ">=1.14.1" },
{ name = "openai", specifier = ">=1.60.2" },
{ name = "program-model", editable = "../program-model" },
@@ -1135,6 +1151,7 @@ requires-dist = [
{ name = "ruff", marker = "extra == 'tests'", specifier = ">=0.9.2" },
{ name = "types-python-dateutil", marker = "extra == 'tests'", specifier = ">=2.9.0.20241206" },
]
provides-extras = ["tests"]
[[package]]
name = "pluggy"
@@ -1760,6 +1777,29 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/94/db/fd0326e331726f07ff7f40675cd86aa804bfd2e5016c727fa761c934990e/xmltodict-0.13.0-py2.py3-none-any.whl", hash = "sha256:aa89e8fd76320154a40d19a0df04a4695fb9dc5ba977cbb68ab3e4eb225e7852", size = 9971 },
]
[[package]]
name = "xxhash"
version = "3.5.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/00/5e/d6e5258d69df8b4ed8c83b6664f2b47d30d2dec551a29ad72a6c69eafd31/xxhash-3.5.0.tar.gz", hash = "sha256:84f2caddf951c9cbf8dc2e22a89d4ccf5d86391ac6418fe81e3c67d0cf60b45f", size = 84241 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/07/0e/1bfce2502c57d7e2e787600b31c83535af83746885aa1a5f153d8c8059d6/xxhash-3.5.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:14470ace8bd3b5d51318782cd94e6f94431974f16cb3b8dc15d52f3b69df8e00", size = 31969 },
{ url = "https://files.pythonhosted.org/packages/3f/d6/8ca450d6fe5b71ce521b4e5db69622383d039e2b253e9b2f24f93265b52c/xxhash-3.5.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:59aa1203de1cb96dbeab595ded0ad0c0056bb2245ae11fac11c0ceea861382b9", size = 30787 },
{ url = "https://files.pythonhosted.org/packages/5b/84/de7c89bc6ef63d750159086a6ada6416cc4349eab23f76ab870407178b93/xxhash-3.5.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:08424f6648526076e28fae6ea2806c0a7d504b9ef05ae61d196d571e5c879c84", size = 220959 },
{ url = "https://files.pythonhosted.org/packages/fe/86/51258d3e8a8545ff26468c977101964c14d56a8a37f5835bc0082426c672/xxhash-3.5.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:61a1ff00674879725b194695e17f23d3248998b843eb5e933007ca743310f793", size = 200006 },
{ url = "https://files.pythonhosted.org/packages/02/0a/96973bd325412feccf23cf3680fd2246aebf4b789122f938d5557c54a6b2/xxhash-3.5.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f2f2c61bee5844d41c3eb015ac652a0229e901074951ae48581d58bfb2ba01be", size = 428326 },
{ url = "https://files.pythonhosted.org/packages/11/a7/81dba5010f7e733de88af9555725146fc133be97ce36533867f4c7e75066/xxhash-3.5.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9d32a592cac88d18cc09a89172e1c32d7f2a6e516c3dfde1b9adb90ab5df54a6", size = 194380 },
{ url = "https://files.pythonhosted.org/packages/fb/7d/f29006ab398a173f4501c0e4977ba288f1c621d878ec217b4ff516810c04/xxhash-3.5.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:70dabf941dede727cca579e8c205e61121afc9b28516752fd65724be1355cc90", size = 207934 },
{ url = "https://files.pythonhosted.org/packages/8a/6e/6e88b8f24612510e73d4d70d9b0c7dff62a2e78451b9f0d042a5462c8d03/xxhash-3.5.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e5d0ddaca65ecca9c10dcf01730165fd858533d0be84c75c327487c37a906a27", size = 216301 },
{ url = "https://files.pythonhosted.org/packages/af/51/7862f4fa4b75a25c3b4163c8a873f070532fe5f2d3f9b3fc869c8337a398/xxhash-3.5.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:3e5b5e16c5a480fe5f59f56c30abdeba09ffd75da8d13f6b9b6fd224d0b4d0a2", size = 203351 },
{ url = "https://files.pythonhosted.org/packages/22/61/8d6a40f288f791cf79ed5bb113159abf0c81d6efb86e734334f698eb4c59/xxhash-3.5.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:149b7914451eb154b3dfaa721315117ea1dac2cc55a01bfbd4df7c68c5dd683d", size = 210294 },
{ url = "https://files.pythonhosted.org/packages/17/02/215c4698955762d45a8158117190261b2dbefe9ae7e5b906768c09d8bc74/xxhash-3.5.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:eade977f5c96c677035ff39c56ac74d851b1cca7d607ab3d8f23c6b859379cab", size = 414674 },
{ url = "https://files.pythonhosted.org/packages/31/5c/b7a8db8a3237cff3d535261325d95de509f6a8ae439a5a7a4ffcff478189/xxhash-3.5.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fa9f547bd98f5553d03160967866a71056a60960be00356a15ecc44efb40ba8e", size = 192022 },
{ url = "https://files.pythonhosted.org/packages/78/e3/dd76659b2811b3fd06892a8beb850e1996b63e9235af5a86ea348f053e9e/xxhash-3.5.0-cp312-cp312-win32.whl", hash = "sha256:f7b58d1fd3551b8c80a971199543379be1cee3d0d409e1f6d8b01c1a2eebf1f8", size = 30170 },
{ url = "https://files.pythonhosted.org/packages/d9/6b/1c443fe6cfeb4ad1dcf231cdec96eb94fb43d6498b4469ed8b51f8b59a37/xxhash-3.5.0-cp312-cp312-win_amd64.whl", hash = "sha256:fa0cafd3a2af231b4e113fba24a65d7922af91aeb23774a8b78228e6cd785e3e", size = 30040 },
{ url = "https://files.pythonhosted.org/packages/0f/eb/04405305f290173acc0350eba6d2f1a794b57925df0398861a20fbafa415/xxhash-3.5.0-cp312-cp312-win_arm64.whl", hash = "sha256:586886c7e89cb9828bcd8a5686b12e161368e0064d040e225e72607b43858ba2", size = 26796 },
]
[[package]]
name = "yarl"
version = "1.18.3"