Added an utility function to ensure that we don't "asanize" the same image twice.
BUG=
Review URL: https://codereview.appspot.com/6817076
git-svn-id: http://sawbuck.googlecode.com/svn/trunk/syzygy@1220 15e8cca8-e42c-11de-a347-f34a4f72eb7d
diff --git a/agent/asan/asan_shadow_unittest.cc b/agent/asan/asan_shadow_unittest.cc
index 1c4ee6f..9add91f 100644
--- a/agent/asan/asan_shadow_unittest.cc
+++ b/agent/asan/asan_shadow_unittest.cc
@@ -11,6 +11,7 @@
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
+
#include "syzygy/agent/asan/asan_shadow.h"
#include "base/rand_util.h"
diff --git a/instrument/transforms/asan_transform.cc b/instrument/transforms/asan_transform.cc
index 9db46d0..6ee0f8c 100644
--- a/instrument/transforms/asan_transform.cc
+++ b/instrument/transforms/asan_transform.cc
@@ -22,6 +22,7 @@
#include "syzygy/block_graph/basic_block_assembler.h"
#include "syzygy/block_graph/block_builder.h"
#include "syzygy/pe/block_util.h"
+#include "syzygy/pe/pe_utils.h"
#include "syzygy/pe/transforms/add_imports_transform.h"
#include "third_party/distorm/files/include/mnemonics.h"
#include "third_party/distorm/files/src/x86defs.h"
@@ -41,6 +42,7 @@
using block_graph::Immediate;
using block_graph::Instruction;
using block_graph::Operand;
+using block_graph::TypedBlock;
using block_graph::Value;
using core::Register;
using core::RegisterCode;
@@ -324,6 +326,18 @@
bool AsanTransform::PreBlockGraphIteration(BlockGraph* block_graph,
BlockGraph::Block* header_block) {
+ bool already_instrumented = false;
+ // Ensure that this image has not already been instrumented.
+ if (!pe::HasImportEntry(header_block, kSyzyAsanDll, &already_instrumented)) {
+ LOG(ERROR) << "Unable to check if the image is already instrumented.";
+ return false;
+ }
+
+ if (already_instrumented) {
+ LOG(ERROR) << "The image is already instrumented.";
+ return false;
+ }
+
// Add an import entry for the ASAN runtime.
AddImportsTransform::ImportedModule import_module(asan_dll_name_.c_str());
diff --git a/pe/pe_utils.cc b/pe/pe_utils.cc
index 319c7b6..7248093 100644
--- a/pe/pe_utils.cc
+++ b/pe/pe_utils.cc
@@ -1,4 +1,4 @@
-// Copyright 2012 Google Inc.
+// Copyright 2012 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
@@ -14,6 +14,7 @@
#include "syzygy/pe/pe_utils.h"
+#include "base/string_util.h"
#include "syzygy/block_graph/typed_block.h"
#include "syzygy/pe/dos_stub.h"
@@ -26,6 +27,16 @@
namespace {
+// A simple struct that can be used to let us access strings using TypedBlock.
+struct StringStruct {
+ const char string[1];
+};
+
+typedef TypedBlock<IMAGE_DOS_HEADER> DosHeader;
+typedef TypedBlock<IMAGE_IMPORT_DESCRIPTOR> ImageImportDescriptor;
+typedef TypedBlock<IMAGE_NT_HEADERS> NtHeaders;
+typedef TypedBlock<StringStruct> String;
+
template <typename BlockPtr>
BlockPtr UncheckedGetNtHeadersBlockFromDosHeaderBlock(
BlockPtr dos_header_block) {
@@ -338,4 +349,70 @@
return true;
}
+bool HasImportEntry(block_graph::BlockGraph::Block* header_block,
+ const base::StringPiece& dll_name,
+ bool* has_import_entry) {
+ DCHECK(header_block != NULL);
+ DCHECK(dll_name != NULL);
+ DCHECK(!dll_name.empty());
+ DCHECK(has_import_entry != NULL);
+
+ *has_import_entry = false;
+
+ DosHeader dos_header;
+ NtHeaders nt_headers;
+ if (!dos_header.Init(0, header_block) ||
+ !dos_header.Dereference(dos_header->e_lfanew, &nt_headers)) {
+ LOG(ERROR) << "Unable to cast image headers.";
+ return false;
+ }
+
+ BlockGraph::Block* image_import_descriptor_block;
+ IMAGE_DATA_DIRECTORY* import_directory =
+ nt_headers->OptionalHeader.DataDirectory + IMAGE_DIRECTORY_ENTRY_IMPORT;
+ DCHECK(nt_headers.HasReference(import_directory->VirtualAddress));
+
+ ImageImportDescriptor image_import_descriptor;
+ if (!nt_headers.Dereference(import_directory->VirtualAddress,
+ &image_import_descriptor)) {
+ // This could happen if the image import descriptor array is empty, and
+ // terminated by a *partial* null entry. However, we've not yet seen that.
+ LOG(ERROR) << "Failed to dereference Image Import Descriptor Array.";
+ return false;
+ }
+
+ image_import_descriptor_block = image_import_descriptor.block();
+
+ ImageImportDescriptor iida;
+ if (!iida.Init(0, image_import_descriptor_block)) {
+ LOG(ERROR) << "Unable to cast Image Import Descriptor.";
+ return false;
+ }
+
+ // The array is NULL terminated with a potentially incomplete descriptor so
+ // we can't use ElementCount - 1.
+ DCHECK_GT(image_import_descriptor_block->size(), 0U);
+ size_t descriptor_count =
+ (common::AlignUp(image_import_descriptor_block->size(),
+ sizeof(IMAGE_IMPORT_DESCRIPTOR)) /
+ sizeof(IMAGE_IMPORT_DESCRIPTOR)) - 1;
+
+ for (size_t iida_index = 0; iida_index < descriptor_count; ++iida_index) {
+ String ref_dll_name;
+ if (!iida.Dereference(iida[iida_index].Name, &ref_dll_name)) {
+ LOG(ERROR) << "Unable to dereference DLL name.";
+ return false;
+ }
+
+ size_t max_len = ref_dll_name.ElementCount();
+ if (base::strncasecmp(ref_dll_name->string, dll_name.data(),
+ max_len) == 0) {
+ *has_import_entry = true;
+ break;
+ }
+ }
+
+ return true;
+}
+
} // namespace pe
diff --git a/pe/pe_utils.h b/pe/pe_utils.h
index 0c1bff7..38757e4 100644
--- a/pe/pe_utils.h
+++ b/pe/pe_utils.h
@@ -1,4 +1,4 @@
-// Copyright 2012 Google Inc.
+// Copyright 2012 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
@@ -114,6 +114,15 @@
bool GetTlsInitializers(block_graph::BlockGraph::Block* dos_header_block,
EntryPointSet* entry_points);
+// Check if an image contains an import entry.
+// @param The image's header-block.
+// @param dll_name The name of the DLL.
+// @param contains_dependence Boolean to indicate if the image contains the
+// import entry.
+// @returns true in case of success, false otherwise.
+bool HasImportEntry(block_graph::BlockGraph::Block* header_block,
+ const base::StringPiece& dll_name,
+ bool* has_import_entry);
} // namespace pe
#endif // SYZYGY_PE_PE_UTILS_H_
diff --git a/pe/pe_utils_unittest.cc b/pe/pe_utils_unittest.cc
index 345e9ca..7ad8369 100644
--- a/pe/pe_utils_unittest.cc
+++ b/pe/pe_utils_unittest.cc
@@ -1,4 +1,4 @@
-// Copyright 2012 Google Inc.
+// Copyright 2012 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
@@ -16,11 +16,13 @@
#include "gmock/gmock.h"
#include "gtest/gtest.h"
+#include "syzygy/pe/transforms/add_imports_transform.h"
namespace pe {
using block_graph::BlockGraph;
using core::RelativeAddress;
+using pe::transforms::AddImportsTransform;
namespace {
@@ -304,4 +306,26 @@
EXPECT_EQ(tls_initializer_block_, entry_points.begin()->first);
}
+TEST_F(PEUtilsTest, HasImportEntry) {
+ // Creates an imported module.
+ AddImportsTransform::ImportedModule module("foo.dll");
+ const char* kFooFunc = "foo_func";
+ size_t function_foo = module.AddSymbol(kFooFunc);
+ ASSERT_EQ(kFooFunc, module.GetSymbolName(function_foo));
+
+ // Apply the transform to add this module import to the block-graph.
+ AddImportsTransform transform;
+ transform.AddModule(&module);
+ ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(
+ &transform, &block_graph_, dos_header_block_));
+
+ // Ensure that we can find this module, and that we can't find a
+ // non-imported module.
+ bool has_import = false;
+ EXPECT_TRUE(HasImportEntry(dos_header_block_, "foo.dll", &has_import));
+ EXPECT_TRUE(has_import);
+ EXPECT_TRUE(HasImportEntry(dos_header_block_, "bar.dll", &has_import));
+ EXPECT_FALSE(has_import);
+}
+
} // namespace pe
diff --git a/pe/transforms/add_imports_transform.cc b/pe/transforms/add_imports_transform.cc
index 02d0c08..5c06a57 100644
--- a/pe/transforms/add_imports_transform.cc
+++ b/pe/transforms/add_imports_transform.cc
@@ -1,4 +1,4 @@
-// Copyright 2012 Google Inc.
+// Copyright 2012 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
@@ -219,6 +219,7 @@
// The array is NULL terminated with a potentially incomplete descriptor so
// we can't use ElementCount - 1.
+ DCHECK_GT(iida_block->size(), 0U);
size_t descriptor_count =
(common::AlignUp(iida_block->size(), sizeof(IMAGE_IMPORT_DESCRIPTOR)) /
sizeof(IMAGE_IMPORT_DESCRIPTOR)) - 1;
@@ -230,7 +231,7 @@
return false;
}
- size_t max_len = dll_name.block()->size() - dll_name.offset();
+ size_t max_len = dll_name.ElementCount();
if (base::strncasecmp(dll_name->string, module_name, max_len) == 0) {
// This should never fail, but we sanity check it nonetheless.
bool result = iid->Init(iida.OffsetOf(iida[iida_index]), iida.block());
@@ -487,7 +488,6 @@
} // namespace
-
const char AddImportsTransform::kTransformName[] = "AddImportsTransform";
AddImportsTransform::AddImportsTransform()
