Make use of the alloc_size attribute

2 files changed, 30 insertions, 0 deletions

diff --git a/include/fortify-headers.h b/include/fortify-headers.h
index 2f81fc0..dde19a2 100644
--- a/include/fortify-headers.h
+++ b/include/fortify-headers.h
@@ -78,6 +78,13 @@
 #define __malloc(...)
 #endif
 
+#if __has_attribute (alloc_size)
+#define __alloc_size(...) __attribute__ ((alloc_size (__VA_ARGS__)))
+#else
+#define __alloc_size(...)
+#endif
+
+
 #endif /* __has_attribute */
 
 
diff --git a/include/stdlib.h b/include/stdlib.h
index b661862..6914f02 100644
--- a/include/stdlib.h
+++ b/include/stdlib.h
@@ -37,13 +37,36 @@ extern "C" {
 #endif
 
 #undef malloc
+#undef realloc
 
 __malloc(malloc (free, 1))
+__alloc_size(1)
 _FORTIFY_FN(malloc) void *malloc(size_t __s)
 {
         return __orig_malloc(__s);
 }
 
+__alloc_size(2)
+_FORTIFY_FN(realloc) void *realloc(void *__p, size_t __s)
+{
+        return __orig_realloc(__p, __s);
+}
+
+__alloc_size(1, 2)
+_FORTIFY_FN(calloc) void *calloc(size_t __n, size_t __s)
+{
+        return __orig_calloc(__n, __s);
+}
+
+#if defined(_GNU_SOURCE) || defined(_BSD_SOURCE)
+#undef reallocarray
+__alloc_size (2, 3)
+_FORTIFY_FN(reallocarray) void* reallocarray(void* __p, size_t __n, size_t __s)
+{
+        return __orig_reallocarray(__p, __n, __s);
+}
+#endif
+
 /* FIXME clang */
 #if (defined(_XOPEN_SOURCE) || defined(_GNU_SOURCE) || defined(_BSD_SOURCE)) && !defined(__clang__)
 #undef realpath