// Copyright (c) 2010 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef APP_WIN_IAT_PATCH_FUNCTION_H_
#define APP_WIN_IAT_PATCH_FUNCTION_H_
#pragma once

#include <windows.h>

#include "base/basictypes.h"

namespace app {
namespace win {

// A class that encapsulates Import Address Table patching helpers and restores
// the original function in the destructor.
//
// It will intercept functions for a specific DLL imported from another DLL.
// This is the case when, for example, we want to intercept
// CertDuplicateCertificateContext function (exported from crypt32.dll) called
// by wininet.dll.
class IATPatchFunction {
 public:
  IATPatchFunction();
  ~IATPatchFunction();

  // Intercept a function in an import table of a specific
  // module. Save the original function and the import
  // table address. These values will be used later
  // during Unpatch
  //
  // Arguments:
  // module                 Module to be intercepted
  // imported_from_module   Module that exports the 'function_name'
  // function_name          Name of the API to be intercepted
  //
  // Returns: Windows error code (winerror.h). NO_ERROR if successful
  //
  // Note: Patching a function will make the IAT patch take some "ownership" on
  // |module|.  It will LoadLibrary(module) to keep the DLL alive until a call
  // to Unpatch(), which will call FreeLibrary() and allow the module to be
  // unloaded.  The idea is to help prevent the DLL from going away while a
  // patch is still active.
  DWORD Patch(const wchar_t* module,
              const char* imported_from_module,
              const char* function_name,
              void* new_function);

  // Unpatch the IAT entry using internally saved original
  // function.
  //
  // Returns: Windows error code (winerror.h). NO_ERROR if successful
  DWORD Unpatch();

  bool is_patched() const {
    return (NULL != intercept_function_);
  }

 private:
  HMODULE module_handle_;
  void* intercept_function_;
  void* original_function_;
  IMAGE_THUNK_DATA* iat_thunk_;

  DISALLOW_COPY_AND_ASSIGN(IATPatchFunction);
};

}  // namespace win
}  // namespace app

#endif  // APP_WIN_IAT_PATCH_FUNCTION_H_