summaryrefslogtreecommitdiff
path: root/lib/interception/tests/interception_win_test.cc
blob: 02c4dd9e595a815ac81ddebba23e4e78dc394f6a (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
//===-- interception_win_test.cc ------------------------------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file is a part of ThreadSanitizer/AddressSanitizer runtime.
// Tests for interception_win.h.
//
//===----------------------------------------------------------------------===//
#include "interception/interception.h"

#include "gtest/gtest.h"

// Too slow for debug build
#if !SANITIZER_DEBUG
#if SANITIZER_WINDOWS

#define WIN32_LEAN_AND_MEAN
#include <windows.h>

namespace {

typedef int (*IdentityFunction)(int);

#if !SANITIZER_WINDOWS64
u8 kIdentityCodeWithPrologue[] = {
    0x55,              // push        ebp
    0x8B, 0xEC,        // mov         ebp,esp
    0x8B, 0x45, 0x08,  // mov         eax,dword ptr [ebp + 8]
    0x5D,              // pop         ebp
    0xC3,              // ret
};

u8 kIdentityCodeWithPushPop[] = {
    0x55,              // push        ebp
    0x8B, 0xEC,        // mov         ebp,esp
    0x53,              // push        ebx
    0x50,              // push        eax
    0x58,              // pop         eax
    0x8B, 0x45, 0x08,  // mov         eax,dword ptr [ebp + 8]
    0x5B,              // pop         ebx
    0x5D,              // pop         ebp
    0xC3,              // ret
};

#endif

// A buffer holding the dynamically generated code under test.
u8* ActiveCode;
size_t ActiveCodeLength = 4096;

bool LoadActiveCode(u8* Code, size_t CodeLength, uptr* EntryPoint) {
  if (ActiveCode == nullptr) {
    ActiveCode =
        (u8*)::VirtualAlloc(nullptr, ActiveCodeLength, MEM_COMMIT | MEM_RESERVE,
                            PAGE_EXECUTE_READWRITE);
    if (ActiveCode == nullptr) return false;
  }

  size_t Position = 0;
  *EntryPoint = (uptr)&ActiveCode[0];

  // Copy the function body.
  for (size_t i = 0; i < CodeLength; ++i)
  	ActiveCode[Position++] = Code[i];

  return true;
}

int InterceptorFunctionCalled;

NOINLINE int InterceptorFunction(int x) {
  ++InterceptorFunctionCalled;
  return x;
}

}  // namespace

namespace __interception {

// Tests for interception_win.h
TEST(Interception, InternalGetProcAddress) {
  HMODULE ntdll_handle = ::GetModuleHandle("ntdll");
  ASSERT_NE(nullptr, ntdll_handle);
  uptr DbgPrint_expected = (uptr)::GetProcAddress(ntdll_handle, "DbgPrint");
  uptr isdigit_expected = (uptr)::GetProcAddress(ntdll_handle, "isdigit");
  uptr DbgPrint_adddress = InternalGetProcAddress(ntdll_handle, "DbgPrint");
  uptr isdigit_address = InternalGetProcAddress(ntdll_handle, "isdigit");

  EXPECT_EQ(DbgPrint_expected, DbgPrint_adddress);
  EXPECT_EQ(isdigit_expected, isdigit_address);
  EXPECT_NE(DbgPrint_adddress, isdigit_address);
}

void TestIdentityFunctionPatching(u8* IdentityCode, size_t IdentityCodeLength) {
  uptr IdentityAddress;
  ASSERT_TRUE(
      LoadActiveCode(IdentityCode, IdentityCodeLength, &IdentityAddress));
  IdentityFunction Identity = (IdentityFunction)IdentityAddress;

  // Validate behavior before dynamic patching.
  InterceptorFunctionCalled = 0;
  EXPECT_EQ(0, Identity(0));
  EXPECT_EQ(42, Identity(42));
  EXPECT_EQ(0, InterceptorFunctionCalled);

  // Patch the function.
  uptr RealIdentityAddress = 0;
  EXPECT_TRUE(OverrideFunction(IdentityAddress, (uptr)&InterceptorFunction,
                               &RealIdentityAddress));
  IdentityFunction RealIdentity = (IdentityFunction)RealIdentityAddress;

  // Calling the redirected function.
  InterceptorFunctionCalled = 0;
  EXPECT_EQ(0, Identity(0));
  EXPECT_EQ(42, Identity(42));
  EXPECT_EQ(2, InterceptorFunctionCalled);

  // Calling the real function.
  InterceptorFunctionCalled = 0;
  EXPECT_EQ(0, RealIdentity(0));
  EXPECT_EQ(42, RealIdentity(42));
  EXPECT_EQ(0, InterceptorFunctionCalled);
}

#if !SANITIZER_WINDOWS64
TEST(Interception, OverrideFunction) {
  TestIdentityFunctionPatching(kIdentityCodeWithPrologue,
                               sizeof(kIdentityCodeWithPrologue));
  TestIdentityFunctionPatching(kIdentityCodeWithPushPop,
                               sizeof(kIdentityCodeWithPushPop));
}
#endif

}  // namespace __interception

#endif  // SANITIZER_WINDOWS
#endif  // #if !SANITIZER_DEBUG