/*

 * Copyright 2001 Sun Microsystems, Inc.  All Rights Reserved.

 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.

 *

 * This code is free software; you can redistribute it and/or modify it

 * under the terms of the GNU General Public License version 2 only, as

 * published by the Free Software Foundation.

 *

 * This code is distributed in the hope that it will be useful, but WITHOUT

 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or

 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License

 * version 2 for more details (a copy is included in the LICENSE file that

 * accompanied this code).

 *

 * You should have received a copy of the GNU General Public License version

 * 2 along with this work; if not, write to the Free Software Foundation,

 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

 *

 * Please contact Sun Microsystems, Inc., 4150 Network Circle, Santa Clara,

 * CA 95054 USA or visit www.sun.com if you need additional information or

 * have any questions.

 *

 */

// Disable too-long symbol warnings

#pragma warning ( disable : 4786 )

#include "libInfo.hpp"

#include "nt4internals.hpp"

#include "isNT4.hpp"

#include "toolHelp.hpp"

#include <assert.h>

using namespace std;

typedef void LibInfoImplFunc(DWORD pid, vector<LibInfo>& info);

static void libInfoImplNT4(DWORD pid, vector<LibInfo>& info);

static void libInfoImplToolHelp(DWORD pid, vector<LibInfo>& info);

void

libInfo(DWORD pid, vector<LibInfo>& info) {

  static LibInfoImplFunc* impl = NULL;

  if (impl == NULL) {

    // See which operating system we're on

    impl = (isNT4() ? &libInfoImplNT4 : &libInfoImplToolHelp);

  }

  assert(impl != NULL);

  (*impl)(pid, info);

}

static ULONG

ModuleCount(NT4::PDEBUG_BUFFER db) {

  return db->ModuleInformation ? *PULONG(db->ModuleInformation) : 0;

}

#define MAX2(a, b) (((a) < (b)) ? (b) : (a))

static void

libInfoImplNT4(DWORD pid, vector<LibInfo>& info) {

  static EnumProcessModulesFunc*   enumFunc = NULL;

  static GetModuleFileNameExFunc*  fnFunc   = NULL;

  static GetModuleInformationFunc* infoFunc = NULL;

  if (enumFunc == NULL) {

    HMODULE dll = loadPSAPIDLL();

    enumFunc = (EnumProcessModulesFunc*)   GetProcAddress(dll, "EnumProcessModules");

    fnFunc   = (GetModuleFileNameExFunc*)  GetProcAddress(dll, "GetModuleFileNameExA");

    infoFunc = (GetModuleInformationFunc*) GetProcAddress(dll, "GetModuleInformation");

    assert(enumFunc != NULL);

    assert(fnFunc   != NULL);

    assert(infoFunc != NULL);

  }

  static HMODULE* mods = new HMODULE[256];

  static int      numMods = 256;

  if (mods == NULL) {

    mods = new HMODULE[numMods];

    if (mods == NULL) {

      return;

    }

  }

  bool done = false;

  HANDLE proc = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);

  if (proc == NULL) {

    return;

  }

  do {

    DWORD bufSize = numMods * sizeof(HMODULE);

    DWORD neededSize;

    if (!(*enumFunc)(proc, mods, bufSize, &neededSize)) {

      // Enum failed

      CloseHandle(proc);

      return;

    }

    int numFetched = neededSize / sizeof(HMODULE);

    if (numMods < numFetched) {

      // Grow buffer

      numMods = MAX2(numFetched, 2 * numMods);

      delete[] mods;

      mods = new HMODULE[numMods];

      if (mods == NULL) {

        CloseHandle(proc);

        return;

      }

    } else {

      char filename[MAX_PATH];

      MODULEINFO modInfo;

      // Iterate through and fetch each one's info

      for (int i = 0; i < numFetched; i++) {

        if (!(*fnFunc)(proc, mods[i], filename, MAX_PATH)) {

          CloseHandle(proc);

          return;

        }

        if (!(*infoFunc)(proc, mods[i], &modInfo, sizeof(MODULEINFO))) {

          CloseHandle(proc);

          return;

        }

        info.push_back(LibInfo(string(filename), (void*) modInfo.lpBaseOfDll));

      }

      done = true;

    }

  } while (!done);

  CloseHandle(proc);

  return;

}

void

libInfoImplToolHelp(DWORD pid, vector<LibInfo>& info) {

  using namespace ToolHelp;

  static CreateToolhelp32SnapshotFunc* snapshotFunc = NULL;

  static Module32FirstFunc*            firstFunc    = NULL;

  static Module32NextFunc*             nextFunc     = NULL;

  if (snapshotFunc == NULL) {

    HMODULE dll = loadDLL();

    snapshotFunc =

      (CreateToolhelp32SnapshotFunc*) GetProcAddress(dll,

                                                     "CreateToolhelp32Snapshot");

    firstFunc = (Module32FirstFunc*) GetProcAddress(dll,

                                                    "Module32First");

    nextFunc = (Module32NextFunc*) GetProcAddress(dll,

                                                  "Module32Next");

    assert(snapshotFunc != NULL);

    assert(firstFunc    != NULL);

    assert(nextFunc     != NULL);

  }

  HANDLE snapshot = (*snapshotFunc)(TH32CS_SNAPMODULE, pid);

  if (snapshot == (HANDLE) -1) {

    // Error occurred during snapshot

    return;

  }

  // Iterate

  MODULEENTRY32 module;

  if ((*firstFunc)(snapshot, &module)) {

    do {

      info.push_back(LibInfo(string(module.szExePath), (void*) module.modBaseAddr));

    } while ((*nextFunc)(snapshot, &module));

  }

  CloseHandle(snapshot);

}