retake from mem if possible

This commit is contained in:
van Hauser
2020-10-15 15:48:39 +02:00
parent 354bda2846
commit f41aafa4f7
3 changed files with 44 additions and 7 deletions

View File

@ -1165,6 +1165,10 @@ u8 *queue_testcase_get(afl_state_t *afl, struct queue_entry *q);
void queue_testcase_retake(afl_state_t *afl, struct queue_entry *q, void queue_testcase_retake(afl_state_t *afl, struct queue_entry *q,
u32 old_len); u32 old_len);
/* If trimming changes the testcase size we have to replace it */
void queue_testcase_retake_mem(afl_state_t *afl, struct queue_entry *q, u8 *in,
u32 len, u32 old_len);
#if TESTCASE_CACHE == 1 #if TESTCASE_CACHE == 1
#error define of TESTCASE_CACHE must be zero or larger than 1 #error define of TESTCASE_CACHE must be zero or larger than 1
#endif #endif

View File

@ -870,8 +870,10 @@ u32 calculate_score(afl_state_t *afl, struct queue_entry *q) {
} }
void queue_testcase_retake(afl_state_t *afl, struct queue_entry *q, /* after a custom trim we need to reload the testcase from disk */
u32 old_len) {
inline void queue_testcase_retake(afl_state_t *afl, struct queue_entry *q,
u32 old_len) {
if (likely(q->testcase_buf)) { if (likely(q->testcase_buf)) {
@ -879,9 +881,9 @@ void queue_testcase_retake(afl_state_t *afl, struct queue_entry *q,
if (len != old_len) { if (len != old_len) {
afl->q_testcase_cache_size = afl->q_testcase_cache_size = afl->q_testcase_cache_size + len - old_len;
afl->q_testcase_cache_size + q->len - old_len;
q->testcase_buf = realloc(q->testcase_buf, len); q->testcase_buf = realloc(q->testcase_buf, len);
if (unlikely(!q->testcase_buf)) { if (unlikely(!q->testcase_buf)) {
PFATAL("Unable to malloc '%s' with len %d", q->fname, len); PFATAL("Unable to malloc '%s' with len %d", q->fname, len);
@ -901,8 +903,35 @@ void queue_testcase_retake(afl_state_t *afl, struct queue_entry *q,
} }
/* after a normal trim we need to replace the testcase with the new data */
inline void queue_testcase_retake_mem(afl_state_t *afl, struct queue_entry *q,
u8 *in, u32 len, u32 old_len) {
if (likely(q->testcase_buf)) {
if (len != old_len) {
afl->q_testcase_cache_size = afl->q_testcase_cache_size + len - old_len;
q->testcase_buf = realloc(q->testcase_buf, len);
if (unlikely(!q->testcase_buf)) {
PFATAL("Unable to malloc '%s' with len %d", q->fname, len);
}
}
memcpy(q->testcase_buf, in, len);
}
}
/* Returns the testcase buf from the file behind this queue entry. /* Returns the testcase buf from the file behind this queue entry.
Increases the refcount. */ Increases the refcount. */
inline u8 *queue_testcase_get(afl_state_t *afl, struct queue_entry *q) { inline u8 *queue_testcase_get(afl_state_t *afl, struct queue_entry *q) {
u32 len = q->len; u32 len = q->len;
@ -913,7 +942,7 @@ inline u8 *queue_testcase_get(afl_state_t *afl, struct queue_entry *q) {
u8 *buf; u8 *buf;
if (q == afl->queue_cur) { if (unlikely(q == afl->queue_cur)) {
buf = afl_realloc((void **)&afl->testcase_buf, len); buf = afl_realloc((void **)&afl->testcase_buf, len);

View File

@ -711,7 +711,11 @@ u8 trim_case(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
}); });
if (orig_len != q->len) { queue_testcase_retake(afl, q, orig_len); } if (orig_len != q->len || custom_trimmed) {
queue_testcase_retake(afl, q, orig_len);
}
if (custom_trimmed) return trimmed_case; if (custom_trimmed) return trimmed_case;
@ -846,7 +850,7 @@ u8 trim_case(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
close(fd); close(fd);
if (orig_len != q->len) queue_testcase_retake(afl, q, orig_len); queue_testcase_retake_mem(afl, q, in_buf, q->len, orig_len);
memcpy(afl->fsrv.trace_bits, afl->clean_trace, afl->fsrv.map_size); memcpy(afl->fsrv.trace_bits, afl->clean_trace, afl->fsrv.map_size);
update_bitmap_score(afl, q); update_bitmap_score(afl, q);