1f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)// Copyright 2013 The Chromium Authors. All rights reserved.
2f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)// Use of this source code is governed by a BSD-style license that can be
3f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)// found in the LICENSE file.
4f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)
5f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)#include "chrome/browser/install_verification/win/loaded_modules_snapshot.h"
6f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)
7f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)#include <Psapi.h>
8f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)
9f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)#include "base/logging.h"
10f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)
11f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)bool GetLoadedModulesSnapshot(std::vector<HMODULE>* snapshot) {
12f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)  DCHECK(snapshot);
13f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)  DCHECK_EQ(0u, snapshot->size());
14f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)  snapshot->resize(1);
15f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)
16f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)  HANDLE process = ::GetCurrentProcess();
17f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)
18f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)  // We will retry at least once after first determining |bytes_required|. If
19f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)  // the list of modules changes after we receive |bytes_required| we may retry
20f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)  // more than once.
21f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)  int retries_remaining = 5;
22f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)  do {
23f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)    DWORD bytes_required = 0;
24f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)    // EnumProcessModules returns 'success' even if the buffer size is too
25f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)    // small.
26f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)    if (!::EnumProcessModules(process,
27f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)                              &(*snapshot)[0],
28f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)                              snapshot->size() * sizeof(HMODULE),
29f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)                              &bytes_required)) {
30f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)      DPLOG(ERROR) << "::EnumProcessModules failed.";
31f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)      return false;
32f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)    }
33f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)    DCHECK_EQ(0u, bytes_required % sizeof(HMODULE));
34f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)    size_t num_modules = bytes_required / sizeof(HMODULE);
35f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)    if (num_modules <= snapshot->size()) {
36f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)      // Buffer size was too big, presumably because a module was unloaded.
37f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)      snapshot->erase(snapshot->begin() + num_modules, snapshot->end());
38f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)      return true;
39f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)    } else if (num_modules == 0) {
40f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)      DLOG(ERROR) << "Can't determine the module list size.";
41f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)      return false;
42f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)    } else {
43f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)      // Buffer size was too small. Try again with a larger buffer.
44f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)      snapshot->resize(num_modules, NULL);
45f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)    }
46f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)  } while (--retries_remaining);
47f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)
48f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)  DLOG(ERROR) << "Failed to enumerate modules.";
49f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)  return false;
50f2477e01787aa58f445919b809d89e252beef54fTorne (Richard Coles)}
51