From a1cb2c371a72af9556a629fb4f89acee92626708 Mon Sep 17 00:00:00 2001 From: Ali Mohammad Pur Date: Mon, 14 Feb 2022 16:49:53 +0330 Subject: [PATCH] AK+Kernel: OOM-harden most parts of Trie The only part of Unveil that can't handle OOM gracefully is the String::formatted() use in the node metadata. --- AK/Trie.h | 171 +++++++++++++++++------------- Kernel/FileSystem/UnveilNode.h | 24 ++++- Kernel/Process.cpp | 5 +- Kernel/Process.h | 4 +- Kernel/ProcessSpecificExposed.cpp | 6 +- Kernel/Syscalls/execve.cpp | 2 +- Kernel/Syscalls/fork.cpp | 2 +- Kernel/Syscalls/unveil.cpp | 17 +-- Tests/AK/TestTrie.cpp | 13 +-- 9 files changed, 145 insertions(+), 99 deletions(-) diff --git a/AK/Trie.h b/AK/Trie.h index 6436d7375c..8c60e2f724 100644 --- a/AK/Trie.h +++ b/AK/Trie.h @@ -6,6 +6,7 @@ #pragma once +#include #include #include #include @@ -29,62 +30,6 @@ template::Type; - class ConstIterator { - - public: - static ConstIterator end() { return {}; } - - bool operator==(const ConstIterator& other) const { return m_current_node == other.m_current_node; } - - const BaseType& operator*() const { return static_cast(*m_current_node); } - const BaseType* operator->() const { return static_cast(m_current_node); } - void operator++() { skip_to_next(); } - - explicit ConstIterator(const Trie& node) - { - m_current_node = &node; - // FIXME: Figure out how to OOM harden this iterator. - MUST(m_state.try_empend(false, node.m_children.begin(), node.m_children.end())); - } - - private: - void skip_to_next() - { - auto& current_state = m_state.last(); - if (current_state.did_generate_root) - ++current_state.it; - else - current_state.did_generate_root = true; - if (current_state.it == current_state.end) - return pop_and_get_next(); - - m_current_node = &*(*current_state.it).value; - - // FIXME: Figure out how to OOM harden this iterator. - MUST(m_state.try_empend(false, m_current_node->m_children.begin(), m_current_node->m_children.end())); - } - void pop_and_get_next() - { - m_state.take_last(); - if (m_state.is_empty()) { - m_current_node = nullptr; - return; - } - - skip_to_next(); - } - - ConstIterator() = default; - - struct State { - bool did_generate_root { false }; - typename HashMap, ValueTraits>::ConstIteratorType it; - typename HashMap, ValueTraits>::ConstIteratorType end; - }; - Vector m_state; - const Trie* m_current_node { nullptr }; - }; - public: using MetadataType = MetadataT; @@ -127,48 +72,56 @@ public: Optional metadata() const requires(!IsNullPointer) { return m_metadata; } void set_metadata(MetadataType metadata) requires(!IsNullPointer) { m_metadata = move(metadata); } const MetadataType& metadata_value() const requires(!IsNullPointer) { return m_metadata.value(); } + MetadataType& metadata_value() requires(!IsNullPointer) { return m_metadata.value(); } const ValueType& value() const { return m_value; } ValueType& value() { return m_value; } - Trie& ensure_child(ValueType value, Optional metadata = {}) + ErrorOr ensure_child(ValueType value, Optional metadata = {}) { auto it = m_children.find(value); if (it == m_children.end()) { - auto node = adopt_nonnull_own_or_enomem(new (nothrow) Trie(value, move(metadata))).release_value_but_fixme_should_propagate_errors(); + auto node = TRY(adopt_nonnull_own_or_enomem(new (nothrow) Trie(value, move(metadata)))); auto& node_ref = *node; - m_children.set(move(value), move(node)); - return static_cast(node_ref); + TRY(m_children.try_set(move(value), move(node))); + return &static_cast(node_ref); } auto& node_ref = *it->value; if (metadata.has_value()) node_ref.m_metadata = move(metadata); - return static_cast(node_ref); + return &static_cast(node_ref); } template - BaseType& insert( + ErrorOr insert( It& it, const It& end, MetadataType metadata, ProvideMetadataFunction provide_missing_metadata) requires(!IsNullPointer) { Trie* last_root_node = &traverse_until_last_accessible_node(it, end); + auto invoke_provide_missing_metadata = [&](Ts && ... args)->ErrorOr> + { + if constexpr (SameAs(args)...))>) + return Optional(provide_missing_metadata(forward(args)...)); + else + return provide_missing_metadata(forward(args)...); + }; for (; it != end; ++it) - last_root_node = static_cast(&last_root_node->ensure_child(*it, provide_missing_metadata(static_cast(*last_root_node), it))); + last_root_node = static_cast(TRY(last_root_node->ensure_child(*it, TRY(invoke_provide_missing_metadata(static_cast(*last_root_node), it))))); last_root_node->set_metadata(move(metadata)); - return static_cast(*last_root_node); + return static_cast(last_root_node); } template - BaseType& insert(It& it, const It& end) requires(IsNullPointer) + ErrorOr insert(It& it, const It& end) requires(IsNullPointer) { Trie* last_root_node = &traverse_until_last_accessible_node(it, end); for (; it != end; ++it) - last_root_node = static_cast(&last_root_node->ensure_child(*it, {})); - return static_cast(*last_root_node); + last_root_node = static_cast(TRY(last_root_node->ensure_child(*it, {}))); + return static_cast(last_root_node); } template - BaseType& insert( + ErrorOr insert( const It& begin, const It& end, MetadataType metadata, ProvideMetadataFunction provide_missing_metadata) requires(!IsNullPointer) { auto it = begin; @@ -176,7 +129,7 @@ public: } template - BaseType& insert(const It& begin, const It& end) requires(IsNullPointer) + ErrorOr insert(const It& begin, const It& end) requires(IsNullPointer) { auto it = begin; return insert(it, end); @@ -185,21 +138,91 @@ public: HashMap, ValueTraits>& children() { return m_children; } HashMap, ValueTraits> const& children() const { return m_children; } - ConstIterator begin() const { return ConstIterator(*this); } - ConstIterator end() const { return ConstIterator::end(); } + template + ErrorOr for_each_node_in_tree_order(Fn callback) const + { + struct State { + bool did_generate_root { false }; + typename HashMap, ValueTraits>::ConstIteratorType it; + typename HashMap, ValueTraits>::ConstIteratorType end; + }; + Vector state; + TRY(state.try_empend(false, m_children.begin(), m_children.end())); + + auto invoke = [&](auto& current_node) -> ErrorOr { + if constexpr (VoidFunction) { + callback(static_cast(current_node)); + return IterationDecision::Continue; + } else if constexpr (IsSpecializationOf())), ErrorOr>) { + return callback(static_cast(current_node)); + } else if constexpr (IteratorFunction) { + return callback(static_cast(current_node)); + } else { + static_assert(DependentFalse, "Invalid iterator function type signature"); + } + return IterationDecision::Continue; + }; + + for (auto* current_node = this; current_node != nullptr;) { + if (TRY(invoke(*current_node)) == IterationDecision::Break) + break; + TRY(skip_to_next_iterator(state, current_node)); + } + return {}; + } [[nodiscard]] bool is_empty() const { return m_children.is_empty(); } void clear() { m_children.clear(); } - BaseType deep_copy() + ErrorOr deep_copy() { - Trie root(m_value, m_metadata); + Trie root(m_value, TRY(copy_metadata(m_metadata))); for (auto& it : m_children) - root.m_children.set(it.key, adopt_nonnull_own_or_enomem(new (nothrow) Trie(it.value->deep_copy())).release_value_but_fixme_should_propagate_errors()); + TRY(root.m_children.try_set(it.key, TRY(adopt_nonnull_own_or_enomem(new (nothrow) Trie(TRY(it.value->deep_copy())))))); return static_cast(move(root)); } private: + static ErrorOr> copy_metadata(Optional const& metadata) + { + if (!metadata.has_value()) + return Optional {}; + + if constexpr (requires(MetadataType t) { { t.copy() } -> SpecializationOf; }) + return Optional { TRY(metadata->copy()) }; +#ifndef KERNEL + else + return Optional { MetadataType(metadata.value()) }; +#endif + } + + static ErrorOr skip_to_next_iterator(auto& state, auto& current_node) + { + auto& current_state = state.last(); + if (current_state.did_generate_root) + ++current_state.it; + else + current_state.did_generate_root = true; + + if (current_state.it == current_state.end) + return pop_and_get_next_iterator(state, current_node); + + current_node = &*(*current_state.it).value; + + TRY(state.try_empend(false, current_node->m_children.begin(), current_node->m_children.end())); + return {}; + } + + static ErrorOr pop_and_get_next_iterator(auto& state, auto& current_node) + { + state.take_last(); + if (state.is_empty()) { + current_node = nullptr; + return {}; + } + return skip_to_next_iterator(state, current_node); + } + ValueType m_value; Optional m_metadata; HashMap, ValueTraits> m_children; diff --git a/Kernel/FileSystem/UnveilNode.h b/Kernel/FileSystem/UnveilNode.h index 7590d9c759..76d3ec63bb 100644 --- a/Kernel/FileSystem/UnveilNode.h +++ b/Kernel/FileSystem/UnveilNode.h @@ -24,9 +24,29 @@ enum UnveilAccess { struct UnveilNode; struct UnveilMetadata { - String full_path; + NonnullOwnPtr full_path; UnveilAccess permissions { None }; bool explicitly_unveiled { false }; + + UnveilMetadata(UnveilMetadata const&) = delete; + UnveilMetadata(UnveilMetadata&&) = default; + + // Note: Intentionally not explicit. + UnveilMetadata(NonnullOwnPtr&& full_path, UnveilAccess permissions = None, bool explicitly_unveiled = false) + : full_path(move(full_path)) + , permissions(permissions) + , explicitly_unveiled(explicitly_unveiled) + { + } + + ErrorOr copy() const + { + return UnveilMetadata { + TRY(full_path->try_clone()), + permissions, + explicitly_unveiled, + }; + } }; struct UnveilNode final : public Trie, UnveilNode> { @@ -34,7 +54,7 @@ struct UnveilNode final : public Trie, Un bool was_explicitly_unveiled() const { return this->metadata_value().explicitly_unveiled; } UnveilAccess permissions() const { return this->metadata_value().permissions; } - const String& path() const { return this->metadata_value().full_path; } + StringView path() const { return this->metadata_value().full_path->view(); } }; } diff --git a/Kernel/Process.cpp b/Kernel/Process.cpp index 13e6076a86..fb30fea499 100644 --- a/Kernel/Process.cpp +++ b/Kernel/Process.cpp @@ -218,17 +218,18 @@ void Process::unprotect_data() ErrorOr> Process::try_create(RefPtr& first_thread, NonnullOwnPtr name, UserID uid, GroupID gid, ProcessID ppid, bool is_kernel_process, RefPtr cwd, RefPtr executable, TTY* tty, Process* fork_parent) { auto space = TRY(Memory::AddressSpace::try_create(fork_parent ? &fork_parent->address_space() : nullptr)); - auto process = TRY(adopt_nonnull_ref_or_enomem(new (nothrow) Process(move(name), uid, gid, ppid, is_kernel_process, move(cwd), move(executable), tty))); + auto process = TRY(adopt_nonnull_ref_or_enomem(new (nothrow) Process(move(name), uid, gid, ppid, is_kernel_process, move(cwd), move(executable), tty, UnveilNode { "/"sv, UnveilMetadata(TRY(KString::try_create("/"sv))) }))); TRY(process->attach_resources(move(space), first_thread, fork_parent)); return process; } -Process::Process(NonnullOwnPtr name, UserID uid, GroupID gid, ProcessID ppid, bool is_kernel_process, RefPtr cwd, RefPtr executable, TTY* tty) +Process::Process(NonnullOwnPtr name, UserID uid, GroupID gid, ProcessID ppid, bool is_kernel_process, RefPtr cwd, RefPtr executable, TTY* tty, UnveilNode unveil_tree) : m_name(move(name)) , m_is_kernel_process(is_kernel_process) , m_executable(move(executable)) , m_cwd(move(cwd)) , m_tty(tty) + , m_unveiled_paths(move(unveil_tree)) , m_wait_blocker_set(*this) { // Ensure that we protect the process data when exiting the constructor. diff --git a/Kernel/Process.h b/Kernel/Process.h index 4e1c6f974f..1c62102dba 100644 --- a/Kernel/Process.h +++ b/Kernel/Process.h @@ -521,7 +521,7 @@ private: bool add_thread(Thread&); bool remove_thread(Thread&); - Process(NonnullOwnPtr name, UserID, GroupID, ProcessID ppid, bool is_kernel_process, RefPtr cwd, RefPtr executable, TTY* tty); + Process(NonnullOwnPtr name, UserID, GroupID, ProcessID ppid, bool is_kernel_process, RefPtr cwd, RefPtr executable, TTY* tty, UnveilNode unveil_tree); static ErrorOr> try_create(RefPtr& first_thread, NonnullOwnPtr name, UserID, GroupID, ProcessID ppid, bool is_kernel_process, RefPtr cwd = nullptr, RefPtr executable = nullptr, TTY* = nullptr, Process* fork_parent = nullptr); ErrorOr attach_resources(NonnullOwnPtr&&, RefPtr& first_thread, Process* fork_parent); static ProcessID allocate_pid(); @@ -800,7 +800,7 @@ private: RefPtr m_alarm_timer; VeilState m_veil_state { VeilState::None }; - UnveilNode m_unveiled_paths { "/", { .full_path = "/" } }; + UnveilNode m_unveiled_paths; OwnPtr m_perf_event_buffer; diff --git a/Kernel/ProcessSpecificExposed.cpp b/Kernel/ProcessSpecificExposed.cpp index 5d2169ee40..fec05d7be1 100644 --- a/Kernel/ProcessSpecificExposed.cpp +++ b/Kernel/ProcessSpecificExposed.cpp @@ -145,9 +145,9 @@ ErrorOr Process::procfs_get_pledge_stats(KBufferBuilder& builder) const ErrorOr Process::procfs_get_unveil_stats(KBufferBuilder& builder) const { JsonArraySerializer array { builder }; - for (auto const& unveiled_path : unveiled_paths()) { + TRY(unveiled_paths().for_each_node_in_tree_order([&](auto const& unveiled_path) { if (!unveiled_path.was_explicitly_unveiled()) - continue; + return; auto obj = array.add_object(); obj.add("path", unveiled_path.path()); StringBuilder permissions_builder; @@ -162,7 +162,7 @@ ErrorOr Process::procfs_get_unveil_stats(KBufferBuilder& builder) const if (unveiled_path.permissions() & UnveilAccess::Browse) permissions_builder.append('b'); obj.add("permissions", permissions_builder.string_view()); - } + })); array.finish(); return {}; } diff --git a/Kernel/Syscalls/execve.cpp b/Kernel/Syscalls/execve.cpp index 2b9de45c04..93d2e16dc3 100644 --- a/Kernel/Syscalls/execve.cpp +++ b/Kernel/Syscalls/execve.cpp @@ -525,7 +525,7 @@ ErrorOr Process::do_exec(NonnullRefPtr main_program_d m_veil_state = VeilState::None; m_unveiled_paths.clear(); - m_unveiled_paths.set_metadata({ "/", UnveilAccess::None, false }); + m_unveiled_paths.set_metadata({ TRY(KString::try_create("/"sv)), UnveilAccess::None, false }); for (auto& property : m_coredump_properties) property = {}; diff --git a/Kernel/Syscalls/fork.cpp b/Kernel/Syscalls/fork.cpp index 5848149871..a9890025cd 100644 --- a/Kernel/Syscalls/fork.cpp +++ b/Kernel/Syscalls/fork.cpp @@ -21,7 +21,7 @@ ErrorOr Process::sys$fork(RegisterState& regs) auto child_name = TRY(m_name->try_clone()); auto child = TRY(Process::try_create(child_first_thread, move(child_name), uid(), gid(), pid(), m_is_kernel_process, m_cwd, m_executable, m_tty, this)); child->m_veil_state = m_veil_state; - child->m_unveiled_paths = m_unveiled_paths.deep_copy(); + child->m_unveiled_paths = TRY(m_unveiled_paths.deep_copy()); TRY(child->m_fds.with_exclusive([&](auto& child_fds) { return m_fds.with_exclusive([&](auto& parent_fds) { diff --git a/Kernel/Syscalls/unveil.cpp b/Kernel/Syscalls/unveil.cpp index 196c7fee54..38cd21c23b 100644 --- a/Kernel/Syscalls/unveil.cpp +++ b/Kernel/Syscalls/unveil.cpp @@ -19,7 +19,7 @@ static void update_intermediate_node_permissions(UnveilNode& root_node, UnveilAc auto& node = static_cast(*entry.value); if (node.was_explicitly_unveiled()) continue; - node.set_metadata({ node.path(), new_permissions, node.was_explicitly_unveiled() }); + node.metadata_value().permissions = new_permissions; update_intermediate_node_permissions(node, new_permissions); } } @@ -109,19 +109,20 @@ ErrorOr Process::sys$unveil(Userspace if (matching_node.permissions() != new_permissions) update_intermediate_node_permissions(matching_node, (UnveilAccess)new_permissions); - matching_node.set_metadata({ matching_node.path(), (UnveilAccess)new_permissions, true }); + matching_node.metadata_value().explicitly_unveiled = true; + matching_node.metadata_value().permissions = (UnveilAccess)new_permissions; m_veil_state = VeilState::Dropped; return 0; } - matching_node.insert( + TRY(matching_node.insert( it, path_parts.end(), - { new_unveiled_path->view(), (UnveilAccess)new_permissions, true }, - [](auto& parent, auto& it) -> Optional { - auto path = String::formatted("{}/{}", parent.path(), *it); - return UnveilMetadata { path, parent.permissions(), false }; - }); + { new_unveiled_path.release_nonnull(), (UnveilAccess)new_permissions, true }, + [](auto& parent, auto& it) -> ErrorOr> { + auto path = TRY(KString::formatted("{}/{}", parent.path(), *it)); + return UnveilMetadata(move(path), parent.permissions(), false); + })); VERIFY(m_veil_state != VeilState::Locked); m_veil_state = VeilState::Dropped; diff --git a/Tests/AK/TestTrie.cpp b/Tests/AK/TestTrie.cpp index 9bcda0d940..7334f3a640 100644 --- a/Tests/AK/TestTrie.cpp +++ b/Tests/AK/TestTrie.cpp @@ -16,12 +16,13 @@ TEST_CASE(normal_behavior) constexpr size_t total_chars = 18; // root (1), 'test' (4), 'example' (7), 'foo' (3), 'foobar' (3, "foo" already stored). for (auto& view : data) { auto it = view.begin(); - dictionary.insert(it, view.end(), view, [](auto& parent, auto& it) -> Optional { return String::formatted("{}{}", parent.metadata_value(), *it); }); + MUST(dictionary.insert(it, view.end(), view, [](auto& parent, auto& it) -> Optional { return String::formatted("{}{}", parent.metadata_value(), *it); })); } size_t i = 0; - for ([[maybe_unused]] auto& node : dictionary) + MUST(dictionary.for_each_node_in_tree_order([&](auto&) { ++i; + })); EXPECT_EQ(i, total_chars); for (auto& view : data) { @@ -49,18 +50,18 @@ TEST_CASE(iterate) for (size_t i = 0; i < input.size(); ++i) input[i] = i; - bunch_of_numbers.insert(input.begin(), input.end()); + MUST(bunch_of_numbers.insert(input.begin(), input.end())); // Iteration order is preorder (order between adjacent nodes is not defined, but parents come before children) // in this case, the tree is linear. size_t i = 0; bool is_root = true; - for (auto& node : bunch_of_numbers) { + MUST(bunch_of_numbers.for_each_node_in_tree_order([&](auto& node) { if (is_root) { is_root = false; - continue; + return; } EXPECT_EQ(input[i], node.value()); ++i; - } + })); }