1 : // Copyright 2013 Google Inc. All Rights Reserved.
2 : //
3 : // Licensed under the Apache License, Version 2.0 (the "License");
4 : // you may not use this file except in compliance with the License.
5 : // You may obtain a copy of the License at
6 : //
7 : // http://www.apache.org/licenses/LICENSE-2.0
8 : //
9 : // Unless required by applicable law or agreed to in writing, software
10 : // distributed under the License is distributed on an "AS IS" BASIS,
11 : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 : // See the License for the specific language governing permissions and
13 : // limitations under the License.
14 :
15 : // ntstatus.h conflicts with windows.h unless this is defined at inclusion.
16 : #define WIN32_NO_STATUS
17 : #include "syzygy/agent/common/dll_notifications.h"
18 :
19 : #include <windows.h>
20 : #include <winternl.h> // For UNICODE_STRING.
21 : #undef WIN32_NO_STATUS
22 : #include <ntstatus.h> // For STATUS_SUCCESS.
23 :
24 : #include "base/logging.h"
25 :
26 :
27 : // These structures and functions are documented in MSDN, see
28 : // http://msdn.microsoft.com/en-us/library/gg547638(v=vs.85).aspx
29 : // there are however no headers or import libraries available in the
30 : // Platform SDK.
31 : enum {
32 : // The DLL was loaded. The NotificationData parameter points to an
33 : // LDR_DLL_LOADED_NOTIFICATION_DATA structure.
34 : LDR_DLL_NOTIFICATION_REASON_LOADED = 1,
35 : // The DLL was unloaded. The NotificationData parameter points to an
36 : // LDR_DLL_UNLOADED_NOTIFICATION_DATA structure.
37 : LDR_DLL_NOTIFICATION_REASON_UNLOADED = 2,
38 : };
39 :
40 : typedef struct _LDR_DLL_LOADED_NOTIFICATION_DATA {
41 : // Reserved.
42 : ULONG Flags;
43 : // The full path name of the DLL module.
44 : PCUNICODE_STRING FullDllName;
45 : // The base file name of the DLL module.
46 : PCUNICODE_STRING BaseDllName;
47 : // A pointer to the base address for the DLL in memory.
48 : PVOID DllBase;
49 : // The size of the DLL image, in bytes.
50 : ULONG SizeOfImage;
51 : } LDR_DLL_LOADED_NOTIFICATION_DATA, *PLDR_DLL_LOADED_NOTIFICATION_DATA;
52 :
53 : typedef struct _LDR_DLL_UNLOADED_NOTIFICATION_DATA {
54 : // Reserved.
55 : ULONG Flags;
56 : // The full path name of the DLL module.
57 : PCUNICODE_STRING FullDllName;
58 : // The base file name of the DLL module.
59 : PCUNICODE_STRING BaseDllName;
60 : // A pointer to the base address for the DLL in memory.
61 : PVOID DllBase;
62 : // The size of the DLL image, in bytes.
63 : ULONG SizeOfImage;
64 : } LDR_DLL_UNLOADED_NOTIFICATION_DATA, *PLDR_DLL_UNLOADED_NOTIFICATION_DATA;
65 :
66 : typedef union _LDR_DLL_NOTIFICATION_DATA {
67 : LDR_DLL_LOADED_NOTIFICATION_DATA Loaded;
68 : LDR_DLL_UNLOADED_NOTIFICATION_DATA Unloaded;
69 : } LDR_DLL_NOTIFICATION_DATA, *PLDR_DLL_NOTIFICATION_DATA;
70 :
71 : typedef VOID (CALLBACK *PLDR_DLL_NOTIFICATION_FUNCTION)(
72 : ULONG notification_reason,
73 : const LDR_DLL_NOTIFICATION_DATA* notification_data,
74 : PVOID context);
75 :
76 : typedef NTSTATUS (NTAPI *LdrRegisterDllNotificationFunc)(
77 : ULONG flags, PLDR_DLL_NOTIFICATION_FUNCTION notification_function,
78 : PVOID context, PVOID *cookie);
79 : typedef NTSTATUS (NTAPI *LdrUnregisterDllNotificationFunc)(PVOID cookie);
80 :
81 : namespace agent {
82 : namespace common {
83 :
84 : namespace {
85 :
86 E : HMODULE GetNtDll() {
87 E : HMODULE ntdll = ::GetModuleHandle(L"ntdll.dll");
88 E : CHECK_NE(static_cast<HMODULE>(NULL), ntdll);
89 E : return ntdll;
90 E : }
91 :
92 : bool Register(PLDR_DLL_NOTIFICATION_FUNCTION notify_fn,
93 : void* context,
94 E : void** cookie) {
95 : LdrRegisterDllNotificationFunc reg_fn =
96 : reinterpret_cast<LdrRegisterDllNotificationFunc>(
97 E : ::GetProcAddress(GetNtDll(), "LdrRegisterDllNotification"));
98 :
99 E : if (reg_fn == NULL)
100 i : return false;
101 :
102 E : NTSTATUS status = reg_fn(0, notify_fn, context, cookie);
103 E : return status == STATUS_SUCCESS;
104 E : }
105 :
106 E : bool Unregister(void* cookie) {
107 : LdrUnregisterDllNotificationFunc unreg_fn =
108 : reinterpret_cast<LdrUnregisterDllNotificationFunc>(
109 E : ::GetProcAddress(GetNtDll(), "LdrUnregisterDllNotification"));
110 :
111 E : if (unreg_fn == NULL)
112 i : return false;
113 :
114 E : NTSTATUS status = unreg_fn(cookie);
115 E : return status == STATUS_SUCCESS;
116 E : }
117 :
118 E : base::StringPiece16 ToStringPiece(const UNICODE_STRING* str) {
119 E : CHECK_NE(static_cast<const UNICODE_STRING*>(NULL), str);
120 E : return base::StringPiece16(str->Buffer, str->Length / sizeof(wchar_t));
121 E : }
122 :
123 : } // namespace
124 :
125 E : DllNotificationWatcher::DllNotificationWatcher() : cookie_(NULL) {
126 E : }
127 :
128 E : DllNotificationWatcher::~DllNotificationWatcher() {
129 E : Reset();
130 E : }
131 :
132 E : bool DllNotificationWatcher::Init(const CallbackType& callback) {
133 E : CHECK_EQ(static_cast<void*>(NULL), cookie_);
134 :
135 E : callback_ = callback;
136 E : if (!Register(NotificationFunction, this, &cookie_)) {
137 i : DCHECK_EQ(static_cast<void*>(NULL), cookie_);
138 i : callback_.Reset();
139 i : return false;
140 : }
141 :
142 E : return true;
143 E : }
144 :
145 E : void DllNotificationWatcher::Reset() {
146 E : if (cookie_ == NULL)
147 E : return;
148 :
149 E : CHECK(Unregister(cookie_));
150 E : cookie_ = NULL;
151 E : callback_.Reset();
152 E : }
153 :
154 : void CALLBACK DllNotificationWatcher::NotificationFunction(
155 E : ULONG reason, const LDR_DLL_NOTIFICATION_DATA* data, void* context) {
156 E : CHECK_NE(static_cast<const LDR_DLL_NOTIFICATION_DATA*>(NULL), data);
157 E : CHECK_NE(static_cast<void*>(NULL), context);
158 :
159 : DllNotificationWatcher* self =
160 E : reinterpret_cast<DllNotificationWatcher*>(context);
161 E : EventType event_type = DllLoaded;
162 E : HMODULE module = NULL;
163 E : size_t module_size = 0;
164 E : base::StringPiece16 dll_path;
165 E : base::StringPiece16 dll_base_name;
166 :
167 E : switch (reason) {
168 : case LDR_DLL_NOTIFICATION_REASON_LOADED:
169 E : event_type = DllLoaded;
170 E : module = reinterpret_cast<HMODULE>(data->Loaded.DllBase);
171 E : module_size = data->Loaded.SizeOfImage;
172 E : dll_path = ToStringPiece(data->Loaded.FullDllName);
173 E : dll_base_name = ToStringPiece(data->Loaded.BaseDllName);
174 E : break;
175 :
176 : case LDR_DLL_NOTIFICATION_REASON_UNLOADED:
177 E : event_type = DllUnloaded;
178 E : module = reinterpret_cast<HMODULE>(data->Unloaded.DllBase);
179 E : module_size = data->Unloaded.SizeOfImage;
180 E : dll_path = ToStringPiece(data->Unloaded.FullDllName);
181 E : dll_base_name = ToStringPiece(data->Unloaded.BaseDllName);
182 E : break;
183 :
184 : default:
185 i : return;
186 : }
187 :
188 E : self->callback_.Run(event_type, module, module_size, dll_path, dll_base_name);
189 E : }
190 :
191 : } // namespace common
192 : } // namespace agent
|