3 files changed, 59 insertions, 0 deletions
diff --git a/sandbox/win/src/service_resolver.h b/sandbox/win/src/service_resolver.h
index 0089692..20486c9 100644
--- a/sandbox/win/src/service_resolver.h
+++ b/sandbox/win/src/service_resolver.h
@@ -46,6 +46,15 @@ class ServiceResolverThunk : public ResolverThunk {
   // Call this to set up ntdll_base_ which will allow for local patches.
   virtual void AllowLocalPatches();
 
+  // Verifies that the function specified by |target_name| in |target_module| is
+  // a service and copies the data from that function into |thunk_storage|. If
+  // |storage_bytes| is too small, then the method fails.
+  virtual NTSTATUS CopyThunk(const void* target_module,
+                             const char* target_name,
+                             BYTE* thunk_storage,
+                             size_t storage_bytes,
+                             size_t* storage_used);
+
  protected:
   // The unit test will use this member to allow local patch on a buffer.
   HMODULE ntdll_base_;
diff --git a/sandbox/win/src/service_resolver_32.cc b/sandbox/win/src/service_resolver_32.cc
index 2e69dbc..9b8bbf9 100644
--- a/sandbox/win/src/service_resolver_32.cc
+++ b/sandbox/win/src/service_resolver_32.cc
@@ -179,6 +179,32 @@ size_t ServiceResolverThunk::GetThunkSize() const {
   return offsetof(ServiceFullThunk, internal_thunk) + GetInternalThunkSize();
 }
 
+NTSTATUS ServiceResolverThunk::CopyThunk(const void* target_module,
+                                         const char* target_name,
+                                         BYTE* thunk_storage,
+                                         size_t storage_bytes,
+                                         size_t* storage_used) {
+  NTSTATUS ret = ResolveTarget(target_module, target_name, &target_);
+  if (!NT_SUCCESS(ret))
+    return ret;
+
+  size_t thunk_bytes = GetThunkSize();
+  if (storage_bytes < thunk_bytes)
+    return STATUS_UNSUCCESSFUL;
+
+  ServiceFullThunk* thunk = reinterpret_cast<ServiceFullThunk*>(thunk_storage);
+
+  if (!IsFunctionAService(&thunk->original) &&
+      (!relaxed_ || !SaveOriginalFunction(&thunk->original, thunk_storage))) {
+    return STATUS_UNSUCCESSFUL;
+  }
+
+  if (NULL != storage_used)
+    *storage_used = thunk_bytes;
+
+  return ret;
+}
+
 bool ServiceResolverThunk::IsFunctionAService(void* local_thunk) const {
   ServiceEntry function_code;
   SIZE_T read;
diff --git a/sandbox/win/src/service_resolver_64.cc b/sandbox/win/src/service_resolver_64.cc
index 9612418..79c6b16 100644
--- a/sandbox/win/src/service_resolver_64.cc
+++ b/sandbox/win/src/service_resolver_64.cc
@@ -116,6 +116,30 @@ size_t ServiceResolverThunk::GetThunkSize() const {
   return sizeof(ServiceFullThunk);
 }
 
+NTSTATUS ServiceResolverThunk::CopyThunk(const void* target_module,
+                                         const char* target_name,
+                                         BYTE* thunk_storage,
+                                         size_t storage_bytes,
+                                         size_t* storage_used) {
+  NTSTATUS ret = ResolveTarget(target_module, target_name, &target_);
+  if (!NT_SUCCESS(ret))
+    return ret;
+
+  size_t thunk_bytes = GetThunkSize();
+  if (storage_bytes < thunk_bytes)
+    return STATUS_UNSUCCESSFUL;
+
+  ServiceFullThunk* thunk = reinterpret_cast<ServiceFullThunk*>(thunk_storage);
+
+  if (!IsFunctionAService(&thunk->original))
+    return STATUS_UNSUCCESSFUL;
+
+  if (NULL != storage_used)
+    *storage_used = thunk_bytes;
+
+  return ret;
+}
+
 bool ServiceResolverThunk::IsFunctionAService(void* local_thunk) const {
   ServiceFullThunk function_code;
   SIZE_T read;