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