1    :  // Copyright 2015 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/asan/iat_patcher.h"
  16    :  
  17    :  #include <vector>
  18    :  
  19    :  #include "base/logging.h"
  20    :  #include "base/win/iat_patch_function.h"
  21    :  #include "base/win/pe_image.h"
  22    :  #include "syzygy/agent/asan/scoped_page_protections.h"
  23    :  #include "syzygy/common/align.h"
  24    :  #include "syzygy/common/com_utils.h"
  25    :  
  26    :  namespace agent {
  27    :  namespace asan {
  28    :  
  29    :  namespace {
  30    :  
  31    :  bool UpdateImportThunk(volatile PIMAGE_THUNK_DATA iat,
  32  i :                         FunctionPointer function) {
  33    :    // Writing to an IAT is inherently racy, as there may be other parties also
  34    :    // writing the same page at the same time. This gets ugly where multiple
  35    :    // parties mess with page protections, as VirtualProtect causes surprising
  36    :    // serialization. We therefore proceed with an abundance of caution, by
  37    :    // running inside an exception handler and using a compare-and-swap to
  38    :    // detect races on VM operations as well as on assignment.
  39    :  
  40  i :    __try {
  41  i :      DWORD old_fn = iat->u1.Function;
  42  i :      DWORD new_fn = reinterpret_cast<DWORD>(function);
  43    :      DWORD prev_fn =
  44  i :          ::InterlockedCompareExchange(&iat->u1.Function, new_fn, old_fn);
  45    :      // Check whether we collided on the assignment.
  46  i :      if (prev_fn != old_fn)
  47  i :        return false;
  48  i :    } __except(EXCEPTION_EXECUTE_HANDLER) {
  49    :      // We took an exception, that goes down as failure.
  50  i :      return false;
  51  i :    }
  52    :  
  53    :    // All shiny!
  54  i :    return true;
  55  i :  }
  56    :  
  57    :  class IATPatchWorker {
  58    :   public:
  59    :    explicit IATPatchWorker(const IATPatchMap& patch);
  60    :  
  61    :    bool PatchImage(base::win::PEImage* image);
  62    :  
  63    :   private:
  64    :    static bool VisitImport(const base::win::PEImage &image, LPCSTR module,
  65    :                            DWORD ordinal, LPCSTR name, DWORD hint,
  66    :                            PIMAGE_THUNK_DATA iat, PVOID cookie);
  67    :    bool OnImport(const char* name, PIMAGE_THUNK_DATA iat);
  68    :  
  69    :    ScopedPageProtections scoped_page_protections_;
  70    :    const IATPatchMap& patch_;
  71    :  
  72    :    DISALLOW_COPY_AND_ASSIGN(IATPatchWorker);
  73    :  };
  74    :  
  75  E :  IATPatchWorker::IATPatchWorker(const IATPatchMap& patch) : patch_(patch) {
  76  E :  }
  77    :  
  78  E :  bool IATPatchWorker::PatchImage(base::win::PEImage* image) {
  79  E :    DCHECK_NE(static_cast<base::win::PEImage*>(nullptr), image);
  80    :  
  81    :    // The IAT patching takes place during enumeration.
  82  E :    bool ret = image->EnumAllImports(&VisitImport, this);
  83    :  
  84    :    // Clean up whatever we soiled, success or failure be damned.
  85  E :    scoped_page_protections_.RestorePageProtections();
  86    :  
  87  E :    return ret;
  88  E :  }
  89    :  
  90    :  bool IATPatchWorker::VisitImport(
  91    :      const base::win::PEImage &image, LPCSTR module, DWORD ordinal,
  92  E :      LPCSTR name, DWORD hint, PIMAGE_THUNK_DATA iat, PVOID cookie) {
  93  E :    if (!name)
  94  E :      return true;
  95    :  
  96  E :    IATPatchWorker* worker = reinterpret_cast<IATPatchWorker*>(cookie);
  97  E :    return worker->OnImport(name, iat);
  98  E :  }
  99    :  
 100  E :  bool IATPatchWorker::OnImport(const char* name, PIMAGE_THUNK_DATA iat) {
 101  E :    auto it = patch_.find(name);
 102    :    // See whether this is a function we care about.
 103  E :    if (it == patch_.end())
 104  E :      return true;
 105    :  
 106    :    // Make the containing page writable.
 107    :    if (!scoped_page_protections_.EnsureContainingPagesWritable(
 108  E :            iat, sizeof(IMAGE_THUNK_DATA))) {
 109  i :      return false;
 110    :    }
 111    :  
 112  E :    return UpdateImportThunk(iat, it->second);
 113  E :  }
 114    :  
 115    :  }  // namespace
 116    :  
 117  E :  bool PatchIATForModule(HMODULE module, const IATPatchMap& patch_map) {
 118  E :    base::win::PEImage image(module);
 119    :  
 120  E :    if (!image.VerifyMagic())
 121  i :      return false;
 122    :  
 123  E :    IATPatchWorker worker(patch_map);
 124    :  
 125  E :    return worker.PatchImage(&image);
 126  E :  }
 127    :  
 128    :  }  // namespace asan
 129    :  }  // namespace agent