Переглянути джерело

Everywhere: Use OOM-safe ByteBuffer APIs where possible

If we can easily communicate failure, let's avoid asserting and report
failure instead.
Ali Mohammad Pur 3 роки тому
батько
коміт
3a9f00c59b

+ 11 - 7
AK/StringBuilder.cpp

@@ -17,18 +17,18 @@
 
 namespace AK {
 
-inline void StringBuilder::will_append(size_t size)
+inline bool StringBuilder::will_append(size_t size)
 {
     Checked<size_t> needed_capacity = m_buffer.size();
     needed_capacity += size;
     VERIFY(!needed_capacity.has_overflow());
     // Prefer to completely use the existing capacity first
     if (needed_capacity <= m_buffer.capacity())
-        return;
+        return true;
     Checked<size_t> expanded_capacity = needed_capacity;
     expanded_capacity *= 2;
     VERIFY(!expanded_capacity.has_overflow());
-    m_buffer.ensure_capacity(expanded_capacity.value());
+    return m_buffer.try_ensure_capacity(expanded_capacity.value());
 }
 
 StringBuilder::StringBuilder(size_t initial_capacity)
@@ -40,8 +40,10 @@ void StringBuilder::append(StringView const& str)
 {
     if (str.is_empty())
         return;
-    will_append(str.length());
-    m_buffer.append(str.characters_without_null_termination(), str.length());
+    auto ok = will_append(str.length());
+    VERIFY(ok);
+    ok = m_buffer.try_append(str.characters_without_null_termination(), str.length());
+    VERIFY(ok);
 }
 
 void StringBuilder::append(char const* characters, size_t length)
@@ -51,8 +53,10 @@ void StringBuilder::append(char const* characters, size_t length)
 
 void StringBuilder::append(char ch)
 {
-    will_append(1);
-    m_buffer.append(&ch, 1);
+    auto ok = will_append(1);
+    VERIFY(ok);
+    ok = m_buffer.try_append(&ch, 1);
+    VERIFY(ok);
 }
 
 void StringBuilder::appendvf(char const* fmt, va_list ap)

+ 1 - 1
AK/StringBuilder.h

@@ -64,7 +64,7 @@ public:
     }
 
 private:
-    void will_append(size_t);
+    bool will_append(size_t);
     u8* data() { return m_buffer.data(); }
     u8 const* data() const { return m_buffer.data(); }
 

+ 52 - 27
Kernel/Coredump.cpp

@@ -198,13 +198,15 @@ KResult Coredump::write_notes_segment(ByteBuffer& notes_segment)
     return KSuccess;
 }
 
-ByteBuffer Coredump::create_notes_process_data() const
+KResultOr<ByteBuffer> Coredump::create_notes_process_data() const
 {
     ByteBuffer process_data;
 
     ELF::Core::ProcessInfo info {};
     info.header.type = ELF::Core::NotesEntryHeader::Type::ProcessInfo;
-    process_data.append((void*)&info, sizeof(info));
+    auto ok = process_data.try_append((void*)&info, sizeof(info));
+    if (!ok)
+        return ENOMEM;
 
     StringBuilder builder;
     {
@@ -227,18 +229,18 @@ ByteBuffer Coredump::create_notes_process_data() const
     }
 
     builder.append(0);
-    process_data.append(builder.string_view().characters_without_null_termination(), builder.length());
+    ok = process_data.try_append(builder.string_view().characters_without_null_termination(), builder.length());
+    if (!ok)
+        return ENOMEM;
 
     return process_data;
 }
 
-ByteBuffer Coredump::create_notes_threads_data() const
+KResultOr<ByteBuffer> Coredump::create_notes_threads_data() const
 {
     ByteBuffer threads_data;
 
     for (auto& thread : m_process->threads_for_coredump({})) {
-        ByteBuffer entry_buff;
-
         ELF::Core::ThreadInfo info {};
         info.header.type = ELF::Core::NotesEntryHeader::Type::ThreadInfo;
         info.tid = thread.tid().value();
@@ -246,20 +248,17 @@ ByteBuffer Coredump::create_notes_threads_data() const
         if (thread.current_trap())
             copy_kernel_registers_into_ptrace_registers(info.regs, thread.get_register_dump_from_stack());
 
-        entry_buff.append((void*)&info, sizeof(info));
-
-        threads_data += entry_buff;
+        if (!threads_data.try_append(&info, sizeof(info)))
+            return ENOMEM;
     }
     return threads_data;
 }
 
