diff --git a/Userland/Utilities/drain.cpp b/Userland/Utilities/drain.cpp index ad4134d0ea3..eb1f1d985fa 100644 --- a/Userland/Utilities/drain.cpp +++ b/Userland/Utilities/drain.cpp @@ -4,14 +4,13 @@ * SPDX-License-Identifier: BSD-2-Clause */ +#include #include #include #include #include #include -constexpr size_t block_size = 256 * KiB; - static ErrorOr seek_and_read(size_t offset, ByteBuffer& buffer, Core::File& file) { TRY(file.seek(offset, SeekMode::SetPosition)); @@ -26,7 +25,7 @@ static ErrorOr seek_and_write(size_t offset, ByteBuffer& buffer, Core::Fil return {}; } -static ErrorOr process_file(Core::File& file, size_t file_size, size_t file_size_rounded) +static ErrorOr process_file(Core::File& file, size_t block_size, size_t file_size, size_t file_size_rounded) { size_t head = 0; size_t tail = file_size_rounded - block_size; @@ -79,11 +78,21 @@ ErrorOr serenity_main(Main::Arguments arguments) TRY(Core::System::pledge("stdio cpath rpath wpath")); StringView path; + size_t block_size_in_kib = 256; + Core::ArgsParser args_parser; args_parser.set_general_help("Print file to stdout, while progressively deleting read segments."); + args_parser.add_option(block_size_in_kib, "Base Block size in KiB, defaults to 256 KiB", "block-size", 'b', "base block size"); args_parser.add_positional_argument(path, "File path", "path", Core::ArgsParser::Required::Yes); args_parser.parse(arguments); + if (block_size_in_kib < 1) + return Error::from_string_literal("Invalid block size"); + if (Checked::multiplication_would_overflow(block_size_in_kib, KiB)) + return Error::from_string_literal("Overflow in block size"); + + size_t block_size = block_size_in_kib * KiB; + if (!FileSystem::exists(path)) return Error::from_errno(ENOENT); @@ -94,7 +103,7 @@ ErrorOr serenity_main(Main::Arguments arguments) size_t file_size_rounded = TRY(file->seek(ceil_div(file_size, block_size) * block_size, SeekMode::SetPosition)); - TRY(process_file(*file, file_size, file_size_rounded)); + TRY(process_file(*file, block_size, file_size, file_size_rounded)); TRY(FileSystem::remove(path, FileSystem::RecursionMode::Disallowed));