base/memory/weak_auto_reset.h

// Copyright 2022 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef BASE_MEMORY_WEAK_AUTO_RESET_H_
#define BASE_MEMORY_WEAK_AUTO_RESET_H_

#include "base/memory/weak_ptr.h"

namespace base {

// Sets a field of an object to a specified value, then returns it to its
// original value when the WeakAutoReset instance goes out of scope. Because a
// weak pointer is used, if the target object is destroyed, no attempt is made
// to restore the original value and no UAF occurs.
//
// Note that as of C++17 we can use CTAD to infer template parameters from
// constructor args; it is valid to write:
//   WeakAutoReset war(myobj->GetWeakPtr(), &MyClass::member_, new_value);
// without specifying explicit types in the classname.
template <class T, class U>
class WeakAutoReset {
 public:
  // Create an empty object that does nothing, you may move a value into this
  // object via assignment.
  WeakAutoReset() = default;

  // Sets member `field` of object pointed to by `ptr` to `new_value`. `ptr`
  // must be valid at time of construction. If `ptr` is still valid when this
  // object goes out of scope, the member will be returned to its original
  // value.
  WeakAutoReset(base::WeakPtr<T> ptr, U T::*field, U new_value)
      : ptr_(ptr),
        field_(field),
        old_value_(std::exchange(ptr.get()->*field, std::move(new_value))) {}

  // Move constructor.
  WeakAutoReset(WeakAutoReset&& other)
      : ptr_(std::move(other.ptr_)),
        field_(std::exchange(other.field_, nullptr)),
        old_value_(std::move(other.old_value_)) {}

  // Move assignment operator.
  WeakAutoReset& operator=(WeakAutoReset&& other) {
    if (this != &other) {
      // If we're already tracking a value, make sure to restore it before
      // overwriting our target.
      Reset();
      ptr_ = std::move(other.ptr_);
      field_ = std::exchange(other.field_, nullptr);
      old_value_ = std::move(other.old_value_);
    }
    return *this;
  }

  ~WeakAutoReset() { Reset(); }

 private:
  void Reset() {
    if (ptr_)
      ptr_.get()->*field_ = std::move(old_value_);
  }

  base::WeakPtr<T> ptr_;
  U T::*field_ = nullptr;
  U old_value_ = U();
};

}  // namespace base

#endif  // BASE_MEMORY_WEAK_AUTO_RESET_H_