1 #include "memprof/memprof_rawprofile.h"
2 
3 #include <cstdint>
4 #include <memory>
5 
6 #include "profile/MemProfData.inc"
7 #include "sanitizer_common/sanitizer_common.h"
8 #include "sanitizer_common/sanitizer_procmaps.h"
9 #include "sanitizer_common/sanitizer_stackdepot.h"
10 #include "sanitizer_common/sanitizer_stacktrace.h"
11 #include "gmock/gmock.h"
12 #include "gtest/gtest.h"
13 
14 namespace {
15 
16 using ::__memprof::MIBMapTy;
17 using ::__memprof::SerializeToRawProfile;
18 using ::__sanitizer::MemoryMappedSegment;
19 using ::__sanitizer::MemoryMappingLayoutBase;
20 using ::__sanitizer::StackDepotPut;
21 using ::__sanitizer::StackTrace;
22 using ::llvm::memprof::MemInfoBlock;
23 using ::testing::_;
24 using ::testing::Action;
25 using ::testing::DoAll;
26 using ::testing::Return;
27 using ::testing::SetArgPointee;
28 
29 class MockMemoryMappingLayout final : public MemoryMappingLayoutBase {
30 public:
31   MOCK_METHOD(bool, Next, (MemoryMappedSegment *), (override));
32   MOCK_METHOD(void, Reset, (), (override));
33 };
34 
35 uint64_t PopulateFakeMap(const MemInfoBlock &FakeMIB, uint64_t StackPCBegin,
36                          MIBMapTy &FakeMap) {
37   constexpr int kSize = 5;
38   uint64_t array[kSize];
39   for (int i = 0; i < kSize; i++) {
40     array[i] = StackPCBegin + i;
41   }
42   StackTrace St(array, kSize);
43   uint32_t Id = StackDepotPut(St);
44 
45   InsertOrMerge(Id, FakeMIB, FakeMap);
46   return Id;
47 }
48 
49 template <class T = uint64_t> T Read(char *&Buffer) {
50   static_assert(std::is_pod<T>::value, "Must be a POD type.");
51   assert(reinterpret_cast<size_t>(Buffer) % sizeof(T) == 0 &&
52          "Unaligned read!");
53   T t = *reinterpret_cast<T *>(Buffer);
54   Buffer += sizeof(T);
55   return t;
56 }
57 
58 TEST(MemProf, Basic) {
59   MockMemoryMappingLayout Layout;
60   MemoryMappedSegment FakeSegment;
61   memset(&FakeSegment, 0, sizeof(FakeSegment));
62   FakeSegment.start = 0x10;
63   FakeSegment.end = 0x20;
64   FakeSegment.offset = 0x10;
65   uint8_t uuid[__sanitizer::kModuleUUIDSize] = {0xC, 0x0, 0xF, 0xF, 0xE, 0xE};
66   memcpy(FakeSegment.uuid, uuid, __sanitizer::kModuleUUIDSize);
67   FakeSegment.protection =
68       __sanitizer::kProtectionExecute | __sanitizer::kProtectionRead;
69 
70   const Action<bool(MemoryMappedSegment *)> SetSegment =
71       DoAll(SetArgPointee<0>(FakeSegment), Return(true));
72   EXPECT_CALL(Layout, Next(_))
73       .WillOnce(SetSegment)
74       .WillOnce(Return(false))
75       .WillOnce(SetSegment)
76       .WillRepeatedly(Return(false));
77 
78   EXPECT_CALL(Layout, Reset).Times(2);
79 
80   MIBMapTy FakeMap;
81   MemInfoBlock FakeMIB;
82   // Since we want to override the constructor set vals to make it easier to
83   // test.
84   memset(&FakeMIB, 0, sizeof(MemInfoBlock));
85   FakeMIB.AllocCount = 0x1;
86   FakeMIB.TotalAccessCount = 0x2;
87 
88   uint64_t FakeIds[2];
89   FakeIds[0] = PopulateFakeMap(FakeMIB, /*StackPCBegin=*/2, FakeMap);
90   FakeIds[1] = PopulateFakeMap(FakeMIB, /*StackPCBegin=*/3, FakeMap);
91 
92   char *Ptr = nullptr;
93   uint64_t NumBytes = SerializeToRawProfile(FakeMap, Layout, Ptr);
94   const char *Buffer = Ptr;
95 
96   ASSERT_GT(NumBytes, 0ULL);
97   ASSERT_TRUE(Ptr);
98 
99   // Check the header.
100   EXPECT_THAT(Read(Ptr), MEMPROF_RAW_MAGIC_64);
101   EXPECT_THAT(Read(Ptr), MEMPROF_RAW_VERSION);
102   const uint64_t TotalSize = Read(Ptr);
103   const uint64_t SegmentOffset = Read(Ptr);
104   const uint64_t MIBOffset = Read(Ptr);
105   const uint64_t StackOffset = Read(Ptr);
106 
107   // ============= Check sizes and padding.
108   EXPECT_EQ(TotalSize, NumBytes);
109   EXPECT_EQ(TotalSize % 8, 0ULL);
110 
111   // Should be equal to the size of the raw profile header.
112   EXPECT_EQ(SegmentOffset, 48ULL);
113 
114   // We expect only 1 segment entry, 8b for the count and 56b for SegmentEntry
115   // in memprof_rawprofile.cpp.
116   EXPECT_EQ(MIBOffset - SegmentOffset, 64ULL);
117 
118   EXPECT_EQ(MIBOffset, 112ULL);
119   // We expect 2 mib entry, 8b for the count and sizeof(uint64_t) +
120   // sizeof(MemInfoBlock) contains stack id + MeminfoBlock.
121   EXPECT_EQ(StackOffset - MIBOffset, 8 + 2 * (8 + sizeof(MemInfoBlock)));
122 
123   EXPECT_EQ(StackOffset, 336ULL);
124   // We expect 2 stack entries, with 5 frames - 8b for total count,
125   // 2 * (8b for id, 8b for frame count and 5*8b for fake frames).
126   // Since this is the last section, there may be additional padding at the end
127   // to make the total profile size 8b aligned.
128   EXPECT_GE(TotalSize - StackOffset, 8ULL + 2 * (8 + 8 + 5 * 8));
129 
130   // ============= Check contents.
131   // The Uuid field is not yet populated on Linux-Elf by the sanitizer procmaps
132   // library, so we expect it to be filled with 0 for now.
133   unsigned char ExpectedSegmentBytes[64] = {
134       0x01, 0, 0, 0, 0, 0, 0, 0, // Number of entries
135       0x10, 0, 0, 0, 0, 0, 0, 0, // Start
136       0x20, 0, 0, 0, 0, 0, 0, 0, // End
137       0x10, 0, 0, 0, 0, 0, 0, 0, // Offset
138       0x0,                       // Uuid
139   };
140   EXPECT_EQ(memcmp(Buffer + SegmentOffset, ExpectedSegmentBytes, 64), 0);
141 
142   // Check that the number of entries is 2.
143   EXPECT_EQ(*reinterpret_cast<const uint64_t *>(Buffer + MIBOffset), 2ULL);
144   // Check that stack id is set.
145   EXPECT_EQ(*reinterpret_cast<const uint64_t *>(Buffer + MIBOffset + 8),
146             FakeIds[0]);
147 
148   // Only check a few fields of the first MemInfoBlock.
149   unsigned char ExpectedMIBBytes[sizeof(MemInfoBlock)] = {
150       0x01, 0, 0, 0, // Alloc count
151       0x02, 0, 0, 0, // Total access count
152   };
153   // Compare contents of 1st MIB after skipping count and stack id.
154   EXPECT_EQ(
155       memcmp(Buffer + MIBOffset + 16, ExpectedMIBBytes, sizeof(MemInfoBlock)),
156       0);
157   // Compare contents of 2nd MIB after skipping count and stack id for the first
158   // and only the id for the second.
159   EXPECT_EQ(memcmp(Buffer + MIBOffset + 16 + sizeof(MemInfoBlock) + 8,
160                    ExpectedMIBBytes, sizeof(MemInfoBlock)),
161             0);
162 
163   // Check that the number of entries is 2.
164   EXPECT_EQ(*reinterpret_cast<const uint64_t *>(Buffer + StackOffset), 2ULL);
165   // Check that the 1st stack id is set.
166   EXPECT_EQ(*reinterpret_cast<const uint64_t *>(Buffer + StackOffset + 8),
167             FakeIds[0]);
168   // Contents are num pcs, value of each pc - 1.
169   unsigned char ExpectedStackBytes[2][6 * 8] = {
170       {
171           0x5, 0, 0, 0, 0, 0, 0, 0, // Number of PCs
172           0x1, 0, 0, 0, 0, 0, 0, 0, // PC ...
173           0x2, 0, 0, 0, 0, 0, 0, 0, 0x3, 0, 0, 0, 0, 0, 0, 0,
174           0x4, 0, 0, 0, 0, 0, 0, 0, 0x5, 0, 0, 0, 0, 0, 0, 0,
175       },
176       {
177           0x5, 0, 0, 0, 0, 0, 0, 0, // Number of PCs
178           0x2, 0, 0, 0, 0, 0, 0, 0, // PC ...
179           0x3, 0, 0, 0, 0, 0, 0, 0, 0x4, 0, 0, 0, 0, 0, 0, 0,
180           0x5, 0, 0, 0, 0, 0, 0, 0, 0x6, 0, 0, 0, 0, 0, 0, 0,
181       },
182   };
183   EXPECT_EQ(memcmp(Buffer + StackOffset + 16, ExpectedStackBytes[0],
184                    sizeof(ExpectedStackBytes[0])),
185             0);
186 
187   // Check that the 2nd stack id is set.
188   EXPECT_EQ(
189       *reinterpret_cast<const uint64_t *>(Buffer + StackOffset + 8 + 6 * 8 + 8),
190       FakeIds[1]);
191 
192   EXPECT_EQ(memcmp(Buffer + StackOffset + 16 + 6 * 8 + 8, ExpectedStackBytes[1],
193                    sizeof(ExpectedStackBytes[1])),
194             0);
195 }
196 
197 } // namespace
198