mirror of
https://github.com/trailofbits/buttercup
synced 2026-06-21 14:11:39 +00:00
patcher: improve context retriever agent (#357)
This commit is contained in:
committed by
GitHub
parent
98e736c194
commit
658f9c64c4
@@ -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
Binary file not shown.
@@ -0,0 +1,83 @@
|
||||
=================================================================
|
||||
[1m[31m==14==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x7f013ccb77f0 at pc 0x555dcfad2fef bp 0x7ffd2207d630 sp 0x7ffd2207d628
|
||||
[1m[0m[1m[34mREAD of size 8 at 0x7f013ccb77f0 thread T0[1m[0m
|
||||
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
|
||||
[1m[32m0x7f013ccb77f0 is located 16 bytes before 1340512-byte region [0x7f013ccb7800,0x7f013cdfec60)
|
||||
[1m[0m[1m[35mallocated by thread T0 here:[1m[0m
|
||||
#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: [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m
|
||||
0x7f013ccb7580: [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m
|
||||
0x7f013ccb7600: [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m
|
||||
0x7f013ccb7680: [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m
|
||||
0x7f013ccb7700: [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m
|
||||
=>0x7f013ccb7780: [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m [1m[31mfa[1m[0m[[1m[31mfa[1m[0m][1m[31mfa[1m[0m
|
||||
0x7f013ccb7800: [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m
|
||||
0x7f013ccb7880: [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m
|
||||
0x7f013ccb7900: [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m
|
||||
0x7f013ccb7980: [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m
|
||||
0x7f013ccb7a00: [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m [1m[0m00[1m[0m
|
||||
Shadow byte legend (one shadow byte represents 8 application bytes):
|
||||
Addressable: [1m[0m00[1m[0m
|
||||
Partially addressable: [1m[0m01[1m[0m [1m[0m02[1m[0m [1m[0m03[1m[0m [1m[0m04[1m[0m [1m[0m05[1m[0m [1m[0m06[1m[0m [1m[0m07[1m[0m
|
||||
Heap left redzone: [1m[31mfa[1m[0m
|
||||
Freed heap region: [1m[35mfd[1m[0m
|
||||
Stack left redzone: [1m[31mf1[1m[0m
|
||||
Stack mid redzone: [1m[31mf2[1m[0m
|
||||
Stack right redzone: [1m[31mf3[1m[0m
|
||||
Stack after return: [1m[35mf5[1m[0m
|
||||
Stack use after scope: [1m[35mf8[1m[0m
|
||||
Global redzone: [1m[31mf9[1m[0m
|
||||
Global init order: [1m[36mf6[1m[0m
|
||||
Poisoned by user: [1m[34mf7[1m[0m
|
||||
Container overflow: [1m[34mfc[1m[0m
|
||||
Array cookie: [1m[31mac[1m[0m
|
||||
Intra object redzone: [1m[33mbb[1m[0m
|
||||
ASan internal: [1m[33mfe[1m[0m
|
||||
Left alloca redzone: [1m[34mca[1m[0m
|
||||
Right alloca redzone: [1m[34mcb[1m[0m
|
||||
==14==ABORTING
|
||||
@@ -0,0 +1,4 @@
|
||||
weight: 1.0
|
||||
package_name: "selinux"
|
||||
harness_name: "secilc-fuzzer"
|
||||
task_id: "selinux-OSV-2021-902-delta-1"
|
||||
@@ -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",
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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]]:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
Generated
+45
-5
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user