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