From 84b981ade7f6e1c6580a152db64126940c10c5e8 Mon Sep 17 00:00:00 2001
From: implicitfield <114500360+implicitfield@users.noreply.github.com>
Date: Wed, 15 Nov 2023 19:13:36 +0400
Subject: [PATCH] Utilities: Add a `drain` utility

---
 Userland/Utilities/CMakeLists.txt |   1 +
 Userland/Utilities/drain.cpp      | 102 ++++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+)
 create mode 100644 Userland/Utilities/drain.cpp

diff --git a/Userland/Utilities/CMakeLists.txt b/Userland/Utilities/CMakeLists.txt
index 87e16bf4ef..947c8f8ca5 100644
--- a/Userland/Utilities/CMakeLists.txt
+++ b/Userland/Utilities/CMakeLists.txt
@@ -89,6 +89,7 @@ target_link_libraries(cpp-parser PRIVATE LibCpp)
 target_link_libraries(cpp-preprocessor PRIVATE LibCpp)
 target_link_libraries(diff PRIVATE LibDiff)
 target_link_libraries(disasm PRIVATE LibX86)
+target_link_libraries(drain PRIVATE LibFileSystem)
 target_link_libraries(expr PRIVATE LibRegex)
 target_link_libraries(fdtdump PRIVATE LibDeviceTree)
 target_link_libraries(file PRIVATE LibGfx LibIPC LibArchive LibCompress LibAudio)
diff --git a/Userland/Utilities/drain.cpp b/Userland/Utilities/drain.cpp
new file mode 100644
index 0000000000..ad4134d0ea
--- /dev/null
+++ b/Userland/Utilities/drain.cpp
@@ -0,0 +1,102 @@
+/*
+ * Copyright (c) 2023, the SerenityOS developers.
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#include <LibCore/ArgsParser.h>
+#include <LibCore/File.h>
+#include <LibCore/System.h>
+#include <LibFileSystem/FileSystem.h>
+#include <LibMain/Main.h>
+
+constexpr size_t block_size = 256 * KiB;
+
+static ErrorOr<void> seek_and_read(size_t offset, ByteBuffer& buffer, Core::File& file)
+{
+    TRY(file.seek(offset, SeekMode::SetPosition));
+    TRY(file.read_until_filled(buffer));
+    return {};
+}
+
+static ErrorOr<void> seek_and_write(size_t offset, ByteBuffer& buffer, Core::File& file)
+{
+    TRY(file.seek(offset, SeekMode::SetPosition));
+    TRY(file.write_until_depleted(buffer));
+    return {};
+}
+
+static ErrorOr<void> process_file(Core::File& file, size_t file_size, size_t file_size_rounded)
+{
+    size_t head = 0;
+    size_t tail = file_size_rounded - block_size;
+    size_t const last_block_size = file_size - tail;
+
+    auto head_buffer = TRY(ByteBuffer::create_uninitialized(block_size));
+    auto tail_buffer = TRY(ByteBuffer::create_uninitialized(last_block_size));
+
+    auto stdout = TRY(Core::File::standard_output());
+
+    // Overwrite the current block (after saving its contents to a temporary buffer) with the last block until we've processed half of the blocks in the file.
+    while (head <= tail) {
+        TRY(seek_and_read(head, head_buffer, file));
+        TRY(seek_and_read(tail, tail_buffer, file));
+
+        TRY(seek_and_write(head, tail_buffer, file));
+        TRY(file.truncate(tail));
+
+        TRY(stdout->write_until_depleted(head_buffer));
+
+        if (tail_buffer.size() != block_size)
+            TRY(tail_buffer.try_resize(block_size));
+
+        head += block_size;
+        tail -= block_size;
+    };
+
+    size_t remaining_size = file_size - head;
+
+    // Note that we iterate downwards from the end of the file, as the above algorithm left all of the remaining blocks in reverse order.
+    while (remaining_size) {
+        size_t to_write = remaining_size >= block_size ? block_size : last_block_size;
+
+        tail_buffer.trim(to_write, true);
+        TRY(seek_and_read(tail, tail_buffer, file));
+
+        TRY(file.truncate(tail));
+
+        TRY(stdout->write_until_depleted(tail_buffer));
+
+        tail -= to_write;
+        remaining_size -= to_write;
+    }
+
+    return {};
+}
+
+ErrorOr<int> serenity_main(Main::Arguments arguments)
+{
+    TRY(Core::System::pledge("stdio cpath rpath wpath"));
+    StringView path;
+
+    Core::ArgsParser args_parser;
+    args_parser.set_general_help("Print file to stdout, while progressively deleting read segments.");
+    args_parser.add_positional_argument(path, "File path", "path", Core::ArgsParser::Required::Yes);
+    args_parser.parse(arguments);
+
+    if (!FileSystem::exists(path))
+        return Error::from_errno(ENOENT);
+
+    auto file = TRY(Core::File::open(path, Core::File::OpenMode::ReadWrite));
+    size_t file_size = TRY(file->seek(0, SeekMode::FromEndPosition));
+    if (file_size < block_size * 2)
+        return Error::from_string_literal("Input file too small");
+
+    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(FileSystem::remove(path, FileSystem::RecursionMode::Disallowed));
+
+    return 0;
+}