include/afl-prealloc.h


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

/*
   american fuzzy lop++ - prealloc a buffer to reuse small elements often
   ----------------------------------------------------------------------

   Originally written by Michal Zalewski

   Now maintained by Marc Heuse <mh@mh-sec.de>,
                     Heiko Eißfeldt <heiko.eissfeldt@hexco.de>,
                     Andrea Fioraldi <andreafioraldi@gmail.com>,
                     Dominik Maier <mail@dmnk.co>

   Copyright 2016, 2017 Google Inc. All rights reserved.
   Copyright 2019-2023 AFLplusplus Project. All rights reserved.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at:

     https://www.apache.org/licenses/LICENSE-2.0

 */

/* If we know we'll reuse small elements often, we'll just preallocate a buffer,
 * then fall back to malloc */
// TODO: Replace free status check with bitmask+CLZ

#ifndef AFL_PREALLOC_H
#define AFL_PREALLOC_H

#include <stdio.h>
#include <stdbool.h>
#include <string.h>

#include "debug.h"
#include "alloc-inl.h"

typedef enum prealloc_status {

  PRE_STATUS_UNUSED = 0,                                     /* free in buf */
  PRE_STATUS_USED,                                           /* used in buf */
  PRE_STATUS_MALLOC                                        /* system malloc */

} pre_status_t;

/* Adds the entry used for prealloc bookkeeping to this struct */

/* prealloc status of this instance */
#define PREALLOCABLE pre_status_t pre_status

/* allocate an element of type *el_ptr, to this variable.
    Uses (and reuses) the given prealloc_buf before hitting libc's malloc.
    prealloc_buf must be the pointer to an array with type `type`.
    `type` must be a struct with uses PREALLOCABLE (a pre_status_t pre_status
   member). prealloc_size must be the array size. prealloc_counter must be a
   variable initialized with 0 (of any name).
    */

#define PRE_ALLOC(el_ptr, prealloc_buf, prealloc_size, prealloc_counter)       \
  do {                                                                         \
                                                                               \
    if ((prealloc_counter) >= (prealloc_size)) {                               \
                                                                               \
      el_ptr = (element_t *)malloc(sizeof(*el_ptr));                           \
      if (!el_ptr) { FATAL("error in list.h -> out of memory for element!"); } \
      el_ptr->pre_status = PRE_STATUS_MALLOC;                                  \
                                                                               \
    } else {                                                                   \
                                                                               \
      /* Find one of our preallocated elements */                              \
      u32 i;                                                                   \
      for (i = 0; i < (prealloc_size); i++) {                                  \
                                                                               \
        el_ptr = &((prealloc_buf)[i]);                                         \
        if (el_ptr->pre_status == PRE_STATUS_UNUSED) {                         \
                                                                               \
          (prealloc_counter)++;                                                \
          el_ptr->pre_status = PRE_STATUS_USED;                                \
          break;                                                               \
                                                                               \
        }                                                                      \
                                                                               \
      }                                                                        \
                                                                               \
    }                                                                          \
                                                                               \
    if (!el_ptr) { FATAL("BUG in list.h -> no element found or allocated!"); } \
                                                                               \
  } while (0);

/* Take a chosen (free) element from the prealloc_buf directly */

#define PRE_ALLOC_FORCE(el_ptr, prealloc_counter)         \
  do {                                                    \
                                                          \
    if ((el_ptr)->pre_status != PRE_STATUS_UNUSED) {      \
                                                          \
      FATAL("PRE_ALLOC_FORCE element already allocated"); \
                                                          \
    }                                                     \
    (el_ptr)->pre_status = PRE_STATUS_USED;               \
    (prealloc_counter)++;                                 \
                                                          \
  } while (0);

/* free an preallocated element */

#define PRE_FREE(el_ptr, prealloc_counter)        \
  do {                                            \
                                                  \
    switch ((el_ptr)->pre_status) {               \
                                                  \
      case PRE_STATUS_USED: {                     \
                                                  \
        (el_ptr)->pre_status = PRE_STATUS_UNUSED; \
        (prealloc_counter)--;                     \
        if ((prealloc_counter) < 0) {             \
                                                  \
          FATAL("Inconsistent data in PRE_FREE"); \
                                                  \
        }                                         \
        break;                                    \
                                                  \
      }                                           \
      case PRE_STATUS_MALLOC: {                   \
                                                  \
        (el_ptr)->pre_status = PRE_STATUS_UNUSED; \
        DFL_ck_free((el_ptr));                    \
        break;                                    \
                                                  \
      }                                           \
      default: {                                  \
                                                  \
        FATAL("Double Free Detected");            \
        break;                                    \
                                                  \
      }                                           \
                                                  \
    }                                             \
                                                  \
  } while (0);

#endif