-ByteBuffer Coredump::create_notes_regions_data() const
+KResultOr<ByteBuffer> Coredump::create_notes_regions_data() const
 {
     ByteBuffer regions_data;
     size_t region_index = 0;
     for (auto& region : m_process->address_space().regions()) {
-
-        ByteBuffer memory_region_info_buffer;
         ELF::Core::MemoryRegionInfo info {};
         info.header.type = ELF::Core::NotesEntryHeader::Type::MemoryRegionInfo;
 
@@ -267,28 +266,32 @@ ByteBuffer Coredump::create_notes_regions_data() const
         info.region_end = region->vaddr().offset(region->size()).get();
         info.program_header_index = region_index++;
 
-        memory_region_info_buffer.append((void*)&info, sizeof(info));
+        if (!regions_data.try_ensure_capacity(regions_data.size() + sizeof(info) + region->name().length() + 1))
+            return ENOMEM;
+
+        auto ok = regions_data.try_append((void*)&info, sizeof(info));
         // NOTE: The region name *is* null-terminated, so the following is ok:
         auto name = region->name();
         if (name.is_empty()) {
             char null_terminator = '\0';
-            memory_region_info_buffer.append(&null_terminator, 1);
+            ok = ok && regions_data.try_append(&null_terminator, 1);
         } else {
-            memory_region_info_buffer.append(name.characters_without_null_termination(), name.length() + 1);
+            ok = ok && regions_data.try_append(name.characters_without_null_termination(), name.length() + 1);
         }
-
-        regions_data += memory_region_info_buffer;
+        VERIFY(ok);
     }
     return regions_data;
 }
 
-ByteBuffer Coredump::create_notes_metadata_data() const
+KResultOr<ByteBuffer> Coredump::create_notes_metadata_data() const
 {
     ByteBuffer metadata_data;
 
     ELF::Core::Metadata metadata {};
     metadata.header.type = ELF::Core::NotesEntryHeader::Type::Metadata;
-    metadata_data.append((void*)&metadata, sizeof(metadata));
+    auto ok = metadata_data.try_append((void*)&metadata, sizeof(metadata));
+    if (!ok)
+        return ENOMEM;
 
     StringBuilder builder;
     {
@@ -298,23 +301,41 @@ ByteBuffer Coredump::create_notes_metadata_data() const
         });
     }
     builder.append(0);
-    metadata_data.append(builder.string_view().characters_without_null_termination(), builder.length());
+    ok = metadata_data.try_append(builder.string_view().characters_without_null_termination(), builder.length());
+    if (!ok)
+        return ENOMEM;
 
     return metadata_data;
 }
 
-ByteBuffer Coredump::create_notes_segment_data() const
+KResultOr<ByteBuffer> Coredump::create_notes_segment_data() const
 {
     ByteBuffer notes_buffer;
 
-    notes_buffer += create_notes_process_data();
-    notes_buffer += create_notes_threads_data();
-    notes_buffer += create_notes_regions_data();
-    notes_buffer += create_notes_metadata_data();
+    if (auto result = create_notes_process_data(); result.is_error())
+        return result;
+    else if (!notes_buffer.try_append(result.value()))
+        return ENOMEM;
+
+    if (auto result = create_notes_threads_data(); result.is_error())
+        return result;
+    else if (!notes_buffer.try_append(result.value()))
+        return ENOMEM;
+
+    if (auto result = create_notes_regions_data(); result.is_error())
+        return result;
+    else if (!notes_buffer.try_append(result.value()))
+        return ENOMEM;
+
+    if (auto result = create_notes_metadata_data(); result.is_error())
+        return result;
+    else if (!notes_buffer.try_append(result.value()))
+        return ENOMEM;
 
     ELF::Core::NotesEntryHeader null_entry {};
     null_entry.type = ELF::Core::NotesEntryHeader::Type::Null;
-    notes_buffer.append(&null_entry, sizeof(null_entry));
+    if (!notes_buffer.try_append(&null_entry, sizeof(null_entry)))
+        return ENOMEM;
 
     return notes_buffer;
 }
@@ -324,7 +345,11 @@ KResult Coredump::write()
     SpinlockLocker lock(m_process->address_space().get_lock());
     ProcessPagingScope scope(m_process);
 
