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
14namespace {
15
16using ::__memprof::MIBMapTy;
17using ::__memprof::SerializeToRawProfile;
18using ::__sanitizer::MemoryMappedSegment;
19using ::__sanitizer::MemoryMappingLayoutBase;
20using ::__sanitizer::StackDepotPut;
21using ::__sanitizer::StackTrace;
22using ::llvm::memprof::MemInfoBlock;
23using ::testing::_;
24using ::testing::Action;
25using ::testing::DoAll;
26using ::testing::Return;
27using ::testing::SetArgPointee;
28
29class MockMemoryMappingLayout final : public MemoryMappingLayoutBase {
30public:
31  MOCK_METHOD(bool, Next, (MemoryMappedSegment *), (override));
32  MOCK_METHOD(void, Reset, (), (override));
33};
34
35uint64_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
49template <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
58TEST(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