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    :  #include "syzygy/agent/common/dll_notifications.h"
  16    :  
  17    :  #include "base/bind.h"
  18    :  #include "base/files/file_path.h"
  19    :  #include "gmock/gmock.h"
  20    :  #include "gtest/gtest.h"
  21    :  #include "syzygy/core/unittest_util.h"
  22    :  
  23    :  namespace agent {
  24    :  namespace common {
  25    :  
  26    :  namespace {
  27    :  
  28    :  using testing::_;
  29    :  using testing::Eq;
  30    :  
  31    :  class NotificationReceiver {
  32    :   public:
  33    :    MOCK_METHOD5(OnNotification,
  34    :                 void(DllNotificationWatcher::EventType type,
  35    :                      HMODULE module,
  36    :                      size_t module_size,
  37    :                      const base::StringPiece16& dll_path,
  38  E :                      const base::StringPiece16& dll_base_name));
  39    :  };
  40    :  
  41    :  class DllNotificationWatcherTest : public testing::Test {
  42    :   public:
  43  E :    DllNotificationWatcherTest() : test_dll_(NULL) {
  44  E :    }
  45    :  
  46  E :    virtual void SetUp() OVERRIDE {
  47  E :      test_dll_path_ = testing::GetExeRelativePath(L"test_dll.dll");
  48  E :    }
  49    :  
  50  E :    virtual void TearDown() OVERRIDE {
  51  E :      if (test_dll_ != NULL) {
  52  i :        ::FreeLibrary(test_dll_);
  53    :      }
  54  E :    }
  55    :  
  56  E :    void LoadTestDll() {
  57  E :      test_dll_ = ::LoadLibrary(test_dll_path_.value().c_str());
  58  E :      ASSERT_NE(static_cast<HMODULE>(NULL), test_dll_);
  59  E :    }
  60    :  
  61  E :    void UnloadTestDll() {
  62  E :      ASSERT_NE(static_cast<HMODULE>(NULL), test_dll_);
  63  E :      ASSERT_TRUE(::FreeLibrary(test_dll_));
  64  E :      test_dll_ = NULL;
  65  E :    }
  66    :  
  67    :    testing::StrictMock<NotificationReceiver> receiver_;
  68    :    base::FilePath test_dll_path_;
  69    :    HMODULE test_dll_;
  70    :  };
  71    :  
  72    :  }  // namespace
  73    :  
  74  E :  TEST_F(DllNotificationWatcherTest, Init) {
  75  E :    DllNotificationWatcher watcher;
  76    :  
  77    :    ASSERT_TRUE(watcher.Init(
  78    :        base::Bind(&NotificationReceiver::OnNotification,
  79  E :                   base::Unretained(&receiver_))));
  80    :  
  81    :  
  82    :    // We expect DLL load notifications for test_dll_ and its import dependency.
  83    :    EXPECT_CALL(receiver_,
  84    :                OnNotification(DllNotificationWatcher::DllLoaded,
  85    :                               _, _,
  86    :                               testing::Eq(test_dll_path_.value()),
  87  E :                               testing::Eq(test_dll_path_.BaseName().value())));
  88    :    EXPECT_CALL(receiver_,
  89    :                OnNotification(DllNotificationWatcher::DllLoaded,
  90    :                               _, _,
  91    :                               _,
  92  E :                               testing::Eq(L"export_dll.dll")));
  93    :  
  94    :    // Load the DLL.
  95  E :    ASSERT_NO_FATAL_FAILURE(LoadTestDll());
  96    :  
  97    :    // Now we should see unload notification for the same DLLs.
  98    :    EXPECT_CALL(receiver_,
  99    :                OnNotification(DllNotificationWatcher::DllUnloaded,
 100    :                               _, _,
 101    :                               testing::Eq(test_dll_path_.value()),
 102  E :                               testing::Eq(test_dll_path_.BaseName().value())));
 103    :    EXPECT_CALL(receiver_,
 104    :                OnNotification(DllNotificationWatcher::DllUnloaded,
 105    :                               _, _,
 106    :                               _,
 107  E :                               testing::Eq(L"export_dll.dll")));
 108  E :    UnloadTestDll();
 109  E :  }
 110    :  
 111  E :  TEST_F(DllNotificationWatcherTest, Reset) {
 112  E :    DllNotificationWatcher watcher;
 113    :  
 114    :    ASSERT_TRUE(watcher.Init(
 115    :        base::Bind(&NotificationReceiver::OnNotification,
 116  E :                   base::Unretained(&receiver_))));
 117    :    // Reset the watcher - this should unregister, and we should get no callbacks.
 118  E :    watcher.Reset();
 119    :  
 120    :    // We should get no notifications after resetting.
 121  E :    ASSERT_NO_FATAL_FAILURE(LoadTestDll());
 122  E :    UnloadTestDll();
 123  E :  }
 124    :  
 125  E :  TEST_F(DllNotificationWatcherTest, ResetOnDeletion) {
 126    :    {
 127  E :      DllNotificationWatcher watcher;
 128    :  
 129    :      ASSERT_TRUE(watcher.Init(
 130    :          base::Bind(&NotificationReceiver::OnNotification,
 131  E :                     base::Unretained(&receiver_))));
 132  E :    }
 133    :  
 134    :    // We should get no notifications after the instance goes out of scope.
 135    :    // We should get no notifications after resetting.
 136  E :    ASSERT_NO_FATAL_FAILURE(LoadTestDll());
 137  E :    UnloadTestDll();
 138  E :  }
 139    :  
 140    :  }  // namespace common
 141    :  }  // namespace agent