1 //===-- bytemap.h -----------------------------------------------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #ifndef SCUDO_BYTEMAP_H_
10 #define SCUDO_BYTEMAP_H_
11 
12 #include "atomic_helpers.h"
13 #include "common.h"
14 #include "mutex.h"
15 
16 namespace scudo {
17 
18 template <uptr Size> class FlatByteMap {
19 public:
20   void initLinkerInitialized() {
21     Map = reinterpret_cast<u8 *>(map(nullptr, Size, "scudo:bytemap"));
22   }
23   void init() { initLinkerInitialized(); }
24 
25   void unmapTestOnly() { unmap(reinterpret_cast<void *>(Map), Size); }
26 
27   void set(uptr Index, u8 Value) {
28     DCHECK_LT(Index, Size);
29     DCHECK_EQ(0U, Map[Index]);
30     Map[Index] = Value;
31   }
32   u8 operator[](uptr Index) {
33     DCHECK_LT(Index, Size);
34     return Map[Index];
35   }
36 
37   void disable() {}
38   void enable() {}
39 
40 private:
41   u8 *Map;
42 };
43 
44 template <uptr Level1Size, uptr Level2Size> class TwoLevelByteMap {
45 public:
46   void initLinkerInitialized() {
47     Level1Map = reinterpret_cast<atomic_uptr *>(
48         map(nullptr, sizeof(atomic_uptr) * Level1Size, "scudo:bytemap"));
49   }
50   void init() {
51     Mutex.init();
52     initLinkerInitialized();
53   }
54 
55   void reset() {
56     for (uptr I = 0; I < Level1Size; I++) {
57       u8 *P = get(I);
58       if (!P)
59         continue;
60       unmap(P, Level2Size);
61     }
62     memset(Level1Map, 0, sizeof(atomic_uptr) * Level1Size);
63   }
64 
65   void unmapTestOnly() {
66     reset();
67     unmap(reinterpret_cast<void *>(Level1Map),
68           sizeof(atomic_uptr) * Level1Size);
69   }
70 
71   uptr size() const { return Level1Size * Level2Size; }
72 
73   void set(uptr Index, u8 Value) {
74     DCHECK_LT(Index, Level1Size * Level2Size);
75     u8 *Level2Map = getOrCreate(Index / Level2Size);
76     DCHECK_EQ(0U, Level2Map[Index % Level2Size]);
77     Level2Map[Index % Level2Size] = Value;
78   }
79 
80   u8 operator[](uptr Index) const {
81     DCHECK_LT(Index, Level1Size * Level2Size);
82     u8 *Level2Map = get(Index / Level2Size);
83     if (!Level2Map)
84       return 0;
85     return Level2Map[Index % Level2Size];
86   }
87 
88   void disable() { Mutex.lock(); }
89   void enable() { Mutex.unlock(); }
90 
91 private:
92   u8 *get(uptr Index) const {
93     DCHECK_LT(Index, Level1Size);
94     return reinterpret_cast<u8 *>(
95         atomic_load(&Level1Map[Index], memory_order_acquire));
96   }
97 
98   u8 *getOrCreate(uptr Index) {
99     u8 *Res = get(Index);
100     if (!Res) {
101       ScopedLock L(Mutex);
102       if (!(Res = get(Index))) {
103         Res = reinterpret_cast<u8 *>(map(nullptr, Level2Size, "scudo:bytemap"));
104         atomic_store(&Level1Map[Index], reinterpret_cast<uptr>(Res),
105                      memory_order_release);
106       }
107     }
108     return Res;
109   }
110 
111   atomic_uptr *Level1Map;
112   HybridMutex Mutex;
113 };
114 
115 } // namespace scudo
116 
117 #endif // SCUDO_BYTEMAP_H_
118