Merge pull request #1287 from trws/dense_resource_type

Dense resource type with DoS protection

resource/libjobspec/jobspec.cpp

@@ -126,7 +126,13 @@ Resource::Resource (const YAML::Node &resnode)
     if (!resnode["type"].IsScalar ()) {
         throw parse_error (resnode["type"], "Value of \"type\" must be a scalar");
     }
-    type = resource_model::resource_type_t{resnode["type"].as<std::string> ()};
+    try {
+        type = resource_model::resource_type_t{resnode["type"].as<std::string> ()};
+    } catch (std::system_error e) {
+        // resource_type must be closed or full
+        throw parse_error (resnode["type"],
+                           "Value of \"type\" must be a resource type known to fluxion");
+    }
     field_count++;
 
     if (!resnode["count"]) {

resource/modules/resource_match.cpp

@@ -1083,6 +1083,7 @@ static int grow_resource_db (std::shared_ptr<resource_ctx_t> &ctx, json_t *resou
     struct idset *grow_set = NULL;
     json_t *r_lite = NULL;
     json_t *jgf = NULL;
+    auto guard = resource_type_t::storage_t::open_for_scope ();
 
     if ((rc = unpack_resources (resources, &grow_set, &r_lite, &jgf, duration)) < 0) {
         flux_log_error (ctx->h, "%s: unpack_resources", __FUNCTION__);
@@ -1461,6 +1462,10 @@ static int init_resource_graph (std::shared_ptr<resource_ctx_t> &ctx)
             }
         }
     }
+
+    // prevent users from consuming unbounded memory with arbitrary resource types
+    subsystem_t::storage_t::finalize ();
+    resource_type_t::storage_t::finalize ();
     return 0;
 }
 

resource/schema/data_std.cpp

@@ -17,6 +17,7 @@ subsystem_t containment_sub{"containment"};
 
 resource_type_t cluster_rt{"cluster"};
 resource_type_t core_rt{"core"};
+resource_type_t socket_rt{"socket"};
 resource_type_t gpu_rt{"gpu"};
 resource_type_t node_rt{"node"};
 resource_type_t rack_rt{"rack"};

resource/schema/data_std.hpp

@@ -36,11 +36,12 @@ extern subsystem_t containment_sub;
 
 constexpr uint64_t resource_type_id{1};
 struct resource_type_tag {};
-using resource_type_t = intern::interned_string<intern::rc_storage<resource_type_tag>>;
+using resource_type_t = intern::interned_string<intern::dense_storage<resource_type_tag, uint16_t>>;
 extern resource_type_t slot_rt;
 extern resource_type_t cluster_rt;
 extern resource_type_t rack_rt;
 extern resource_type_t node_rt;
+extern resource_type_t socket_rt;
 extern resource_type_t gpu_rt;
 extern resource_type_t core_rt;
 

resource/utilities/command.cpp

@@ -429,13 +429,18 @@ static int update_run (std::shared_ptr<resource_context_t> &ctx,
     }
 
     gettimeofday (&st, NULL);
-    if ((rc = ctx->traverser->run (str, ctx->writers, rd, id, at, d)) != 0) {
-        std::cerr << "ERROR: traverser run () returned error " << std::endl;
-        if (ctx->traverser->err_message () != "") {
-            std::cerr << "ERROR: " << ctx->traverser->err_message ();
-            ctx->traverser->clear_err_message ();
+
+    {
+        auto guard = resource_type_t::storage_t::open_for_scope ();
+        if ((rc = ctx->traverser->run (str, ctx->writers, rd, id, at, d)) != 0) {
+            std::cerr << "ERROR: traverser run () returned error " << std::endl;
+            if (ctx->traverser->err_message () != "") {
+                std::cerr << "ERROR: " << ctx->traverser->err_message ();
+                ctx->traverser->clear_err_message ();
+            }
         }
     }
+
     ctx->writers->emit (o);
     std::ostream &out = (ctx->params.r_fname != "") ? ctx->params.r_out : std::cout;
     out << o.str ();

resource/utilities/resource-query.cpp

@@ -587,6 +587,9 @@ static int init_resource_graph (std::shared_ptr<resource_context_t> &ctx)
         return -1;
     }
 
+    // prevent users from consuming unbounded memory with arbitrary resource types
+    subsystem_t::storage_t::finalize ();
+    resource_type_t::storage_t::finalize ();
     return rc;
 }
 

src/common/libintern/interner.cpp

@@ -14,6 +14,7 @@
 #include <memory>
 #include <mutex>
 #include <shared_mutex>
+#include <thread>
 #include <tuple>
 #include <unordered_map>
 #include <utility>
@@ -47,6 +48,10 @@ struct dense_inner_storage {
     std::unordered_map<std::string, size_t, string_hash, std::equal_to<>> ids_by_string;
     // reader/writer lock to protect entries
     std::unique_ptr<std::shared_mutex> mtx = std::make_unique<std::shared_mutex> ();
+    // if the object is finalized and not opened, reject new strings
+    bool finalized = false;
+    // whether this object is allowed to accept new strings after finalization from each thread
+    std::unordered_map<std::thread::id, bool> open_map;
 };
 
 dense_inner_storage &get_dense_inner_storage (size_t unique_id)
