[GWP-ASan] Implement stack frame compression.

[hctim]

Summary

This patch introduces stack frame compression to GWP-ASan. Each stack frame is
variable-length integer encoded as the difference between frame[i] and
frame[i - 1]. Furthermore, we use zig-zag encoding on the difference to ensure
that negative differences are also encoded into a relatively small number of
bytes.

Examples of what the compression looks like can be seen in
gwp_asan/tests/compression.cpp.

This compression can reduce the memory consumption cost of stack traces by
~50%.

[vlad.tsyrklevich]

compiler-rt/lib/gwp_asan/guarded_pool_allocator.cpp

78 ↗

(On Diff #214990)

Use a VLA instead of allocas, here and below.

compiler-rt/lib/gwp_asan/options.inc

28 ↗

(On Diff #214990)

Why make this an option? I don't see why users would need to tune it, I'd set it to something reasonable (64?) and just leave it as a constant.

compiler-rt/lib/gwp_asan/stack_trace_compressor.cpp

10 ↗

(On Diff #214990)

Most of these implementations vary from the ones in Chrome, I'd consider just using the same ones (modulo style differences) for consistency.

compiler-rt/lib/gwp_asan/stack_trace_compressor.h

9 ↗

(On Diff #214990)

COMPRESSION -> COMPRESSOR

[hctim]

• Updated with Vlad's comments.

compiler-rt/lib/gwp_asan/guarded_pool_allocator.cpp

78 ↗

(On Diff #214990)

Unfortunately VLA with non-constexpr size is C++14 only, and IMHO adding explicit C++14 requirements in the CMake file is more ugly than just using alloca here. LMK if you're really keen on a different strategy.

compiler-rt/lib/gwp_asan/options.inc

28 ↗

(On Diff #214990)

Agreed that runtime users shouldn't generally need to tune this, but this is more for the supporting allocator to tune. I'd expect that bionic on Android has vastly different requirements for scudo on fuchsia for example, and this allows them to tune it in the supporting allocator. Scudo currently exposes this to the user (as they use optional/options_parser.cpp from sanitizer-common to enumaret the options), but allocators don't have to expose this option to the user.

compiler-rt/lib/gwp_asan/stack_trace_compressor.cpp

10 ↗

(On Diff #214990)

I've done this and updated the tests. Original code didn't zigzag encode the first pointer, only the subsequent diffs.

[vlad.tsyrklevich]

compiler-rt/lib/gwp_asan/options.inc

28 ↗

(On Diff #214990)

I don't understand how the bionic/scudo examples have different requirements. Why would an embedding allocator need to tune this? The actual amount of space (and hence the approximate number of frames) is already constant, all they can adjust is collecting fewer frames.

[hctim]

compiler-rt/lib/gwp_asan/options.inc

28 ↗

(On Diff #214990)

That's a good point.

Maybe they should also be able to control the number of storage bytes as well?

[vlad.tsyrklevich]

compiler-rt/lib/gwp_asan/options.inc

28 ↗

(On Diff #214990)

Maybe in the future if we implement MESH and allocations pages don't dominate memory consumption, as it stands now the stack traces are a small part of the memory consumption so I would just leave it constant until a need arises. (This also has the side benefit of being able to replace the allocas with fixed-size arrays.)

[hctim]

• Removed ability to set trace length.

compiler-rt/lib/gwp_asan/options.inc

28 ↗

(On Diff #214990)

Ack. Have gone ahead and done this.

[leonardchan]

Hi. We believe this patch is causing a CMake error for out toolchian build:

-- Compiler-RT supported architectures: i386
CMake Error at /b/s/w/ir/k/llvm-project/compiler-rt/lib/gwp_asan/CMakeLists.txt:111 (target_link_options):
  Unknown CMake command "target_link_options".

Could you look into this? Thanks.

[hctim]

In D66189#1632302, @leonardchan wrote:

Hi. We believe this patch is causing a CMake error for out toolchian build:

-- Compiler-RT supported architectures: i386
CMake Error at /b/s/w/ir/k/llvm-project/compiler-rt/lib/gwp_asan/CMakeLists.txt:111 (target_link_options):
  Unknown CMake command "target_link_options".

Could you look into this? Thanks.

Hey Leonard, that should have been fixed in rL369051. Can you apply that patch and let me know?

[leonardchan]

In D66189#1632318, @hctim wrote:

In D66189#1632302, @leonardchan wrote:

Hi. We believe this patch is causing a CMake error for out toolchian build:

-- Compiler-RT supported architectures: i386
CMake Error at /b/s/w/ir/k/llvm-project/compiler-rt/lib/gwp_asan/CMakeLists.txt:111 (target_link_options):
  Unknown CMake command "target_link_options".

Could you look into this? Thanks.

Hey Leonard, that should have been fixed in rL369051. Can you apply that patch and let me know?

Looks like it was applied for us, but we're still seeing the same error.

[leonardchan]

Let me know if there's anything that you need from me to help diagnose this.

[hctim]

Looks like rL369051 was accidentally reverted by rL369055. Have fixed it in rL369067.

[eugenis]

This build is on r369069 and still has the problem:
http://lab.llvm.org:8011/builders/sanitizer-x86_64-linux-android/builds/23577

[2195/2261] Performing build step for 'compiler-rt'
[0/1] Re-running CMake...

• Compiler-RT supported architectures: x86_64
• Builtin supported architectures: x86_64

CMake Error at lib/gwp_asan/CMakeLists.txt:104 (add_llvm_executable):

Unknown CMake command "add_llvm_executable".

[hctim]

Commit is still causing some build issues. I'll take a look shortly and rollback if necessary.

[aheejin]

This comment was removed by aheejin.

[aheejin]

I posted a bug report but it turned out to be a bug with -D_GLIBCXX_DEBUG. Sorry, please ignore my previous comment.

compiler-rt/trunk/lib/gwp_asan/CMakeLists.txt

 add_compiler_rt_component(gwp_asan)
 
 include_directories(..)
 
 set(GWP_ASAN_SOURCES
   platform_specific/guarded_pool_allocator_posix.cpp
   platform_specific/mutex_posix.cpp
   guarded_pool_allocator.cpp
   random.cpp
+  stack_trace_compressor.cpp
 )
 
 set(GWP_ASAN_HEADERS
   definitions.h
   guarded_pool_allocator.h
   mutex.h
   options.h
   options.inc
   random.h
+  stack_trace_compressor.h
 )
 
 # Ensure that GWP-ASan meets the delegated requirements of some supporting
 # allocators. Some supporting allocators (e.g. scudo standalone) cannot use any
 # parts of the C++ standard library.
 set(GWP_ASAN_CFLAGS -fno-rtti -fno-exceptions -nostdinc++ -pthread)
 append_list_if(COMPILER_RT_HAS_FPIC_FLAG -fPIC GWP_ASAN_CFLAGS)
 append_list_if(COMPILER_RT_HAS_OMIT_FRAME_POINTER_FLAG -fno-omit-frame-pointer
add_compiler_rt_object_libraries(RTGwpAsanBacktraceLibc
       SOURCES optional/backtrace_linux_libc.cpp
       ADDITIONAL_HEADERS ${GWP_ASAN_BACKTRACE_HEADERS}
       CFLAGS ${GWP_ASAN_CFLAGS})
   add_compiler_rt_object_libraries(RTGwpAsanBacktraceSanitizerCommon
       ARCHS ${GWP_ASAN_SUPPORTED_ARCH}
       SOURCES optional/backtrace_sanitizer_common.cpp
       ADDITIONAL_HEADERS ${GWP_ASAN_BACKTRACE_HEADERS}
       CFLAGS ${GWP_ASAN_CFLAGS} ${SANITIZER_COMMON_CFLAGS})
+
+  # Build the stack trace compressor fuzzer.
+  add_llvm_executable(stack_trace_compressor_fuzzer
+    stack_trace_compressor_fuzzer.cpp
+    ${GWP_ASAN_SOURCES}
+    ${GWP_ASAN_HEADERS})
+  set_target_properties(stack_trace_compressor_fuzzer
+    PROPERTIES FOLDER "Fuzzers")
+  target_compile_options(stack_trace_compressor_fuzzer
+    PRIVATE -fsanitize=fuzzer-no-link)
+  target_link_options(stack_trace_compressor_fuzzer PRIVATE -fsanitize=fuzzer)
+  add_dependencies(gwp_asan stack_trace_compressor_fuzzer)
 endif()
 
 if(COMPILER_RT_INCLUDE_TESTS)
   add_subdirectory(tests)
 endif()

compiler-rt/trunk/lib/gwp_asan/guarded_pool_allocator.h

 //===-- guarded_pool_allocator.h --------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef GWP_ASAN_GUARDED_POOL_ALLOCATOR_H_
 #define GWP_ASAN_GUARDED_POOL_ALLOCATOR_H_
 
 #include "gwp_asan/definitions.h"
 #include "gwp_asan/mutex.h"
 #include "gwp_asan/options.h"
 #include "gwp_asan/random.h"
+#include "gwp_asan/stack_trace_compressor.h"
 
 #include <stddef.h>
 #include <stdint.h>
 
 namespace gwp_asan {
 // This class is the primary implementation of the allocator portion of GWP-
 // ASan. It is the sole owner of the pool of sequentially allocated guarded
 // slots. It should always be treated as a singleton.
enum class Error {
     USE_AFTER_FREE,
     DOUBLE_FREE,
     INVALID_FREE,
     BUFFER_OVERFLOW,
     BUFFER_UNDERFLOW
   };
 
   struct AllocationMetadata {
-    // Maximum number of stack trace frames to collect for allocations + frees.
-    // TODO(hctim): Implement stack frame compression, a-la Chromium.
-    static constexpr size_t kMaximumStackFrames = 64;
+    // The number of bytes used to store a compressed stack frame. On 64-bit
+    // platforms, assuming a compression ratio of 50%, this should allow us to
+    // store ~64 frames per trace.
+    static constexpr size_t kStackFrameStorageBytes = 256;
+
+    // Maximum number of stack frames to collect on allocation/deallocation. The
+    // actual number of collected frames may be less than this as the stack
+    // frames are compressed into a fixed memory range.
+    static constexpr size_t kMaxTraceLengthToCollect = 128;
 
     // Records the given allocation metadata into this struct.
     void RecordAllocation(uintptr_t Addr, size_t Size,
                           options::Backtrace_t Backtrace);
 
     // Record that this allocation is now deallocated.
     void RecordDeallocation(options::Backtrace_t Backtrace);
 
     struct CallSiteInfo {
-      // The backtrace to the allocation/deallocation.
-      uintptr_t Trace[kMaximumStackFrames] = {};
+      // The compressed backtrace to the allocation/deallocation.
+      uint8_t CompressedTrace[kStackFrameStorageBytes];
       // The thread ID for this trace, or kInvalidThreadID if not available.
       uint64_t ThreadID = kInvalidThreadID;
-      // The length of the trace. Zero indicates that no trace was collected.
-      size_t TraceLength = 0;
+      // The size of the compressed trace (in bytes). Zero indicates that no
+      // trace was collected.
+      size_t TraceSize = 0;
     };
 
     // The address of this allocation.
     uintptr_t Addr = 0;
     // Represents the actual size of the allocation.
     size_t Size = 0;
 
     CallSiteInfo AllocationTrace;

compiler-rt/trunk/lib/gwp_asan/guarded_pool_allocator.cpp

void GuardedPoolAllocator::AllocationMetadata::RecordAllocation(
     uintptr_t AllocAddr, size_t AllocSize, options::Backtrace_t Backtrace) {
   Addr = AllocAddr;
   Size = AllocSize;
   IsDeallocated = false;
 
   // TODO(hctim): Ask the caller to provide the thread ID, so we don't waste
   // other thread's time getting the thread ID under lock.
   AllocationTrace.ThreadID = getThreadID();
-  AllocationTrace.TraceLength = 0;
-  DeallocationTrace.TraceLength = 0;
+  AllocationTrace.TraceSize = 0;
+  DeallocationTrace.TraceSize = 0;
   DeallocationTrace.ThreadID = kInvalidThreadID;
-  if (Backtrace)
-    AllocationTrace.TraceLength =
-        Backtrace(AllocationTrace.Trace, kMaximumStackFrames);
+
+  if (Backtrace) {
+    uintptr_t UncompressedBuffer[kMaxTraceLengthToCollect];
+    size_t BacktraceLength =
+        Backtrace(UncompressedBuffer, kMaxTraceLengthToCollect);
+    AllocationTrace.TraceSize = compression::pack(
+        UncompressedBuffer, BacktraceLength, AllocationTrace.CompressedTrace,
+        kStackFrameStorageBytes);
+  }
 }
 
 void GuardedPoolAllocator::AllocationMetadata::RecordDeallocation(
     options::Backtrace_t Backtrace) {
   IsDeallocated = true;
   // Ensure that the unwinder is not called if the recursive flag is set,
   // otherwise non-reentrant unwinders may deadlock.
-  DeallocationTrace.TraceLength = 0;
+  DeallocationTrace.TraceSize = 0;
   if (Backtrace && !ThreadLocals.RecursiveGuard) {
     ScopedBoolean B(ThreadLocals.RecursiveGuard);
-    DeallocationTrace.TraceLength =
-        Backtrace(DeallocationTrace.Trace, kMaximumStackFrames);
+
+    uintptr_t UncompressedBuffer[kMaxTraceLengthToCollect];
+    size_t BacktraceLength =
+        Backtrace(UncompressedBuffer, kMaxTraceLengthToCollect);
+    DeallocationTrace.TraceSize = compression::pack(
+        UncompressedBuffer, BacktraceLength, DeallocationTrace.CompressedTrace,
+        kStackFrameStorageBytes);
   }
   DeallocationTrace.ThreadID = getThreadID();
 }
 
 void GuardedPoolAllocator::init(const options::Options &Opts) {
   // Note: We return from the constructor here if GWP-ASan is not available.
   // This will stop heap-allocation of class members, as well as mmap() of the
   // guarded slots.
void printAllocDeallocTraces(uintptr_t AccessPtr, AllocationMetadata *Meta,
   if (Meta->IsDeallocated) {
     if (Meta->DeallocationTrace.ThreadID ==
         GuardedPoolAllocator::kInvalidThreadID)
       Printf("0x%zx was deallocated by thread <unknown> here:\n", AccessPtr);
     else
       Printf("0x%zx was deallocated by thread %zu here:\n", AccessPtr,
              Meta->DeallocationTrace.ThreadID);
 
-    PrintBacktrace(Meta->DeallocationTrace.Trace,
-                   Meta->DeallocationTrace.TraceLength, Printf);
+    uintptr_t UncompressedTrace[AllocationMetadata::kMaxTraceLengthToCollect];
+    size_t UncompressedLength = compression::unpack(
+        Meta->DeallocationTrace.CompressedTrace,
+        Meta->DeallocationTrace.TraceSize, UncompressedTrace,
+        AllocationMetadata::kMaxTraceLengthToCollect);
+
+    PrintBacktrace(UncompressedTrace, UncompressedLength, Printf);
   }
 
   if (Meta->AllocationTrace.ThreadID == GuardedPoolAllocator::kInvalidThreadID)
     Printf("0x%zx was allocated by thread <unknown> here:\n", Meta->Addr);
   else
     Printf("0x%zx was allocated by thread %zu here:\n", Meta->Addr,
            Meta->AllocationTrace.ThreadID);
 
-  PrintBacktrace(Meta->AllocationTrace.Trace, Meta->AllocationTrace.TraceLength,
-                 Printf);
+  uintptr_t UncompressedTrace[AllocationMetadata::kMaxTraceLengthToCollect];
+  size_t UncompressedLength = compression::unpack(
+      Meta->AllocationTrace.CompressedTrace, Meta->AllocationTrace.TraceSize,
+      UncompressedTrace, AllocationMetadata::kMaxTraceLengthToCollect);
+
+  PrintBacktrace(UncompressedTrace, UncompressedLength, Printf);
 }
 
 struct ScopedEndOfReportDecorator {
   ScopedEndOfReportDecorator(options::Printf_t Printf) : Printf(Printf) {}
   ~ScopedEndOfReportDecorator() { Printf("*** End GWP-ASan report ***\n"); }
   options::Printf_t Printf;
 };
 } // anonymous namespace

compiler-rt/trunk/lib/gwp_asan/stack_trace_compressor.h

+//===-- stack_trace_compressor.h --------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef GWP_ASAN_STACK_TRACE_COMPRESSOR_
+#define GWP_ASAN_STACK_TRACE_COMPRESSOR_
+
+#include <stddef.h>
+#include <stdint.h>
+
+// These functions implement stack frame compression and decompression. We store
+// the zig-zag encoded pointer difference between frame[i] and frame[i - 1] as
+// a variable-length integer. This can reduce the memory overhead of stack
+// traces by 50%.
+
+namespace gwp_asan {
+namespace compression {
+
+// For the stack trace in `Unpacked` with length `UnpackedSize`, pack it into
+// the buffer `Packed` maximum length `PackedMaxSize`. The return value is the
+// number of bytes that were written to the output buffer.
+size_t pack(const uintptr_t *Unpacked, size_t UnpackedSize, uint8_t *Packed,
+            size_t PackedMaxSize);
+
+// From the packed stack trace in `Packed` of length `PackedSize`, write the
+// unpacked stack trace of maximum length `UnpackedMaxSize` into `Unpacked`.
+// Returns the number of full entries unpacked, or zero on error.
+size_t unpack(const uint8_t *Packed, size_t PackedSize, uintptr_t *Unpacked,
+              size_t UnpackedMaxSize);
+
+} // namespace compression
+} // namespace gwp_asan
+
+#endif // GWP_ASAN_STACK_TRACE_COMPRESSOR_

compiler-rt/trunk/lib/gwp_asan/stack_trace_compressor.cpp

+//===-- stack_trace_compressor.cpp ------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "gwp_asan/stack_trace_compressor.h"
+
+namespace gwp_asan {
+namespace compression {
+namespace {
+// Encodes `Value` as a variable-length integer to `Out`. Returns zero if there
+// was not enough space in the output buffer to write the complete varInt.
+// Otherwise returns the length of the encoded integer.
+size_t varIntEncode(uintptr_t Value, uint8_t *Out, size_t OutLen) {
+  for (size_t i = 0; i < OutLen; ++i) {
+    Out[i] = Value & 0x7f;
+    Value >>= 7;
+    if (!Value)
+      return i + 1;
+
+    Out[i] |= 0x80;
+  }
+
+  return 0;
+}
+
+// Decodes a variable-length integer to `Out`. Returns zero if the integer was
+// too large to be represented in a uintptr_t, or if the input buffer finished
+// before the integer was decoded (either case meaning that the `In` does not
+// point to a valid varInt buffer). Otherwise, returns the number of bytes that
+// were used to store the decoded integer.
+size_t varIntDecode(const uint8_t *In, size_t InLen, uintptr_t *Out) {
+  *Out = 0;
+  uint8_t Shift = 0;
+
+  for (size_t i = 0; i < InLen; ++i) {
+    *Out |= (static_cast<uintptr_t>(In[i]) & 0x7f) << Shift;
+
+    if (In[i] < 0x80)
+      return i + 1;
+
+    Shift += 7;
+
+    // Disallow overflowing the range of the output integer.
+    if (Shift >= sizeof(uintptr_t) * 8)
+      return 0;
+  }
+  return 0;
+}
+
+uintptr_t zigzagEncode(uintptr_t Value) {
+  uintptr_t Encoded = Value << 1;
+  if (static_cast<intptr_t>(Value) >= 0)
+    return Encoded;
+  return ~Encoded;
+}
+
+uintptr_t zigzagDecode(uintptr_t Value) {
+  uintptr_t Decoded = Value >> 1;
+  if (!(Value & 1))
+    return Decoded;
+  return ~Decoded;
+}
+} // anonymous namespace
+
+size_t pack(const uintptr_t *Unpacked, size_t UnpackedSize, uint8_t *Packed,
+            size_t PackedMaxSize) {
+  size_t Index = 0;
+  for (size_t CurrentDepth = 0; CurrentDepth < UnpackedSize; CurrentDepth++) {
+    uintptr_t Diff = Unpacked[CurrentDepth];
+    if (CurrentDepth > 0)
+      Diff -= Unpacked[CurrentDepth - 1];
+    size_t EncodedLength =
+        varIntEncode(zigzagEncode(Diff), Packed + Index, PackedMaxSize - Index);
+    if (!EncodedLength)
+      break;
+
+    Index += EncodedLength;
+  }
+
+  return Index;
+}
+
+size_t unpack(const uint8_t *Packed, size_t PackedSize, uintptr_t *Unpacked,
+              size_t UnpackedMaxSize) {
+  size_t CurrentDepth;
+  size_t Index = 0;
+  for (CurrentDepth = 0; CurrentDepth < UnpackedMaxSize; CurrentDepth++) {
+    uintptr_t EncodedDiff;
+    size_t DecodedLength =
+        varIntDecode(Packed + Index, PackedSize - Index, &EncodedDiff);
+    if (!DecodedLength)
+      break;
+    Index += DecodedLength;
+
+    Unpacked[CurrentDepth] = zigzagDecode(EncodedDiff);
+    if (CurrentDepth > 0)
+      Unpacked[CurrentDepth] += Unpacked[CurrentDepth - 1];
+  }
+
+  if (Index != PackedSize && CurrentDepth != UnpackedMaxSize)
+    return 0;
+
+  return CurrentDepth;
+}
+
+} // namespace compression
+} // namespace gwp_asan

compiler-rt/trunk/lib/gwp_asan/stack_trace_compressor_fuzzer.cpp

+#include <cstddef>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <vector>
+
+#include "gwp_asan/stack_trace_compressor.h"
+
+constexpr size_t kBytesForLargestVarInt = (sizeof(uintptr_t) * 8) / 7 + 1;
+
+extern "C" int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size) {
+  size_t BufferSize = kBytesForLargestVarInt * Size / sizeof(uintptr_t);
+  std::vector<uint8_t> Buffer(BufferSize);
+  std::vector<uint8_t> Buffer2(BufferSize);
+
+  // Unpack the fuzz bytes.
+  gwp_asan::compression::unpack(Data, Size,
+                                reinterpret_cast<uintptr_t *>(Buffer2.data()),
+                                BufferSize / sizeof(uintptr_t));
+
+  // Pack the fuzz bytes.
+  size_t BytesWritten = gwp_asan::compression::pack(
+      reinterpret_cast<const uintptr_t *>(Data), Size / sizeof(uintptr_t),
+      Buffer.data(), BufferSize);
+
+  // Unpack the compressed buffer.
+  size_t DecodedElements = gwp_asan::compression::unpack(
+      Buffer.data(), BytesWritten,
+      reinterpret_cast<uintptr_t *>(Buffer2.data()),
+      BufferSize / sizeof(uintptr_t));
+
+  // Ensure that every element was encoded and decoded properly.
+  if (DecodedElements != Size / sizeof(uintptr_t))
+    abort();
+
+  // Ensure that the compression and uncompression resulted in the same trace.
+  const uintptr_t *FuzzPtrs = reinterpret_cast<const uintptr_t *>(Data);
+  const uintptr_t *DecodedPtrs =
+      reinterpret_cast<const uintptr_t *>(Buffer2.data());
+  for (size_t i = 0; i < Size / sizeof(uintptr_t); ++i) {
+    if (FuzzPtrs[i] != DecodedPtrs[i]) {
+      fprintf(stderr, "FuzzPtrs[%zu] != DecodedPtrs[%zu] (0x%zx vs. 0x%zx)", i,
+              i, FuzzPtrs[i], DecodedPtrs[i]);
+      abort();
+    }
+  }
+
+  return 0;
+}

compiler-rt/trunk/lib/gwp_asan/tests/compression.cpp

+//===-- compression.cpp -----------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "gwp_asan/stack_trace_compressor.h"
+#include "gtest/gtest.h"
+
+namespace gwp_asan {
+namespace compression {
+
+TEST(GwpAsanCompressionTest, SingleByteVarInt) {
+  uint8_t Compressed[1];
+
+  uintptr_t Uncompressed = 0x00;
+  EXPECT_EQ(1u, pack(&Uncompressed, 1u, Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[0], 0x00);
+
+  Uncompressed = 0x01;
+  EXPECT_EQ(1u, pack(&Uncompressed, 1u, Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[0], 0x02); // +1 => 2 in zigzag.
+
+  Uncompressed = 0x3f;
+  EXPECT_EQ(1u, pack(&Uncompressed, 1u, Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[0], 0x7e); // +63 => 127 in zigzag.
+}
+
+TEST(GwpAsanCompressionTest, MultiByteVarInt) {
+  uint8_t Compressed[sizeof(uintptr_t) + 1];
+
+  uintptr_t Uncompressed = 0x40;
+  EXPECT_EQ(2u, pack(&Uncompressed, 1u, Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[0], 0x80); // +64 => 128 in zigzag.
+  EXPECT_EQ(Compressed[1], 0x01);
+
+  Uncompressed = 0x41;
+  EXPECT_EQ(2u, pack(&Uncompressed, 1u, Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[0], 0x82); // +65 => 130 in zigzag
+  EXPECT_EQ(Compressed[1], 0x01);
+
+  Uncompressed = 0x1fff;
+  EXPECT_EQ(2u, pack(&Uncompressed, 1u, Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[0], 0xfe); // +8191 => 16382 in zigzag
+  EXPECT_EQ(Compressed[1], 0x7f);
+
+  Uncompressed = 0x2000;
+  EXPECT_EQ(3u, pack(&Uncompressed, 1u, Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[0], 0x80); // +8192 => 16384 in zigzag
+  EXPECT_EQ(Compressed[1], 0x80);
+  EXPECT_EQ(Compressed[2], 0x01);
+
+  Uncompressed = 0xff010ff0;
+  EXPECT_EQ(5u, pack(&Uncompressed, 1u, Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[0], 0xe0); // +0xff010ff0 => 0x1FE021FE0 in zigzag
+  EXPECT_EQ(Compressed[1], 0xbf);
+  EXPECT_EQ(Compressed[2], 0x88);
+  EXPECT_EQ(Compressed[3], 0xf0);
+  EXPECT_EQ(Compressed[4], 0x1f);
+}
+
+TEST(GwpAsanCompressionTest, CorrectDifference) {
+  uint8_t Compressed[10];
+  uintptr_t Uncompressed[2] = {0x00, 0x00};
+
+  EXPECT_EQ(2u, pack(Uncompressed, sizeof(Uncompressed) / sizeof(uintptr_t),
+                     Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[1], 0x00); // +0 difference => 0 in zigzag.
+
+  Uncompressed[1] = 0x01;
+  EXPECT_EQ(2u, pack(Uncompressed, sizeof(Uncompressed) / sizeof(uintptr_t),
+                     Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[1], 0x02); // +1 difference => 2 in zigzag.
+
+  Uncompressed[1] = 0x02;
+  EXPECT_EQ(2u, pack(Uncompressed, sizeof(Uncompressed) / sizeof(uintptr_t),
+                     Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[1], 0x04); // +2 difference => 4 in zigzag.
+
+  Uncompressed[1] = 0x80;
+  EXPECT_EQ(3u, pack(Uncompressed, sizeof(Uncompressed) / sizeof(uintptr_t),
+                     Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[1], 0x80); // +128 difference => +256 in zigzag (note the
+  EXPECT_EQ(Compressed[2], 0x02); // varint encoding here).
+
+  Uncompressed[0] = 0x01;
+  Uncompressed[1] = 0x00;
+  EXPECT_EQ(2u, pack(Uncompressed, sizeof(Uncompressed) / sizeof(uintptr_t),
+                     Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[1], 0x01); // -1 difference => +1 in zigzag.
+
+  Uncompressed[0] = 0x02;
+  EXPECT_EQ(2u, pack(Uncompressed, sizeof(Uncompressed) / sizeof(uintptr_t),
+                     Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[1], 0x03); // -2 difference => +3 in zigzag.
+
+  Uncompressed[0] = 0x80;
+  EXPECT_EQ(4u, pack(Uncompressed, sizeof(Uncompressed) / sizeof(uintptr_t),
+                     Compressed, sizeof(Compressed)));
+  EXPECT_EQ(Compressed[2], 0xff); // -128 difference => +255 in zigzag (note the
+  EXPECT_EQ(Compressed[3], 0x01); // varint encoding here).
+}
+
+// Space needed to encode the biggest uintptr_t as a varint is ceil((8 / 7) *
+// sizeof(uintptr_t)), as each 7 bits requires 8 bits of space.
+constexpr size_t kBytesForLargestVarInt = (sizeof(uintptr_t) * 8) / 7 + 1;
+
+// Ensures that when the closest diff between two pointers is via. underflow,
+// we take the underflow option.
+TEST(GwpAsanCompressionTest, ClosestDiffIsUnderflow) {
+  uint8_t Compressed[2];
+  uintptr_t Uncompressed[2] = {0x00, UINTPTR_MAX};
+
+  EXPECT_EQ(2u, pack(Uncompressed, sizeof(Uncompressed) / sizeof(uintptr_t),
+                     Compressed, sizeof(Compressed)));
+  // -1 difference => +1 in zigzag.
+  EXPECT_EQ(Compressed[1], 0x01);
+}
+
+// Ensures that when the closest diff between two pointers is via. overflow,
+// that we take this option.
+TEST(GwpAsanCompressionTest, ClosestDiffIsOverflow) {
+  uint8_t Compressed[2];
+  uintptr_t Uncompressed[2] = {UINTPTR_MAX, 0x00};
+
+  // Note here that the first element is encoded as the difference from zero.
+  EXPECT_EQ(2u, pack(Uncompressed, sizeof(Uncompressed) / sizeof(uintptr_t),
+                     Compressed, sizeof(Compressed)));
+  // -1 difference => +1 in zigzag (the first pointer is encoded as -1).
+  EXPECT_EQ(Compressed[0], 0x01);
+  // +1 difference => +2 in zigzag.
+  EXPECT_EQ(Compressed[1], 0x02);
+}
+
+void runPackUnpack(uintptr_t *Test, size_t NumEntries) {
+  // Setup the input/output buffers based on the maximum possible size.
+  uintptr_t *Uncompressed =
+      static_cast<uintptr_t *>(alloca(NumEntries * sizeof(uintptr_t)));
+  size_t CompressedBufferSize = NumEntries * kBytesForLargestVarInt;
+  uint8_t *Compressed = static_cast<uint8_t *>(alloca(CompressedBufferSize));
+
+  // Pack the provided testcase, recoding the number of bytes it took for
+  // storage.
+  size_t BytesUsedForPacking =
+      pack(Test, NumEntries, Compressed, CompressedBufferSize);
+  EXPECT_NE(BytesUsedForPacking, 0u);
+
+  // Unpack the testcase and ensure that the correct number of entries was
+  // unpacked.
+  EXPECT_EQ(NumEntries,
+            unpack(Compressed, BytesUsedForPacking, Uncompressed, NumEntries));
+
+  // Ensure that the unpacked trace is the same as the original testcase.
+  for (size_t i = 0; i < NumEntries; ++i) {
+    EXPECT_EQ(Uncompressed[i], Test[i]);
+  }
+}
+
+TEST(GwpAsanCompressionTest, UncompressVarInt) {
+  uint8_t Compressed[] = {0x00, 0xaa, 0xaf, 0xd0, 0xda, 0x24};
+  uintptr_t Uncompressed[2];
+
+  EXPECT_EQ(2u, unpack(Compressed, sizeof(Compressed), Uncompressed, 2u));
+  EXPECT_EQ(Uncompressed[0], 0x00u);
+  EXPECT_EQ(Uncompressed[1], 0x125aa0bd5u);
+}
+
+TEST(GwpAsanCompressionTest, CompressUncompressAscending) {
+  uintptr_t Test[] = {1, 2, 3};
+  runPackUnpack(Test, sizeof(Test) / sizeof(uintptr_t));
+}
+
+TEST(GwpAsanCompressionTest, CompressUncompressDescending) {
+  uintptr_t Test[] = {3, 2, 1};
+  runPackUnpack(Test, sizeof(Test) / sizeof(uintptr_t));
+}
+
+TEST(GwpAsanCompressionTest, CompressUncompressRepeated) {
+  uintptr_t Test[] = {3, 3, 3};
+  runPackUnpack(Test, sizeof(Test) / sizeof(uintptr_t));
+}
+
+TEST(GwpAsanCompressionTest, CompressUncompressZigZag) {
+  uintptr_t Test[] = {1, 3, 2, 4, 1, 2};
+  runPackUnpack(Test, sizeof(Test) / sizeof(uintptr_t));
+}
+
+TEST(GwpAsanCompressionTest, CompressUncompressVarInt) {
+  uintptr_t Test[] = {0x1981561, 0x18560, 0x125ab9135, 0x1232562};
+  runPackUnpack(Test, sizeof(Test) / sizeof(uintptr_t));
+}
+
+TEST(GwpAsanCompressionTest, CompressUncompressLargestDifference) {
+  uintptr_t Test[] = {0x00, INTPTR_MAX, UINTPTR_MAX, INTPTR_MAX, 0x00};
+  runPackUnpack(Test, sizeof(Test) / sizeof(uintptr_t));
+}
+
+TEST(GwpAsanCompressionTest, CompressUncompressBigPointers) {
+  uintptr_t Test[] = {UINTPTR_MAX, UINTPTR_MAX - 10};
+  runPackUnpack(Test, sizeof(Test) / sizeof(uintptr_t));
+
+  uintptr_t Test2[] = {UINTPTR_MAX - 10, UINTPTR_MAX};
+  runPackUnpack(Test2, sizeof(Test2) / sizeof(uintptr_t));
+}
+
+TEST(GwpAsanCompressionTest, UncompressFailsWithOutOfBoundsVarInt) {
+  uint8_t Compressed[kBytesForLargestVarInt + 1];
+  for (size_t i = 0; i < kBytesForLargestVarInt; ++i) {
+    Compressed[i] = 0x80;
+  }
+  Compressed[kBytesForLargestVarInt] = 0x00;
+
+  uintptr_t Uncompressed;
+  EXPECT_EQ(unpack(Compressed, kBytesForLargestVarInt + 1, &Uncompressed, 1),
+            0u);
+}
+
+TEST(GwpAsanCompressionTest, UncompressFailsWithTooSmallBuffer) {
+  uint8_t Compressed[] = {0x80, 0x00};
+
+  uintptr_t Uncompressed;
+  EXPECT_EQ(unpack(Compressed, 1u, &Uncompressed, 1), 0u);
+}
+
+TEST(GwpAsanCompressionTest, CompressPartiallySucceedsWithTooSmallBuffer) {
+  uintptr_t Uncompressed[] = {
+      0x80,  // Requires 2 bytes for varint.
+      0x100, // Requires two bytes for varint difference of 0x80.
+      0xff,  // Requires single byte for varint difference of -0x01
+  };
+  uint8_t Compressed[3 * kBytesForLargestVarInt];
+
+  // Zero and one byte buffers shouldn't encode anything (see above for size
+  // requirements).
+  EXPECT_EQ(pack(Uncompressed, 3u, Compressed, 0u), 0u);
+  EXPECT_EQ(pack(Uncompressed, 3u, Compressed, 1u), 0u);
+
+  // Two byte buffer should hold a single varint-encoded value.
+  EXPECT_EQ(pack(Uncompressed, 3u, Compressed, 2u), 2u);
+
+  // Three bytes isn't enough to cover the first two pointers, as both take two
+  // bytes each to store. Expect a single value to be compressed.
+  EXPECT_EQ(pack(Uncompressed, 3u, Compressed, 3u), 2u);
+
+  // Four bytes is enough for the first two pointers to be stored.
+  EXPECT_EQ(pack(Uncompressed, 3u, Compressed, 4u), 4u);
+
+  // And five is enough for all three pointers to be stored.
+  EXPECT_EQ(pack(Uncompressed, 3u, Compressed, 5u), 5u);
+  // And a buffer that's bigger than five bytes should still only write five
+  // bytes.
+  EXPECT_EQ(pack(Uncompressed, 3u, Compressed, 6u), 5u);
+  EXPECT_EQ(pack(Uncompressed, 3u, Compressed, 3 * kBytesForLargestVarInt), 5u);
+}
+} // namespace compression
+} // namespace gwp_asan