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
24namespace 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//
42template<typename Traits>
43class 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
110class 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
121class 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
134class 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
145class 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
158class 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