-    ByteBuffer notes_segment = create_notes_segment_data();
+    auto notes_segment_result = create_notes_segment_data();
+    if (notes_segment_result.is_error())
+        return notes_segment_result.error();
+
+    auto& notes_segment = notes_segment_result.value();
 
     auto result = write_elf_header();
     if (result.is_error())

+ 5 - 5
Kernel/Coredump.h

@@ -29,11 +29,11 @@ private:
     [[nodiscard]] KResult write_regions();
     [[nodiscard]] KResult write_notes_segment(ByteBuffer&);
 
-    ByteBuffer create_notes_segment_data() const;
-    ByteBuffer create_notes_process_data() const;
-    ByteBuffer create_notes_threads_data() const;
-    ByteBuffer create_notes_regions_data() const;
-    ByteBuffer create_notes_metadata_data() const;
+    KResultOr<ByteBuffer> create_notes_segment_data() const;
+    KResultOr<ByteBuffer> create_notes_process_data() const;
+    KResultOr<ByteBuffer> create_notes_threads_data() const;
+    KResultOr<ByteBuffer> create_notes_regions_data() const;
+    KResultOr<ByteBuffer> create_notes_metadata_data() const;
 
     NonnullRefPtr<Process> m_process;
     NonnullRefPtr<FileDescription> m_fd;

+ 4 - 1
Tests/LibTLS/TestTLSHandshake.cpp

@@ -95,7 +95,10 @@ TEST_CASE(test_TLS_hello_handshake)
             loop.quit(1);
         } else {
             //            print_buffer(data.value(), 16);
-            contents.append(data.value().data(), data.value().size());
+            if (!contents.try_append(data.value().data(), data.value().size())) {
+                FAIL("Allocation failure");
+                loop.quit(1);
+            }
         }
     };
     tls->on_tls_finished = [&] {

+ 2 - 1
Userland/Applications/Debugger/main.cpp

@@ -65,7 +65,8 @@ static bool handle_disassemble_command(const String& command, void* first_instru
         auto value = g_debug_session->peek(reinterpret_cast<u32*>(first_instruction) + i);
         if (!value.has_value())
             break;
-        code.append(&value, sizeof(u32));
+        if (!code.try_append(&value, sizeof(u32)))
+            break;
     }
 
     X86::SimpleInstructionStream stream(code.data(), code.size());

+ 6 - 4
Userland/Libraries/LibC/netdb.cpp

@@ -282,7 +282,7 @@ hostent* gethostbyaddr(const void* addr, socklen_t addr_size, int type)
 
 struct servent* getservent()
 {
-    //If the services file is not open, attempt to open it and return null if it fails.
+    // If the services file is not open, attempt to open it and return null if it fails.
     if (!services_file) {
         services_file = fopen(services_path, "r");
 
@@ -424,7 +424,7 @@ void endservent()
 // false if failure occurs.
 static bool fill_getserv_buffers(const char* line, ssize_t read)
 {
-    //Splitting the line by tab delimiter and filling the servent buffers name, port, and protocol members.
+    // Splitting the line by tab delimiter and filling the servent buffers name, port, and protocol members.
     String string_line = String(line, read);
     string_line.replace(" ", "\t", true);
     auto split_line = string_line.split('\t');
@@ -465,7 +465,8 @@ static bool fill_getserv_buffers(const char* line, ssize_t read)
                 break;
             }
             auto alias = split_line[i].to_byte_buffer();
-            alias.append("\0", sizeof(char));
+            if (!alias.try_append("\0", sizeof(char)))
+                return false;
             __getserv_alias_list_buffer.append(move(alias));
         }
     }
@@ -635,7 +636,8 @@ static bool fill_getproto_buffers(const char* line, ssize_t read)
             if (split_line[i].starts_with('#'))
                 break;
             auto alias = split_line[i].to_byte_buffer();
-            alias.append("\0", sizeof(char));
+            if (!alias.try_append("\0", sizeof(char)))
+                return false;
             __getproto_alias_list_buffer.append(move(alias));
         }
     }

+ 4 - 1
Userland/Libraries/LibCrypto/ASN1/PEM.cpp

@@ -35,7 +35,10 @@ ByteBuffer decode_pem(ReadonlyBytes data)
                 break;
             }
             auto b64decoded = decode_base64(lexer.consume_line().trim_whitespace(TrimMode::Right));
-            decoded.append(b64decoded.data(), b64decoded.size());
+            if (!decoded.try_append(b64decoded.data(), b64decoded.size())) {
+                dbgln("Failed to decode PEM, likely OOM condition");
+                return {};
+            }
             break;
         }
         case Ended:

