mozglue/tests/TestStackCookie.cpp

   1 /* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
   2 /* vim: set ts=8 sts=2 et sw=2 tw=80: */
   3 /* This Source Code Form is subject to the terms of the Mozilla Public
   4  * License, v. 2.0. If a copy of the MPL was not distributed with this
   5  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
   6
   7 #include <stdio.h>
   8 #include <windows.h>
   9
  10 #include <array>
  11 #include <utility>
  12
  13 #include "mozilla/WindowsStackCookie.h"
  14
  15 #if defined(DEBUG) && defined(_M_X64) && !defined(__MINGW64__)
  16
  17 uint64_t NoStackCookie(const uint64_t* aArray, size_t aSize) {
  18   uint64_t result = 0;
  19   for (size_t i = 0; i < aSize; ++i) {
  20     result += aArray[i];
  21   }
  22   result /= aSize;
  23   return result;
  24 }
  25
  26 // We expect the following instructions to be generated:
  27 //
  28 // 48 8b 05 XX XX XX XX  mov     rax,qword ptr [__security_cookie]
  29 // 48 31 e0              xor     rax,rsp 48
  30 // 89 44 24 38           mov     qword ptr [rsp+38h],rax
  31 uint64_t StackCookieWithSmallStackSpace(const uint64_t* aArray, size_t aSize) {
  32   uint64_t array[0x2]{};
  33   for (size_t i = 0; i < aSize; ++i) {
  34     array[aArray[i]] += aArray[i];
  35   }
  36   return array[0] + array[1];
  37 }
  38
  39 // We expect the following instructions to be generated:
  40 //
  41 // 48 8b 05 XX XX XX XX     mov     rax,qword ptr [__security_cookie]
  42 // 48 31 e0                 xor     rax,rsp
  43 // 48 89 84 24 28 40 00 00  mov     qword ptr [rsp+4028h],rax
  44 uint64_t StackCookieWithLargeStackSpace(const uint64_t* aArray, size_t aSize) {
  45   uint64_t array[0x800]{};
  46   for (size_t i = 0; i < aSize; ++i) {
  47     array[aArray[i]] += aArray[i];
  48   }
  49   return array[0] + array[0x7FF];
  50 }
  51
  52 bool TestStackCookieCheck() {
  53   std::array<std::pair<uintptr_t, bool>, 3> testCases{
  54       std::make_pair<uintptr_t, bool>(
  55           reinterpret_cast<uintptr_t>(NoStackCookie), false),
  56       std::make_pair<uintptr_t, bool>(
  57           reinterpret_cast<uintptr_t>(StackCookieWithSmallStackSpace), true),
  58       std::make_pair<uintptr_t, bool>(
  59           reinterpret_cast<uintptr_t>(StackCookieWithLargeStackSpace), true),
  60   };
  61   for (auto [functionAddress, expectStackCookieCheck] : testCases) {
  62     if (mozilla::HasStackCookieCheck(functionAddress) !=
  63         expectStackCookieCheck) {
  64       printf(
  65           "TEST-FAILED | StackCookie | Wrong output from HasStackCookieCheck "
  66           "for function at %p (expected %d).\n",
  67           reinterpret_cast<void*>(functionAddress), expectStackCookieCheck);
  68       return false;
  69     }
  70     printf(
  71         "TEST-PASS | StackCookie | Correct output from HasStackCookieCheck for "
  72         "function at %p (expected %d).\n",
  73         reinterpret_cast<void*>(functionAddress), expectStackCookieCheck);
  74   }
  75   return true;
  76 }
  77
  78 #endif  // defined(DEBUG) && defined(_M_X64) && !defined(__MINGW64__)
  79
  80 int wmain(int argc, wchar_t* argv[]) {
  81 #if defined(DEBUG) && defined(_M_X64) && !defined(__MINGW64__)
  82   if (!TestStackCookieCheck()) {
  83     return 1;
  84   }
  85 #endif  // defined(DEBUG) && defined(_M_X64) && !defined(__MINGW64__)
  86
  87   printf("TEST-PASS | StackCookie | All tests ran successfully\n");
  88   return 0;
  89 }