[RFC PATCH 02/11] refcount: Implement inc/decrement-and-return functions

[David Howells]

Implement functions that increment or decrement a refcount_t object and
return the value.  The dec-and-ret function can be used to maintain a
counter in a cache where 1 means the object is unused, but available and
the garbage collector can use refcount_dec_if_one() to make the object
unavailable.  Further, both functions can be used to accurately trace the
refcount (refcount_inc() followed by refcount_read() can't be considered
accurate).

The interface is as follows:

	unsigned int refcount_dec_return(refcount_t *r);
	unsigned int refcount_inc_return(refcount_t *r);

instead.

Signed-off-by: David Howells <dhowells@xxxxxxxxxx>
cc: Peter Zijlstra <peterz@xxxxxxxxxxxxx>
cc: Kees Cook <keescook@xxxxxxxxxxxx>
---

 include/linux/refcount.h |   12 ++++++++
 lib/refcount.c           |   67 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/include/linux/refcount.h b/include/linux/refcount.h
index 591792c8e5b0..566c0cea7343 100644
--- a/include/linux/refcount.h
+++ b/include/linux/refcount.h
@@ -52,6 +52,8 @@ extern __must_check bool refcount_sub_and_test(unsigned int i, refcount_t *r);
 
 extern __must_check bool refcount_dec_and_test(refcount_t *r);
 extern void refcount_dec(refcount_t *r);
+extern __must_check unsigned int refcount_inc_return(refcount_t *r);
+extern __must_check unsigned int refcount_dec_return(refcount_t *r);
 #else
 static inline __must_check bool refcount_add_not_zero(unsigned int i, refcount_t *r)
 {
@@ -87,6 +89,16 @@ static inline void refcount_dec(refcount_t *r)
 {
 	atomic_dec(&r->refs);
 }
+
+static inline unsigned int refcount_inc_return(refcount_t *r)
+{
+	return atomic_inc_return(&r->refs);
+}
+
+static inline unsigned int refcount_dec_return(refcount_t *r)
+{
+	return atomic_dec_return(&r->refs);
+}
 #endif /* CONFIG_REFCOUNT_FULL */
 
 extern __must_check bool refcount_dec_if_one(refcount_t *r);
diff --git a/lib/refcount.c b/lib/refcount.c
index 5d0582a9480c..3a1d800bf830 100644
--- a/lib/refcount.c
+++ b/lib/refcount.c
@@ -154,6 +154,40 @@ void refcount_inc(refcount_t *r)
 EXPORT_SYMBOL(refcount_inc);
 
 /**
+ * refcount_inc_return - increment a refcount and return the new value
+ * @r: the refcount to increment
+ *
+ * Similar to atomic_inc_return(), but will saturate at UINT_MAX and WARN.
+ *
+ * Provides no memory ordering, it is assumed the caller has guaranteed the
+ * object memory to be stable (RCU, etc.). It does provide a control dependency
+ * and thereby orders future stores. See the comment on top.
+ *
+ * Return: the new value.
+ */
+unsigned int refcount_inc_return(refcount_t *r)
+{
+	unsigned int new, val = atomic_read(&r->refs);
+
+	do {
+		new = val + 1;
+
+		if (!val) {
+			WARN_ONCE(!val, "refcount_t: increment on 0; use-after-free.\n");
+			return 0;
+		}
+
+		if (unlikely(!new))
+			return UINT_MAX;
+
+	} while (!atomic_try_cmpxchg_relaxed(&r->refs, &val, new));
+
+	WARN_ONCE(new == UINT_MAX, "refcount_t: saturated; leaking memory.\n");
+	return new;
+}
+EXPORT_SYMBOL(refcount_inc_return);
+
+/**
  * refcount_sub_and_test - subtract from a refcount and test if it is 0
  * @i: amount to subtract from the refcount
  * @r: the refcount
@@ -227,6 +261,39 @@ void refcount_dec(refcount_t *r)
 	WARN_ONCE(refcount_dec_and_test(r), "refcount_t: decrement hit 0; leaking memory.\n");
 }
 EXPORT_SYMBOL(refcount_dec);
+
+/**
+ * refcount_dec_return - Decrement a refcount and return the new value.
+ * @r: the refcount
+ *
+ * Similar to atomic_dec_return(), it will WARN on underflow and fail to
+ * decrement when saturated at UINT_MAX.  It isn't permitted to use this to
+ * decrement a counter to 0.
+ *
+ * Provides release memory ordering, such that prior loads and stores are done
+ * before.
+ */
+unsigned int refcount_dec_return(refcount_t *r)
+{
+	unsigned int new, val = atomic_read(&r->refs);
+
+	do {
+		if (unlikely(val == UINT_MAX))
+			return val;
+
+		new = val - 1;
+		if (unlikely(val == 0)) {
+			WARN_ONCE(val == 0, "refcount_t: underflow; use-after-free.\n");
+			return val;
+		}
+
+		WARN_ONCE(val == 1, "refcount_t: decrement hit 0; leaking memory.\n");
+
+	} while (!atomic_try_cmpxchg_release(&r->refs, &val, new));
+
+	return new;
+}
+EXPORT_SYMBOL(refcount_dec);
 #endif /* CONFIG_REFCOUNT_FULL */
 
 /**