1  #pragma once
  2  
  3  #include <inttypes.h>
  4  #include <algorithm>
  5  #include <cstdint>
  6  #include <fcntl.h>
  7  #include <string>
  8  #include <sys/stat.h>
  9  #include <unistd.h>
 10  #include <vector>
 11  #include <cstdio>
 12  #include <cstdlib>
 13  
 14  #include "KittyUtils.hpp"
 15  
 16  #ifdef __APPLE__
 17  #include <mach/mach.h>
 18  
 19  /**
 20   * @brief Validates the memory access rights for a given address range.
 21   *
 22   * This class uses Mach kernel APIs to determine the read, write, and execute permissions of a memory address.
 23   */
 24  class KittyPtrValidator
 25  {
 26  private:
 27      struct RegionInfo
 28      {
 29          uintptr_t start;
 30          uintptr_t end;
 31          bool readable;
 32          bool writable;
 33          bool executable;
 34  
 35          RegionInfo(uintptr_t s, uintptr_t e, bool r, bool w, bool x)
 36              : start(s), end(e), readable(r), writable(w), executable(x)
 37          {
 38          }
 39  
 40          inline bool canMergeWith(const RegionInfo &other) const
 41          {
 42              return end == other.start && readable == other.readable && writable == other.writable &&
 43                     executable == other.executable;
 44          }
 45      };
 46  
 47      std::vector<RegionInfo> cachedRegions_;
 48      const mach_port_t task_ = mach_task_self();
 49      const size_t page_size_ = sysconf(_SC_PAGESIZE);
 50      bool use_cache_ = true;
 51      size_t last_region_index_ = 0;
 52  
 53      bool _findRegion(uintptr_t addr, RegionInfo *region);
 54  
 55  public:
 56      KittyPtrValidator()
 57          : task_(mach_task_self()), page_size_(sysconf(_SC_PAGESIZE)), use_cache_(false), last_region_index_(0)
 58      {
 59      }
 60  
 61      /**
 62       * @brief Constructs a KittyPtrValidator object with the specified task and whether to use a cache.
 63       *
 64       * @param task process task.
 65       * @param use_cache Determines if the cache should be used.
 66       */
 67      KittyPtrValidator(mach_port_t task, bool use_cache)
 68          : task_(task), page_size_(sysconf(_SC_PAGESIZE)), use_cache_(use_cache), last_region_index_(0)
 69      {
 70          if (use_cache_)
 71              refreshRegionCache();
 72      }
 73  
 74      /**
 75       * @brief Sets whether to use cached region information.
 76       *
 77       * @param use_cache True to use cached region information, false to clear the cache.
 78       */
 79      inline void setUseCache(bool use_cache)
 80      {
 81          use_cache_ = use_cache;
 82          if (!use_cache_)
 83          {
 84              cachedRegions_.clear();
 85              last_region_index_ = 0;
 86          }
 87          else
 88          {
 89              refreshRegionCache();
 90          }
 91      }
 92  
 93      /**
 94       * @brief Checks if a pointer is readable.
 95       *
 96       * @param ptr The memory address to check.
 97       * @param len The length of the memory range to check.
 98       * @return true if the memory address is readable, false otherwise.
 99       */
100      bool isPtrReadable(uintptr_t ptr, size_t len = sizeof(void *));
101  
102      /**
103       * @brief Checks if a pointer is writable.
104       *
105       * @param ptr The memory address to check.
106       * @param len The length of the memory range to check.
107       * @return true if the memory address is writable, false otherwise.
108       */
109      bool isPtrWritable(uintptr_t ptr, size_t len = sizeof(void *));
110  
111      /**
112       * @brief Checks if a pointer is executable.
113       *
114       * @param ptr The memory address to check.
115       * @param len The length of the memory range to check.
116       * @return true if the memory address is executable, false otherwise.
117       */
118      bool isPtrExecutable(uintptr_t ptr, size_t len = sizeof(void *));
119  
120      /**
121       * @brief Checks if a pointer is within the address space of the current task.
122       *
123       * @param ptr The memory address to check.
124       * @return true if the pointer is within the address space, false otherwise.
125       */
126      inline bool isPtrInAddressSpace(uintptr_t ptr)
127      {
128          if (ptr == 0)
129              return false;
130          RegionInfo region(0, 0, false, false, false);
131          return _findRegion(ptr, &region);
132      }
133  
134      inline bool isPtrReadable(const void *ptr, size_t len = sizeof(void *))
135      {
136          return ptr && isPtrReadable(uintptr_t(ptr), len);
137      }
138      inline bool isPtrWritable(const void *ptr, size_t len = sizeof(void *))
139      {
140          return ptr && isPtrWritable(uintptr_t(ptr), len);
141      }
142      inline bool isPtrExecutable(const void *ptr, size_t len = sizeof(void *))
143      {
144          return ptr && isPtrExecutable(uintptr_t(ptr), len);
145      }
146      inline bool isPtrInAddressSpace(const void *ptr)
147      {
148          return ptr && isPtrInAddressSpace(uintptr_t(ptr));
149      }
150  
151      /**
152       * @brief Clears the cached region information.
153       */
154      inline void clearCache()
155      {
156          cachedRegions_.clear();
157          last_region_index_ = 0;
158      }
159  
160      /**
161       * @brief Refreshes the cached region information.
162       */
163      void refreshRegionCache();
164  
165      /**
166       * @brief Retrieves the cached region information.
167       *
168       * @return A vector of RegionInfo objects representing the cached regions.
169       */
170      inline std::vector<RegionInfo> cachedRegions() const
171      {
172          return cachedRegions_;
173      }
174  };
175  
176  #else
177  
178  /**
179   * @brief Validates the memory access rights for a given address range.
180   *
181   * This class uses process maps to determine the read, write, and execute permissions of a memory address.
182   */
183  class KittyPtrValidator
184  {
185  private:
186      struct RegionInfo
187      {
188          uintptr_t start;
189          uintptr_t end;
190          bool readable;
191          bool writable;
192          bool executable;
193  
194          RegionInfo(uintptr_t s, uintptr_t e, bool r, bool w, bool x)
195              : start(s), end(e), readable(r), writable(w), executable(x)
196          {
197          }
198  
199          inline bool canMergeWith(const RegionInfo &other) const
200          {
201              return end == other.start && readable == other.readable && writable == other.writable &&
202                     executable == other.executable;
203          }
204      };
205  
206      std::vector<RegionInfo> cachedRegions_;
207      pid_t pid_ = getpid();
208      const size_t page_size_ = sysconf(_SC_PAGESIZE);
209      bool use_cache_ = true;
210      size_t last_region_index_ = 0;
211  
212      bool _parseMapsLine(const char *line, RegionInfo *region);
213  
214      bool _findRegion(uintptr_t addr, RegionInfo *region);
215  
216  public:
217      KittyPtrValidator() : pid_(getpid()), page_size_(sysconf(_SC_PAGESIZE)), use_cache_(false), last_region_index_(0)
218      {
219      }
220  
221      /**
222       * @brief Constructs a KittyPtrValidator object with the specified process ID and whether to use a cache.
223       *
224       * @param pid The process ID.
225       * @param use_cache Determines if the cache should be used.
226       */
227      KittyPtrValidator(pid_t pid, bool use_cache)
228          : pid_(pid), page_size_(sysconf(_SC_PAGESIZE)), use_cache_(use_cache), last_region_index_(0)
229      {
230          if (use_cache_)
231              refreshRegionCache();
232      }
233  
234      /**
235       * @brief Sets whether to use cached region information.
236       *
237       * @param use_cache True to use cached region information, false to clear the cache.
238       */
239      inline void setUseCache(bool use_cache)
240      {
241          use_cache_ = use_cache;
242          if (!use_cache_)
243          {
244              cachedRegions_.clear();
245              last_region_index_ = 0;
246          }
247          else
248          {
249              refreshRegionCache();
250          }
251      }
252  
253      /**
254       * @brief Sets the process ID to query memory regions.
255       *
256       * @param pid The process ID to query.
257       */
258      inline void setPID(pid_t pid)
259      {
260          cachedRegions_.clear();
261          last_region_index_ = 0;
262          pid_ = pid;
263  
264          if (use_cache_)
265          {
266              refreshRegionCache();
267          }
268      }
269  
270      /**
271       * @brief Checks if a pointer is readable.
272       *
273       * @param ptr The memory address to check.
274       * @param len The length of the memory range to check.
275       * @return true if the memory address is readable, false otherwise.
276       */
277      bool isPtrReadable(uintptr_t ptr, size_t len = sizeof(void *));
278  
279      /**
280       * @brief Checks if a pointer is writable.
281       *
282       * @param ptr The memory address to check.
283       * @param len The length of the memory range to check.
284       * @return true if the memory address is writable, false otherwise.
285       */
286      bool isPtrWritable(uintptr_t ptr, size_t len = sizeof(void *));
287  
288      /**
289       * @brief Checks if a pointer is executable.
290       *
291       * @param ptr The memory address to check.
292       * @param len The length of the memory range to check.
293       * @return true if the memory address is executable, false otherwise.
294       */
295      bool isPtrExecutable(uintptr_t ptr, size_t len = sizeof(void *));
296  
297      /**
298       * @brief Checks if a pointer is within the address space of the current process.
299       *
300       * @param ptr The memory address to check.
301       * @return true if the pointer is within the address space, false otherwise.
302       */
303      inline bool isPtrInAddressSpace(uintptr_t ptr)
304      {
305          if (ptr == 0)
306              return false;
307          ptr = KittyUtils::untagHeepPtr(ptr);
308          RegionInfo region(0, 0, false, false, false);
309          return _findRegion(ptr, &region);
310      }
311  
312      inline bool isPtrReadable(const void *ptr, size_t len = sizeof(void *))
313      {
314          return ptr && isPtrReadable(uintptr_t(ptr), len);
315      }
316      inline bool isPtrWritable(const void *ptr, size_t len = sizeof(void *))
317      {
318          return ptr && isPtrWritable(uintptr_t(ptr), len);
319      }
320      inline bool isPtrExecutable(const void *ptr, size_t len = sizeof(void *))
321      {
322          return ptr && isPtrExecutable(uintptr_t(ptr), len);
323      }
324      inline bool isPtrInAddressSpace(const void *ptr)
325      {
326          return ptr && isPtrInAddressSpace(uintptr_t(ptr));
327      }
328  
329      /**
330       * @brief Clears the cached region information.
331       */
332      inline void clearCache()
333      {
334          cachedRegions_.clear();
335          last_region_index_ = 0;
336      }
337  
338      /**
339       * @brief Refreshes the cached region information.
340       */
341      void refreshRegionCache();
342  
343      /**
344       * @brief Retrieves the cached region information.
345       *
346       * @return A vector of RegionInfo objects representing the cached regions.
347       */
348      inline std::vector<RegionInfo> cachedRegions() const
349      {
350          return cachedRegions_;
351      }
352  };
353  
354  #endif