1 files changed, 22 insertions, 1 deletions
diff --git a/sandbox/src/interception.cc b/sandbox/src/interception.cc
index 198cdbf..8ee56f4 100644
--- a/sandbox/src/interception.cc
+++ b/sandbox/src/interception.cc
@@ -30,6 +30,9 @@ namespace sandbox {
 
 SANDBOX_INTERCEPT SharedMemory* g_interceptions;
 
+// Magic constant that identifies that this function is not to be patched.
+const char kUnloadDLLDummyFunction[] = "@";
+
 InterceptionManager::InterceptionManager(TargetProcess* child_process,
                                          bool relaxed)
     : child_(child_process), names_used_(false), relaxed_(relaxed) {
@@ -49,7 +52,6 @@ bool InterceptionManager::AddToPatchedFunctions(
   function.interceptor_address = replacement_code_address;
 
   interceptions_.push_back(function);
-
   return true;
 }
 
@@ -65,7 +67,19 @@ bool InterceptionManager::AddToPatchedFunctions(
 
   interceptions_.push_back(function);
   names_used_ = true;
+  return true;
+}
 
+bool InterceptionManager::AddToUnloadModules(const wchar_t* dll_name) {
+  InterceptionData module_to_unload;
+  module_to_unload.type = INTERCEPTION_UNLOAD_MODULE;
+  module_to_unload.dll = dll_name;
+  // The next two are dummy values that make the structures regular, instead
+  // of having special cases. They should not be used.
+  module_to_unload.function = kUnloadDLLDummyFunction;
+  module_to_unload.interceptor_address = reinterpret_cast<void*>(1);
+
+  interceptions_.push_back(module_to_unload);
   return true;
 }
 
@@ -204,6 +218,7 @@ bool InterceptionManager::SetupDllInfo(const InterceptionData& data,
   *buffer = reinterpret_cast<char*>(*buffer) + required;
 
   // set up the dll info to be what we know about it at this time
+  dll_info->unload_module = (data.type == INTERCEPTION_UNLOAD_MODULE);
   dll_info->record_bytes = required;
   dll_info->offset_to_functions = required;
   dll_info->num_functions = 0;
@@ -221,6 +236,12 @@ bool InterceptionManager::SetupInterceptionInfo(const InterceptionData& data,
   DCHECK(buffer);
   DCHECK(*buffer);
 
+  if ((dll_info->unload_module) && 
+      (data.function != kUnloadDLLDummyFunction)) {
+    // Can't specify a dll for both patch and unload.
+    NOTREACHED();
+  }
+
   FunctionInfo* function = reinterpret_cast<FunctionInfo*>(*buffer);
 
   size_t name_bytes = data.function.size();