+ 4 - 1
Userland/Libraries/LibCrypto/PK/Code/EMSA_PSS.h

@@ -151,7 +151,10 @@ public:
         for (size_t counter = 0; counter < length / HashFunction::DigestSize - 1; ++counter) {
             hash_fn.update(seed);
             hash_fn.update((u8*)&counter, 4);
-            T.append(hash_fn.digest().data, HashFunction::DigestSize);
+            if (!T.try_append(hash_fn.digest().data, HashFunction::DigestSize)) {
+                dbgln("EMSA_PSS: MGF1 digest failed, not enough space");
+                return;
+            }
         }
         out.overwrite(0, T.data(), length);
     }

+ 2 - 1
Userland/Libraries/LibGfx/BMPWriter.cpp

@@ -135,7 +135,8 @@ ByteBuffer BMPWriter::dump(const RefPtr<Bitmap> bitmap, DibHeader dib_header)
         }
     }
 
-    buffer.append(pixel_data.data(), pixel_data.size());
+    if (!buffer.try_append(pixel_data.data(), pixel_data.size()))
+        dbgln("Failed to write {} bytes of pixel data to buffer", pixel_data.size());
     return buffer;
 }
 

+ 1 - 3
Userland/Libraries/LibGfx/PNGWriter.cpp

@@ -72,9 +72,7 @@ PNGChunk::PNGChunk(String type)
 
 void PNGChunk::store_type()
 {
-    for (auto character : type()) {
-        m_data.append(&character, sizeof(character));
-    }
+    m_data.append(type().bytes());
 }
 
 void PNGChunk::store_data_length()

+ 2 - 1
Userland/Libraries/LibLine/Editor.cpp

@@ -370,7 +370,8 @@ void Editor::insert(const u32 cp)
     StringBuilder builder;
     builder.append(Utf32View(&cp, 1));
     auto str = builder.build();
-    m_pending_chars.append(str.characters(), str.length());
+    if (!m_pending_chars.try_append(str.characters(), str.length()))
+        return;
 
     readjust_anchored_styles(m_cursor, ModificationKind::Insertion);
 

+ 4 - 2
Userland/Libraries/LibPDF/Parser.cpp

@@ -263,8 +263,10 @@ bool Parser::initialize_hint_tables()
         auto total_size = primary_size + overflow_size;
 
         possible_merged_stream_buffer = ByteBuffer::create_uninitialized(total_size);
