weak_auto_reset.h 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071
  1. // Copyright 2022 The Chromium Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #ifndef BASE_MEMORY_WEAK_AUTO_RESET_H_
  5. #define BASE_MEMORY_WEAK_AUTO_RESET_H_
  6. #include "base/memory/weak_ptr.h"
  7. namespace base {
  8. // Sets a field of an object to a specified value, then returns it to its
  9. // original value when the WeakAutoReset instance goes out of scope. Because a
  10. // weak pointer is used, if the target object is destroyed, no attempt is made
  11. // to restore the original value and no UAF occurs.
  12. //
  13. // Note that as of C++17 we can use CTAD to infer template parameters from
  14. // constructor args; it is valid to write:
  15. // WeakAutoReset war(myobj->GetWeakPtr(), &MyClass::member_, new_value);
  16. // without specifying explicit types in the classname.
  17. template <class T, class U>
  18. class WeakAutoReset {
  19. public:
  20. // Create an empty object that does nothing, you may move a value into this
  21. // object via assignment.
  22. WeakAutoReset() = default;
  23. // Sets member `field` of object pointed to by `ptr` to `new_value`. `ptr`
  24. // must be valid at time of construction. If `ptr` is still valid when this
  25. // object goes out of scope, the member will be returned to its original
  26. // value.
  27. WeakAutoReset(base::WeakPtr<T> ptr, U T::*field, U new_value)
  28. : ptr_(ptr),
  29. field_(field),
  30. old_value_(std::exchange(ptr.get()->*field, std::move(new_value))) {}
  31. // Move constructor.
  32. WeakAutoReset(WeakAutoReset&& other)
  33. : ptr_(std::move(other.ptr_)),
  34. field_(std::exchange(other.field_, nullptr)),
  35. old_value_(std::move(other.old_value_)) {}
  36. // Move assignment operator.
  37. WeakAutoReset& operator=(WeakAutoReset&& other) {
  38. if (this != &other) {
  39. // If we're already tracking a value, make sure to restore it before
  40. // overwriting our target.
  41. Reset();
  42. ptr_ = std::move(other.ptr_);
  43. field_ = std::exchange(other.field_, nullptr);
  44. old_value_ = std::move(other.old_value_);
  45. }
  46. return *this;
  47. }
  48. ~WeakAutoReset() { Reset(); }
  49. private:
  50. void Reset() {
  51. if (ptr_)
  52. ptr_.get()->*field_ = std::move(old_value_);
  53. }
  54. base::WeakPtr<T> ptr_;
  55. U T::*field_ = nullptr;
  56. U old_value_ = U();
  57. };
  58. } // namespace base
  59. #endif // BASE_MEMORY_WEAK_AUTO_RESET_H_