[PATCH] include: detect use-after-free errors using the reference counts

[Justus Winter]

* include/refcount.h (refcount_init): There must be at least one
reference at initialization time.
(refcounts_init): Likewise.
(refcount_unsafe_ref): New function retaining the previous
functionality of refcount_ref.  It is occasionally useful to raise the
reference count again after it dropped to zero.
(refcounts_unsafe_ref): Likewise.
(refcounts_unsafe_weak_ref): Likewise.
(refcount_ref): Detect use-after-free errors.
(refcounts_ref): Likewise.
(refcounts_ref_weak): Likewise.
* libtrivfs/protid-clean.c (trivfs_clean_protid): Use refcount_unsafe_ref.
---
 include/refcount.h       | 73 ++++++++++++++++++++++++++++++++++++++++++------
 libtrivfs/protid-clean.c |  2 +-
 2 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/include/refcount.h b/include/refcount.h
index 785b052..ebde42d 100644
--- a/include/refcount.h
+++ b/include/refcount.h
@@ -31,18 +31,23 @@
 /* An opaque type.  You must not access these values directly.  */
 typedef unsigned int refcount_t;
 
-/* Initialize REF with REFERENCES.  */
+/* Initialize REF with REFERENCES.  REFERENCES must not be zero.  */
 static inline void
 refcount_init (refcount_t *ref, unsigned int references)
 {
+  assert (references > 0 || !"references must not be zero!");
   *ref = references;
 }
 
 /* Increment REF.  Return the result of the operation.  This function
    uses atomic operations.  It is not required to serialize calls to
-   this function.  */
+   this function.
+
+   This is the unsafe version of refcount_ref.  refcount_ref also
+   checks for use-after-free errors.  When in doubt, use that one
+   instead.  */
 static inline unsigned int
-refcount_ref (refcount_t *ref)
+refcount_unsafe_ref (refcount_t *ref)
 {
   unsigned int r;
   r = __atomic_add_fetch (ref, 1, __ATOMIC_RELAXED);
@@ -50,6 +55,18 @@ refcount_ref (refcount_t *ref)
   return r;
 }
 
+/* Increment REF.  Return the result of the operation.  This function
+   uses atomic operations.  It is not required to serialize calls to
+   this function.  */
+static inline unsigned int
+refcount_ref (refcount_t *ref)
+{
+  unsigned int r;
+  r = refcount_unsafe_ref (ref);
+  assert (r != 1 || !"refcount detected use-after-free!");
+  return r;
+}
+
 /* Decrement REF.  Return the result of the operation.  This function
    uses atomic operations.  It is not required to serialize calls to
    this function.  */
@@ -101,19 +118,25 @@ union _references {
   uint64_t value;
 };
 
-/* Initialize REF with HARD and WEAK references.  */
+/* Initialize REF with HARD and WEAK references.  HARD and WEAK must
+   not both be zero.  */
 static inline void
 refcounts_init (refcounts_t *ref, uint32_t hard, uint32_t weak)
 {
+  assert ((hard != 0 || weak != 0) || !"references must not both be zero!");
   ref->references = (struct references) { .hard = hard, .weak = weak };
 }
 
 /* Increment the hard reference count of REF.  If RESULT is not NULL,
    the result of the operation is written there.  This function uses
    atomic operations.  It is not required to serialize calls to this
-   function.  */
+   function.
+
+   This is the unsafe version of refcounts_ref.  refcounts_ref also
+   checks for use-after-free errors.  When in doubt, use that one
+   instead.  */
 static inline void
-refcounts_ref (refcounts_t *ref, struct references *result)
+refcounts_unsafe_ref (refcounts_t *ref, struct references *result)
 {
   const union _references op = { .references = { .hard = 1 } };
   union _references r;
@@ -123,6 +146,21 @@ refcounts_ref (refcounts_t *ref, struct references *result)
     *result = r.references;
 }
 
+/* Increment the hard reference count of REF.  If RESULT is not NULL,
+   the result of the operation is written there.  This function uses
+   atomic operations.  It is not required to serialize calls to this
+   function.  */
+static inline void
+refcounts_ref (refcounts_t *ref, struct references *result)
+{
+  struct references r;
+  refcounts_unsafe_ref (ref, &r);
+  assert (! (r.hard == 1 && r.weak == 0)
+          || !"refcount detected use-after-free!");
+  if (result)
+    *result = r;
+}
+
 /* Decrement the hard reference count of REF.  If RESULT is not NULL,
    the result of the operation is written there.  This function uses
    atomic operations.  It is not required to serialize calls to this
@@ -200,9 +238,13 @@ refcounts_demote (refcounts_t *ref, struct references 
*result)
 /* Increment the weak reference count of REF.  If RESULT is not NULL,
    the result of the operation is written there.  This function uses
    atomic operations.  It is not required to serialize calls to this
-   function.  */
+   function.
+
+   This is the unsafe version of refcounts_ref_weak.
+   refcounts_ref_weak also checks for use-after-free errors.  When in
+   doubt, use that one instead.  */
 static inline void
-refcounts_ref_weak (refcounts_t *ref, struct references *result)
+refcounts_unsafe_ref_weak (refcounts_t *ref, struct references *result)
 {
   const union _references op = { .references = { .weak = 1 } };
   union _references r;
@@ -212,6 +254,21 @@ refcounts_ref_weak (refcounts_t *ref, struct references 
*result)
     *result = r.references;
 }
 
+/* Increment the weak reference count of REF.  If RESULT is not NULL,
+   the result of the operation is written there.  This function uses
+   atomic operations.  It is not required to serialize calls to this
+   function.  */
+static inline void
+refcounts_ref_weak (refcounts_t *ref, struct references *result)
+{
+  struct references r;
+  refcounts_unsafe_ref_weak (ref, &r);
+  assert (! (r.hard == 0 && r.weak == 1)
+          || !"refcount detected use-after-free!");
+  if (result)
+    *result = r;
+}
+
 /* Decrement the weak reference count of REF.  If RESULT is not NULL,
    the result of the operation is written there.  This function uses
    atomic operations.  It is not required to serialize calls to this
diff --git a/libtrivfs/protid-clean.c b/libtrivfs/protid-clean.c
index adc5e98..ff6cc16 100644
--- a/libtrivfs/protid-clean.c
+++ b/libtrivfs/protid-clean.c
@@ -36,7 +36,7 @@ trivfs_clean_protid (void *arg)
       if (refcount_deref (&cred->po->refcnt) == 0)
         {
           /* Reacquire a reference while we call the hook.  */
-          refcount_ref (&cred->po->refcnt);
+          refcount_unsafe_ref (&cred->po->refcnt);
           (*trivfs_peropen_destroy_hook) (cred->po);
           if (refcount_deref (&cred->po->refcnt) == 0)
             {
-- 
2.1.0