-        possible_merged_stream_buffer.append(primary_hint_stream->bytes());
-        possible_merged_stream_buffer.append(overflow_hint_stream->bytes());
+        auto ok = possible_merged_stream_buffer.try_append(primary_hint_stream->bytes());
+        ok = ok && possible_merged_stream_buffer.try_append(overflow_hint_stream->bytes());
+        if (!ok)
+            return false;
         hint_stream_bytes = possible_merged_stream_buffer.bytes();
     } else {
         hint_stream_bytes = primary_hint_stream->bytes();

+ 2 - 1
Userland/Libraries/LibSQL/Heap.cpp

@@ -75,7 +75,8 @@ bool Heap::write_block(u32 block, ByteBuffer& buffer)
     VERIFY(buffer.size() <= BLOCKSIZE);
     auto sz = buffer.size();
     if (sz < BLOCKSIZE) {
-        buffer.resize(BLOCKSIZE);
+        if (!buffer.try_resize(BLOCKSIZE))
+            return false;
         memset(buffer.offset_pointer((int)sz), 0, BLOCKSIZE - sz);
     }
     dbgln_if(SQL_DEBUG, "{:02x} {:02x} {:02x} {:02x} {:02x} {:02x} {:02x} {:02x}",

+ 4 - 1
Userland/Libraries/LibTLS/HandshakeClient.cpp

@@ -119,7 +119,10 @@ bool TLSv12::compute_master_secret_from_pre_master_secret(size_t length)
         return false;
     }
 
-    m_context.master_key.resize(length);
+    if (!m_context.master_key.try_resize(length)) {
+        dbgln("Couldn't allocate enough space for the master key :(");
+        return false;
+    }
 
     pseudorandom_function(
         m_context.master_key,

+ 10 - 2
Userland/Libraries/LibTLS/Record.cpp

@@ -37,7 +37,11 @@ void TLSv12::alert(AlertLevel level, AlertDescription code)
 
 void TLSv12::write_packet(ByteBuffer& packet)
 {
-    m_context.tls_buffer.append(packet.data(), packet.size());
+    auto ok = m_context.tls_buffer.try_append(packet.data(), packet.size());
+    if (!ok) {
+        // Toooooo bad, drop the record on the ground.
+        return;
+    }
     if (m_context.connection_status > ConnectionStatus::Disconnected) {
         if (!m_has_scheduled_write_flush) {
             dbgln_if(TLS_DEBUG, "Scheduling write of {}", m_context.tls_buffer.size());
@@ -451,7 +455,11 @@ ssize_t TLSv12::handle_message(ReadonlyBytes buffer)
         } else {
             dbgln_if(TLS_DEBUG, "application data message of size {}", plain.size());
 
-            m_context.application_buffer.append(plain.data(), plain.size());
+            if (!m_context.application_buffer.try_append(plain.data(), plain.size())) {
+                payload_res = (i8)Error::DecryptionFailed;
+                auto packet = build_alert(true, (u8)AlertDescription::DecryptionFailed);
+                write_packet(packet);
+            }
         }
         break;
     case MessageType::Handshake:

+ 4 - 1
Userland/Libraries/LibTLS/TLSv12.cpp

@@ -35,7 +35,10 @@ void TLSv12::consume(ReadonlyBytes record)
 
     dbgln_if(TLS_DEBUG, "Consuming {} bytes", record.size());
 
-    m_context.message_buffer.append(record.data(), record.size());
+    if (!m_context.message_buffer.try_append(record.data(), record.size())) {
+        dbgln("Not enough space in message buffer, dropping the record");
+        return;
+    }
 
     size_t index { 0 };
     size_t buffer_length = m_context.message_buffer.size();

+ 2 - 1
Userland/Libraries/LibWasm/AbstractMachine/AbstractMachine.h

@@ -349,7 +349,8 @@ public:
                 return false;
         }
         auto previous_size = m_size;
-        m_data.resize(new_size);
+        if (!m_data.try_resize(new_size))
+            return false;
         m_size = new_size;
         // The spec requires that we zero out everything on grow
         __builtin_memset(m_data.offset_pointer(previous_size), 0, size_to_grow);

+ 2 - 1
Userland/Libraries/LibWasm/Parser/Parser.cpp

@@ -744,7 +744,8 @@ ParseResult<CustomSection> CustomSection::parse(InputStream& stream)
         auto size = stream.read({ buf, 16 });
         if (size == 0)
             break;
-        data_buffer.append(buf, size);
+        if (!data_buffer.try_append(buf, size))
+            return with_eof_check(stream, ParseError::HugeAllocationRequested);
     }
 
     return CustomSection(name.release_value(), move(data_buffer));

+ 4 - 1
Userland/Services/InspectorServer/InspectableProcess.cpp

@@ -57,7 +57,10 @@ String InspectableProcess::wait_for_response()
         auto packet = m_socket->read(remaining_bytes);
         if (packet.size() == 0)
             break;
-        data.append(packet.data(), packet.size());
+        if (!data.try_append(packet.data(), packet.size())) {
+            dbgln("Failed to append {} bytes to data buffer", packet.size());
+            break;
+        }
         remaining_bytes -= packet.size();
     }
 

+ 2 - 1
Userland/Utilities/pro.cpp

@@ -122,7 +122,8 @@ private:
     {
         if (!m_condition()) {
         write_to_buffer:;
-            m_buffer.append(bytes.data(), bytes.size());
+            if (!m_buffer.try_append(bytes.data(), bytes.size()))
+                return 0;
             return bytes.size();
         }
 

+ 7 - 3
Userland/Utilities/test-crypto.cpp

@@ -165,8 +165,9 @@ static void tls(const char* message, size_t len)
             g_loop.quit(0);
         };
     }
-    write.append(message, len);
-    write.append("\r\n", 2);
+    auto ok = write.try_append(message, len);
+    ok = ok && write.try_append("\r\n", 2);
+    VERIFY(ok);
 }
 
 static void aes_cbc(const char* message, size_t len)
@@ -2037,7 +2038,10 @@ static void tls_test_client_hello()
             loop.quit(1);
         } else {
             //            print_buffer(data.value(), 16);
-            contents.append(data.value().data(), data.value().size());
+            if (!contents.try_append(data.value().data(), data.value().size())) {
+                FAIL(Allocation failed);
+                loop.quit(1);
+            }
         }
     };
     tls->on_tls_finished = [&] {