@@ -56,6 +61,22 @@ dense_inner_storage &get_dense_inner_storage (size_t unique_id)
     return groups[unique_id];
 }
 
+void dense_storage_finalize (dense_inner_storage &storage)
+{
+    auto ul = std::unique_lock (*storage.mtx);
+    storage.finalized = true;
+}
+void dense_storage_open (dense_inner_storage &storage)
+{
+    auto ul = std::unique_lock (*storage.mtx);
+    storage.open_map[std::this_thread::get_id ()] = true;
+}
+void dense_storage_close (dense_inner_storage &storage)
+{
+    auto ul = std::unique_lock (*storage.mtx);
+    storage.open_map[std::this_thread::get_id ()] = false;
+}
+
 struct sparse_inner_storage;
 struct sparse_string_node {
     // This must be carefully managed, it holds onto the sparse_inner_storage that is keeping the
@@ -153,6 +174,15 @@ view_and_id get_both (dense_inner_storage &storage, const std::string_view s, ch
             throw std::system_error (ENOMEM,
                                      std::generic_category (),
                                      "Too many strings for configured size");
+        // if storage is finalized and the thread isn't in the open map, we must not add another
+        // here, check while under the shared lock
+        if (storage.finalized && !storage.open_map[std::this_thread::get_id ()]) {
+            using namespace std::string_literals;
+            std::string err =
+                "This interner is finalized and must be open to add strings, found new string: '"s
+                + std::string (s) + "'"s;
+            throw std::system_error (ENOMEM, std::generic_category (), err);
+        }
     }  // release shared lock
 
     // writer lock scope
@@ -167,6 +197,7 @@ view_and_id get_both (dense_inner_storage &storage, const std::string_view s, ch
 
 const std::string *get_by_id (dense_inner_storage &storage, size_t string_id)
 {
+    auto sl = std::shared_lock (*storage.mtx);
     return storage.strings_by_id.at (string_id);
 }
 

src/common/libintern/interner.hpp

@@ -20,6 +20,7 @@
 #include <ostream>
 #include <math.h>
 
+#include "scope_guard.hpp"
 #include <boost/container/small_vector.hpp>
 #include <boost/optional/optional.hpp>
 #include <stdexcept>
@@ -36,12 +37,20 @@ struct dense_inner_storage;
 view_and_id get_both (dense_inner_storage &ds, std::string_view s, char bytes_supported);
 const std::string *get_by_id (dense_inner_storage &ds, size_t string_id);
 
+void dense_storage_finalize (dense_inner_storage &storage);
+void dense_storage_open (dense_inner_storage &storage);
+void dense_storage_close (dense_inner_storage &storage);
+
 dense_inner_storage &get_dense_inner_storage (size_t unique_id);
 
 std::size_t hash_combine (std::size_t lhs, std::size_t rhs);
 };  // namespace detail
 
 /// interner storage class providing dense, in-order IDs of configurable size
+///
+/// allows addition of strings until finalized, then after that only when explicitly opened for
+/// additions, this is not for thread safety, it is to prevent denial of service from addition of
+/// invalid types through interfaces that take user input
 template<class Tag, class Id>
     requires (sizeof (Id) <= sizeof (size_t))
 struct dense_storage {
@@ -55,6 +64,18 @@ struct dense_storage {
     dense_storage &operator= (dense_storage const &) = delete;
     dense_storage &operator= (dense_storage &&) = delete;
 
+   private:
+    static detail::dense_inner_storage &get_storage ()
+    {
+        // C++ guarantees atomic initialization of a static like this, template expansion makes
+        // __PRETTY_FUNCTION__ unique for valid programs, using function here to avoid having to
+        // include heavy headers in this header
+        static detail::dense_inner_storage &s =
+            ::intern::detail::get_dense_inner_storage (dense_storage::unique_id ());
+        return s;
+    };
+
+   public:
     static detail::view_and_id get_both (std::string_view s);
     static id_instance_t get_id (std::string_view s);
     static const std::string *get_by_id (id_storage_t string_id);
@@ -65,16 +86,28 @@ struct dense_storage {
         return uid;
     }
 
-   private:
-    static detail::dense_inner_storage &get_storage ()
+    /// stop new strings from being added unless the object is explicitly opened
+    static void finalize ()
     {
-        // C++ guarantees atomic initialization of a static like this, template expansion makes
-        // __PRETTY_FUNCTION__ unique for valid programs, using function here to avoid having to
-        // include heavy headers in this header
-        static detail::dense_inner_storage &s =
-            ::intern::detail::get_dense_inner_storage (dense_storage::unique_id ());
-        return s;
-    };
+        detail::dense_storage_finalize (get_storage ());
+    }
+    /// open the interner for additions and auto-close on scope exit
+    [[nodiscard]] static auto open_for_scope ()
+    {
+        open ();
+        return sg::make_scope_guard ([] () { close (); });
+    }
+
+    /// open the interner for additions
+    static auto open ()
+    {
+        detail::dense_storage_open (get_storage ());
+    }
+    /// close the interner for additions
+    static void close ()
+    {
+        detail::dense_storage_close (get_storage ());
+    }
 };
 
 template<class Tag, class Id>