11#include " unicornafl.h"
22#include " config.h"
33#include " priv.h"
4+ #include " cmplog.h"
45
56#include < cstdio>
67#include < cstdlib>
@@ -398,9 +399,7 @@ class UCAFL {
398399 }
399400 }
400401
401- void _uc_hook_sub_impl (uint64_t cur_loc, uint64_t arg1, uint64_t arg2,
402- uint32_t size) {
403-
402+ void _uc_hook_cmpcov (uint64_t cur_loc, uint64_t arg1, uint64_t arg2, uint32_t size) {
404403 if (size >= 64 ) {
405404 if (unlikely (MAP_SIZE - cur_loc < 8 ))
406405 cur_loc -= 8 ;
@@ -416,6 +415,53 @@ class UCAFL {
416415 }
417416 }
418417
418+ void _uc_hook_cmplog (uint64_t cur_loc, uint64_t arg1, uint64_t arg2, uint32_t size) {
419+ struct cmp_map * map = this ->cmpmap_ ;
420+
421+ uintptr_t k = (uintptr_t )cur_loc & (CMP_MAP_W - 1 );
422+ uint32_t shape = 0 ;
423+ switch (size) {
424+ case 16 :
425+ shape = 1 ;
426+ break ;
427+ case 32 :
428+ shape = 3 ;
429+ break ;
430+ case 64 :
431+ shape = 7 ;
432+ break ;
433+ default :
434+ break ;
435+ }
436+
437+ // get hits count
438+ uint16_t hits;
439+ if (map->headers [k].type != CMP_TYPE_INS) {
440+ map->headers [k].type = CMP_TYPE_INS;
441+ map->headers [k].hits = 1 ;
442+ map->headers [k].shape = shape;
443+ hits = 0 ;
444+ } else {
445+ hits = map->headers [k].hits ++;
446+ if (map->headers [k].shape < shape) {
447+ map->headers [k].shape = shape;
448+ }
449+ }
450+
451+ hits &= CMP_MAP_H - 1 ;
452+ map->log [k][hits].v0 = arg1;
453+ map->log [k][hits].v1 = arg2;
454+ }
455+
456+ void _uc_hook_sub_impl (uint64_t cur_loc, uint64_t arg1, uint64_t arg2,
457+ uint32_t size) {
458+ if (this ->has_cmplog_ ) {
459+ return _uc_hook_cmplog (cur_loc, arg1, arg2, size);
460+ }
461+
462+ _uc_hook_cmpcov (cur_loc, arg1, arg2, size);
463+ }
464+
419465 static void _uc_hook_sub_cmp (uc_engine* uc, uint64_t address, uint64_t arg1,
420466 uint64_t arg2, uint32_t size,
421467 void * user_data) {
@@ -462,23 +508,27 @@ class UCAFL {
462508 exit (1 );
463509 }
464510
465- // These two hooks are for compcov and may not be supported by the arch.
466- err = uc_hook_add (this ->uc_ , &this ->h3_ , UC_HOOK_TCG_OPCODE,
467- (void *)_uc_hook_sub, (void *)this , 1 , 0 , UC_TCG_OP_SUB,
468- UC_TCG_OP_FLAG_DIRECT);
511+ // If we should use Laf-Intel or Red-Queen do add CMP hooks
512+ if (this ->has_cmpcov_ || this ->has_cmplog_ ) {
469513
470- if (err) {
471- ERR ( " Failed to setup UC_TCG_OP_SUB direct hook. \n " );
472- exit ( 1 );
473- }
514+ // These two hooks may not be supported by the arch.
515+ err = uc_hook_add ( this -> uc_ , & this -> h3_ , UC_HOOK_TCG_OPCODE,
516+ ( void *)_uc_hook_sub, ( void *) this , 1 , 0 , UC_TCG_OP_SUB,
517+ UC_TCG_OP_FLAG_DIRECT);
474518
475- err = uc_hook_add (this ->uc_ , &this ->h4_ , UC_HOOK_TCG_OPCODE,
476- (void *)_uc_hook_sub_cmp, (void *)this , 1 , 0 ,
477- UC_TCG_OP_SUB, UC_TCG_OP_FLAG_CMP);
519+ if (err) {
520+ ERR (" Failed to setup UC_TCG_OP_SUB direct hook.\n " );
521+ exit (1 );
522+ }
478523
479- if (err) {
480- ERR (" Failed to setup UC_TCG_OP_SUB cmp hook.\n " );
481- exit (1 );
524+ err = uc_hook_add (this ->uc_ , &this ->h4_ , UC_HOOK_TCG_OPCODE,
525+ (void *)_uc_hook_sub_cmp, (void *)this , 1 , 0 ,
526+ UC_TCG_OP_SUB, UC_TCG_OP_FLAG_CMP);
527+
528+ if (err) {
529+ ERR (" Failed to setup UC_TCG_OP_SUB cmp hook.\n " );
530+ exit (1 );
531+ }
482532 }
483533 }
484534
@@ -528,6 +578,29 @@ class UCAFL {
528578
529579 this ->has_afl_ = false ;
530580 }
581+
582+ char * cmplog_map_id_str = getenv (CMPLOG_SHM_ENV_VAR);
583+ this ->has_cmpcov_ = !!getenv (" UNICORN_AFL_CMPCOV" );
584+
585+ if (cmplog_map_id_str) {
586+ if (this ->has_cmpcov_ ) {
587+ ERR (" CMPLOG and CMPCOV turned on at the same time!\n " );
588+ ERR (" I'll turn off CMPCOV.\n " );
589+ this ->has_cmpcov_ = false ;
590+ }
591+
592+ int cmplog_map_id = atoi (cmplog_map_id_str);
593+ this ->cmpmap_ = (struct cmp_map *)shmat (cmplog_map_id, NULL , 0 );
594+
595+ if (this ->cmpmap_ == (void *)-1 ) {
596+ ERR (" Can't get the afl cmp-mapping area.\n " );
597+ exit (0 );
598+ }
599+
600+ this ->has_cmplog_ = true ;
601+ } else {
602+ this ->has_cmplog_ = false ;
603+ }
531604 }
532605
533606 uc_afl_ret _fksrv_start () {
@@ -543,7 +616,7 @@ class UCAFL {
543616 if (this ->afl_testcase_ptr_ ) {
544617 /* Parent supports testcases via shared map - and the user wants to
545618 * use it. Tell AFL. */
546- status = (FS_OPT_ENABLED | FS_OPT_SHDMEM_FUZZ);
619+ status = (FS_OPT_ENABLED | FS_OPT_SHDMEM_FUZZ | FS_OPT_NEWCMPLOG );
547620 /* Phone home and tell the parent that we're OK. If parent isn't
548621 there, assume we're not running in forkserver mode and just
549622 execute program. */
@@ -867,6 +940,10 @@ class UCAFL {
867940 uint32_t afl_inst_rms_;
868941 uint64_t afl_prev_loc_;
869942
943+ bool has_cmpcov_;
944+ bool has_cmplog_;
945+ struct cmp_map * cmpmap_;
946+
870947 // Fake signal value
871948 int wifsignaled_;
872949
0 commit comments