Set maps size at compile time with env vars

libafl_targets/build.rs

@@ -1,7 +1,6 @@
 //! build.rs for `libafl_targets`
 
-use std::env;
+use std::{env, fs::File, io::Write, path::Path};
-use std::path::Path;
 
 fn main() {
     let out_dir = env::var_os("OUT_DIR").unwrap();
@@ -9,6 +8,44 @@ fn main() {
     //let out_dir_path = Path::new(&out_dir);
     let _src_dir = Path::new("src");
 
+    let dest_path = Path::new(&out_dir).join("constants.rs");
+    let mut constants_file = File::create(&dest_path).expect("Could not create file");
+
+    let edges_map_size: usize = option_env!("LIBAFL_EDGES_MAP_SIZE")
+        .map_or(Ok(65536), str::parse)
+        .expect("Could not parse LIBAFL_EDGES_MAP_SIZE");
+    let cmp_map_size: usize = option_env!("LIBAFL_CMP_MAP_SIZE")
+        .map_or(Ok(65536), str::parse)
+        .expect("Could not parse LIBAFL_CMP_MAP_SIZE");
+    let cmplog_map_w: usize = option_env!("LIBAFL_CMPLOG_MAP_W")
+        .map_or(Ok(65536), str::parse)
+        .expect("Could not parse LIBAFL_CMPLOG_MAP_W");
+    let cmplog_map_h: usize = option_env!("LIBAFL_CMPLOG_MAP_H")
+        .map_or(Ok(32), str::parse)
+        .expect("Could not parse LIBAFL_CMPLOG_MAP_H");
+
+    write!(
+        &mut constants_file,
+        "// These constants are autogenerated by build.rs
+
+/// The size of the edges map
+pub const EDGES_MAP_SIZE: usize = {};
+/// The size of the cmps map
+pub const CMP_MAP_SIZE: usize = {};
+/// The width of the CmpLog map
+pub const CMPLOG_MAP_W: usize = {};
+/// The height of the CmpLog map
+pub const CMPLOG_MAP_H: usize = {};
+",
+        edges_map_size, cmp_map_size, cmplog_map_w, cmplog_map_h
+    )
+    .expect("Could not write file");
+
+    println!("cargo:rerun-if-env-changed=LIBAFL_EDGES_MAP_SIZE");
+    println!("cargo:rerun-if-env-changed=LIBAFL_CMP_MAP_SIZE");
+    println!("cargo:rerun-if-env-changed=LIBAFL_CMPLOG_MAP_W");
+    println!("cargo:rerun-if-env-changed=LIBAFL_CMPLOG_MAP_H");
+
     //std::env::set_var("CC", "clang");
     //std::env::set_var("CXX", "clang++");
 
@@ -31,6 +68,9 @@ fn main() {
         }
 
         sancov_cmp
+            .define("CMP_MAP_SIZE", format!("{}", cmp_map_size))
+            .define("CMPLOG_MAP_W", format!("{}", cmplog_map_w))
+            .define("CMPLOG_MAP_H", format!("{}", cmplog_map_h))
             .file(_src_dir.join("sancov_cmp.c"))
             .compile("sancov_cmp");
     }

libafl_targets/src/cmplog.h

@@ -3,8 +3,12 @@
 
 #include "common.h"
 
+#ifndef CMPLOG_MAP_W
 #define CMPLOG_MAP_W 65536
+#endif
+#ifndef CMPLOG_MAP_H
 #define CMPLOG_MAP_H 32
+#endif
 
 #define CMPLOG_KIND_INS 0
 #define CMPLOG_KIND_RTN 1

libafl_targets/src/cmplog.rs

@@ -11,11 +11,8 @@ use libafl::{
 
 use serde::{Deserialize, Serialize};
 
-// TODO compile time flag
+use crate::{CMPLOG_MAP_H, CMPLOG_MAP_W};
-/// The `CmpLogMap` W value
+
-pub const CMPLOG_MAP_W: usize = 65536;
-/// The `CmpLogMap` H value
-pub const CMPLOG_MAP_H: usize = 32;
 /// The `CmpLog` map size
 pub const CMPLOG_MAP_SIZE: usize = CMPLOG_MAP_W * CMPLOG_MAP_H;
 

libafl_targets/src/coverage.rs

@@ -1,8 +1,6 @@
 //! Coverage maps as static mut array
 
-// TODO compile time flag
+use crate::EDGES_MAP_SIZE;
-/// The map size for the edges map.
-pub const EDGES_MAP_SIZE: usize = 65536;
 
 /// The map for edges.
 pub static mut EDGES_MAP: [u8; EDGES_MAP_SIZE] = [0; EDGES_MAP_SIZE];

libafl_targets/src/lib.rs

@@ -3,6 +3,8 @@
 #[macro_use]
 extern crate serde_big_array;
 
+include!(concat!(env!("OUT_DIR"), "/constants.rs"));
+
 #[cfg(any(feature = "sancov_pcguard_edges", feature = "sancov_pcguard_hitcounts"))]
 pub mod sancov_pcguard;
 #[cfg(any(feature = "sancov_pcguard_edges", feature = "sancov_pcguard_hitcounts"))]

libafl_targets/src/sancov_cmp.c

@@ -14,7 +14,7 @@ void __sanitizer_cov_trace_cmp1(uint8_t arg1, uint8_t arg2) {
   k = (k >> 4) ^ (k << 8);
 
 #ifdef SANCOV_VALUE_PROFILE
-  k &= MAP_SIZE - 1;
+  k &= CMP_MAP_SIZE - 1;
   __libafl_targets_value_profile1(k, arg1, arg2);
 #endif
 #ifdef SANCOV_CMPLOG
@@ -30,7 +30,7 @@ void __sanitizer_cov_trace_cmp2(uint16_t arg1, uint16_t arg2) {
   k = (k >> 4) ^ (k << 8);
 
 #ifdef SANCOV_VALUE_PROFILE
-  k &= MAP_SIZE - 1;
+  k &= CMP_MAP_SIZE - 1;
   __libafl_targets_value_profile2(k, arg1, arg2);
 #endif
 #ifdef SANCOV_CMPLOG
@@ -46,7 +46,7 @@ void __sanitizer_cov_trace_cmp4(uint32_t arg1, uint32_t arg2) {
   k = (k >> 4) ^ (k << 8);
 
 #ifdef SANCOV_VALUE_PROFILE
-  k &= MAP_SIZE - 1;
+  k &= CMP_MAP_SIZE - 1;
   __libafl_targets_value_profile4(k, arg1, arg2);
 #endif
 #ifdef SANCOV_CMPLOG
@@ -62,7 +62,7 @@ void __sanitizer_cov_trace_cmp8(uint64_t arg1, uint64_t arg2) {
   k = (k >> 4) ^ (k << 8);
 
 #ifdef SANCOV_VALUE_PROFILE
-  k &= MAP_SIZE - 1;
+  k &= CMP_MAP_SIZE - 1;
   __libafl_targets_value_profile8(k, arg1, arg2);
 #endif
 #ifdef SANCOV_CMPLOG
@@ -90,7 +90,7 @@ void __sanitizer_cov_trace_switch(uint64_t val, uint64_t *cases) {
     k = (k >> 4) ^ (k << 8);
     // val , cases[i + 2]
 #ifdef SANCOV_VALUE_PROFILE
-    k &= MAP_SIZE - 1;
+    k &= CMP_MAP_SIZE - 1;
     switch (cases[1]) {
         case 8:
         __libafl_targets_value_profile1(k, (uint8_t)val, (uint8_t)cases[i + 2]);

libafl_targets/src/sancov_pcguard.rs

@@ -41,7 +41,10 @@ pub unsafe extern "C" fn __sanitizer_cov_trace_pc_guard_init(mut start: *mut u32
 
     while start < stop {
         MAX_EDGES_NUM = MAX_EDGES_NUM.wrapping_add(1);
-        *start = (MAX_EDGES_NUM & (EDGES_MAP_SIZE - 1)) as u32;
+        if MAX_EDGES_NUM >= EDGES_MAP.len() {
+            panic!(format!("The number of edges reported by SanitizerCoverage exceed the size of the edges map ({}). Use the LIBAFL_EDGES_MAP_SIZE env to increase it at compile time.", EDGES_MAP.len()));
+        }
+        *start = MAX_EDGES_NUM as u32;
         start = start.offset(1);
     }
 }

libafl_targets/src/value_profile.h

@@ -3,10 +3,11 @@
 
 #include "common.h"
 
-// TODO compile time flag
+#ifndef CMP_MAP_SIZE
-#define MAP_SIZE 65536
+#define CMP_MAP_SIZE 65536
+#endif
 
-extern uint8_t libafl_cmp_map[MAP_SIZE];
+extern uint8_t libafl_cmp_map[CMP_MAP_SIZE];
 
 #ifdef _MSC_VER
   #include <intrin.h>

libafl_targets/src/value_profile.rs

@@ -1,8 +1,6 @@
 //! Value profile support for `LibAFL`
 
-// TODO compile time flag
+use crate::CMP_MAP_SIZE;
-/// The Cmp map size.
-pub const CMP_MAP_SIZE: usize = 65536;
 
 /// The constant cmplog map for the current `LibAFL` target
 #[no_mangle]