introduced checked variant of strdup()

wjakob · Nov 1, 2023 · 176e4f7 · 176e4f7
1 parent 3c27ce8
commit 176e4f7
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 8 deletions.
diff --git a/include/nanobind/nb_defs.h b/include/nanobind/nb_defs.h
@@ -28,7 +28,6 @@
 #  define NB_INLINE        __forceinline
 #  define NB_INLINE_LAMBDA
 #  define NB_NOINLINE      __declspec(noinline)
-# define  NB_STRDUP        _strdup
 #else
 #  define NB_EXPORT        __attribute__ ((visibility("default")))
 #  define NB_IMPORT        NB_EXPORT
@@ -39,7 +38,6 @@
 #else
 #    define NB_INLINE_LAMBDA
 #endif
-#  define NB_STRDUP        strdup
 #endif
 
 #if defined(__GNUC__)

diff --git a/src/error.cpp b/src/error.cpp
@@ -46,7 +46,7 @@ python_error::python_error(const python_error &e)
         Py_INCREF(m_value);
     }
     if (e.m_what)
-        m_what = NB_STRDUP(e.m_what);
+        m_what = detail::strdup_check(e.m_what);
 }
 
 python_error::python_error(python_error &&e) noexcept
@@ -95,7 +95,7 @@ python_error::python_error(const python_error &e)
         Py_XINCREF(m_traceback);
     }
     if (e.m_what)
-        m_what = NB_STRDUP(e.m_what);
+        m_what = detail::strdup_check(e.m_what);
 }
 
 python_error::python_error(python_error &&e) noexcept
@@ -147,7 +147,7 @@ const char *python_error::what() const noexcept {
 #if defined(Py_LIMITED_API) || defined(PYPY_VERSION)
     object mod = module_::import_("traceback"),
            result = mod.attr("format_exception")(exc_type, exc_value, exc_traceback);
-    m_what = NB_STRDUP(borrow<str>(str("\n").attr("join")(result)).c_str());
+    m_what = detail::strdup_check(borrow<str>(str("\n").attr("join")(result)).c_str());
 #else
     buf.clear();
     if (exc_traceback.is_valid()) {

diff --git a/src/nb_func.cpp b/src/nb_func.cpp
@@ -167,6 +167,18 @@ void *malloc_check(size_t size) {
     return ptr;
 }
 
+char *strdup_check(const char *s) {
+    char *result;
+    #if defined(_WIN32)
+        result = _strdup(s);
+    #else
+        result = strdup(s);
+    #endif
+    if (!result)
+        fail("nanobind: strdup() failed!");
+    return result;
+}
+
 /**
  * \brief Wrap a C++ function into a Python function object
  *
@@ -1108,7 +1120,7 @@ NB_NOINLINE char *type_name(const std::type_info *t) {
     int status = 0;
     char *name = abi::__cxa_demangle(name_in, nullptr, nullptr, &status);
 #else
-    char *name = NB_STRDUP(name_in);
+    char *name = strdup_check(name_in);
     strexc(name, "class ");
     strexc(name, "struct ");
     strexc(name, "enum ");

diff --git a/src/nb_internals.h b/src/nb_internals.h
@@ -320,5 +320,7 @@ template <typename T> struct scoped_pymalloc {
     T *ptr{ nullptr };
 };
 
+extern char *strdup_check(const char *);
+
 NAMESPACE_END(detail)
 NAMESPACE_END(NB_NAMESPACE)
diff --git a/src/nb_type.cpp b/src/nb_type.cpp
@@ -388,7 +388,7 @@ static int nb_type_init(PyObject *self, PyObject *args, PyObject *kwds) {
     t->flags |=  (uint32_t) type_flags::is_python_type;
     t->flags &= ~((uint32_t) type_flags::has_implicit_conversions);
     PyObject *name = nb_type_name(self);
-    t->name = NB_STRDUP(PyUnicode_AsUTF8AndSize(name, nullptr));
+    t->name = strdup_check(PyUnicode_AsUTF8AndSize(name, nullptr));
     Py_DECREF(name);
     t->type_py = (PyTypeObject *) self;
     t->implicit = nullptr;
@@ -846,7 +846,7 @@ PyObject *nb_type_new(const type_init_data *t) noexcept {
     bool base_intrusive_ptr =
         tb && (tb->flags & (uint32_t) type_flags::intrusive_ptr);
 
-    char *name_copy = NB_STRDUP(name.c_str());
+    char *name_copy = strdup_check(name.c_str());
 
     constexpr size_t nb_type_max_slots = 10,
                      nb_extra_slots = 80,