Home | History | Annotate | Download | only in base 
      1 // Copyright 2010 Google Inc. All Rights Reserved.
      2 
      3 
      4 #ifndef TALK_BASE_WIN32TOOLHELP_H_
      5 #define TALK_BASE_WIN32TOOLHELP_H_
      6 
      7 #ifndef WIN32
      8 #error WIN32 Only
      9 #endif
     10 
     11 #include "talk/base/win32.h"
     12 
     13 // Should be included first, but that causes redefinitions.
     14 #include <tlhelp32.h>
     15 
     16 #include "talk/base/constructormagic.h"
     17 
     18 namespace talk_base {
     19 
     20 // The toolhelp api used to enumerate processes and their modules
     21 // on Windows is very repetetive and clunky to use. This little
     22 // template wraps it to make it a little more programmer friendly.
     23 //
     24 // Traits: Traits type that adapts the enumerator to the corresponding
     25 //         win32 toolhelp api. Each traits class need to:
     26 //         - define the type of the enumerated data as a public symbol Type
     27 //
     28 //         - implement bool First(HANDLE, T*) normally calls a
     29 //           Xxxx32First method in the toolhelp API. Ex Process32First(...)
     30 //
     31 //         - implement bool Next(HANDLE, T*) normally calls a
     32 //           Xxxx32Next method in the toolhelp API. Ex Process32Next(...)
     33 //
     34 //         - implement bool CloseHandle(HANDLE)
     35 //
     36 template<typename Traits>
     37 class ToolhelpEnumeratorBase {
     38  public:
     39   ToolhelpEnumeratorBase(HANDLE snapshot)
     40       : snapshot_(snapshot), broken_(false), first_(true) {
     41 
     42     // Clear out the Traits::Type structure instance.
     43     Zero(&current_);
     44   }
     45 
     46   virtual ~ToolhelpEnumeratorBase() {
     47     Close();
     48   }
     49 
     50   // Moves forward to the next object using the First and Next
     51   // pointers. If either First or Next ever indicates an failure
     52   // all subsequent calls to this method will fail; the enumerator
     53   // object is considered broken.
     54   bool Next() {
     55     if (!Valid()) {
     56       return false;
     57     }
     58 
     59     // Move the iteration forward.
     60     current_.dwSize = sizeof(typename Traits::Type);
     61     bool incr_ok = false;
     62     if (first_) {
     63       incr_ok = Traits::First(snapshot_, &current_);
     64       first_ = false;
     65     } else {
     66       incr_ok = Traits::Next(snapshot_, &current_);
     67     }
     68 
     69     if (!incr_ok) {
     70       Zero(&current_);
     71       broken_ = true;
     72     }
     73 
     74     return incr_ok;
     75   }
     76 
     77   const typename Traits::Type& current() const {
     78     return current_;
     79   }
     80 
     81   void Close() {
     82     if (snapshot_ != INVALID_HANDLE_VALUE) {
     83       Traits::CloseHandle(snapshot_);
     84       snapshot_ = INVALID_HANDLE_VALUE;
     85     }
     86   }
     87 
     88  private:
     89   // Checks the state of the snapshot handle.
     90   bool Valid() {
     91     return snapshot_ != INVALID_HANDLE_VALUE && !broken_;
     92   }
     93 
     94   static void Zero(typename Traits::Type* buff) {
     95     ZeroMemory(buff, sizeof(typename Traits::Type));
     96   }
     97 
     98   HANDLE snapshot_;
     99   typename Traits::Type current_;
    100   bool broken_;
    101   bool first_;
    102 };
    103 
    104 class ToolhelpTraits {
    105  public:
    106   static HANDLE CreateSnapshot(uint32 flags, uint32 process_id) {
    107     return CreateToolhelp32Snapshot(flags, process_id);
    108   }
    109 
    110   static bool CloseHandle(HANDLE handle) {
    111     return ::CloseHandle(handle) == TRUE;
    112   }
    113 };
    114 
    115 class ToolhelpProcessTraits : public ToolhelpTraits {
    116  public:
    117   typedef PROCESSENTRY32 Type;
    118 
    119   static bool First(HANDLE handle, Type* t) {
    120     return ::Process32First(handle, t) == TRUE;
    121   }
    122 
    123   static bool Next(HANDLE handle, Type* t) {
    124     return ::Process32Next(handle, t) == TRUE;
    125   }
    126 };
    127 
    128 class ProcessEnumerator : public ToolhelpEnumeratorBase<ToolhelpProcessTraits> {
    129  public:
    130   ProcessEnumerator()
    131       : ToolhelpEnumeratorBase(
    132            ToolhelpProcessTraits::CreateSnapshot(TH32CS_SNAPPROCESS, 0)) {
    133   }
    134 
    135  private:
    136   DISALLOW_EVIL_CONSTRUCTORS(ProcessEnumerator);
    137 };
    138 
    139 class ToolhelpModuleTraits : public ToolhelpTraits {
    140  public:
    141   typedef MODULEENTRY32 Type;
    142 
    143   static bool First(HANDLE handle, Type* t) {
    144     return ::Module32First(handle, t) == TRUE;
    145   }
    146 
    147   static bool Next(HANDLE handle, Type* t) {
    148     return ::Module32Next(handle, t) == TRUE;
    149   }
    150 };
    151 
    152 class ModuleEnumerator : public ToolhelpEnumeratorBase<ToolhelpModuleTraits> {
    153  public:
    154   explicit ModuleEnumerator(uint32 process_id)
    155       : ToolhelpEnumeratorBase(
    156             ToolhelpModuleTraits::CreateSnapshot(TH32CS_SNAPMODULE,
    157                                                  process_id)) {
    158   }
    159 
    160  private:
    161   DISALLOW_EVIL_CONSTRUCTORS(ModuleEnumerator);
    162 };
    163 
    164 }  // namespace talk_base
    165 
    166 #endif  // TALK_BASE_WIN32TOOLHELP_H_
    167