hexdump: Support truncating output with -n

1 files changed, 31 insertions, 12 deletions

diff --git a/Userland/Utilities/hexdump.cpp b/Userland/Utilities/hexdump.cpp
index 4b2e21d03a..6ccfd8d7c7 100644
--- a/Userland/Utilities/hexdump.cpp
+++ b/Userland/Utilities/hexdump.cpp
@@ -27,8 +27,11 @@ ErrorOr<int> serenity_main(Main::Arguments args)
     Core::ArgsParser args_parser;
     char const* path = nullptr;
     bool verbose = false;
+    Optional<size_t> max_bytes;
+
     args_parser.add_positional_argument(path, "Input", "input", Core::ArgsParser::Required::No);
     args_parser.add_option(verbose, "Display all input data", "verbose", 'v');
+    args_parser.add_option(max_bytes, "Truncate to a fixed number of bytes", nullptr, 'n', "bytes");
 
     args_parser.parse(args);
 
@@ -67,18 +70,34 @@ ErrorOr<int> serenity_main(Main::Arguments args)
     Array<u8, BUFSIZ> contents;
     Span<u8> previous_line;
     static_assert(LINE_LENGTH_BYTES * 2 <= contents.size(), "Buffer is too small?!");
-    size_t contents_size = 0;
+    size_t bytes_in_buffer = 0;
+    size_t total_bytes_read = 0;
+    size_t bytes_remaining = 0;
+    size_t bytes_to_read = 0;
 
     int nread;
     auto state = State::Print;
-    while (true) {
-        nread = file->read(&contents[contents_size], BUFSIZ - contents_size);
+    bool is_input_remaining = true;
+    while (is_input_remaining) {
+        bytes_to_read = BUFSIZ - bytes_in_buffer;
+
+        if (max_bytes.has_value()) {
+            bytes_remaining = max_bytes.value() - total_bytes_read;
+            if (bytes_remaining < bytes_to_read) {
+                bytes_to_read = bytes_remaining;
+                is_input_remaining = false;
+            }
+        }
+
+        nread = file->read(&contents[bytes_in_buffer], (int)bytes_to_read);
         if (nread <= 0)
             break;
-        contents_size += nread;
+
+        total_bytes_read += nread;
+        bytes_in_buffer += nread;
 
         size_t offset;
-        for (offset = 0; offset + LINE_LENGTH_BYTES - 1 < contents_size; offset += LINE_LENGTH_BYTES) {
+        for (offset = 0; offset + LINE_LENGTH_BYTES - 1 < bytes_in_buffer; offset += LINE_LENGTH_BYTES) {
             if (verbose) {
                 print_line(&contents[offset], LINE_LENGTH_BYTES);
                 continue;
@@ -106,17 +125,17 @@ ErrorOr<int> serenity_main(Main::Arguments args)
             previous_line = current_line;
         }
 
-        contents_size -= offset;
-        VERIFY(contents_size < LINE_LENGTH_BYTES);
+        bytes_in_buffer -= offset;
+        VERIFY(bytes_in_buffer < LINE_LENGTH_BYTES);
         // If we managed to make the buffer exactly full, &contents[BUFSIZ] would blow up.
-        if (contents_size > 0) {
+        if (bytes_in_buffer > 0) {
             // Regions cannot overlap due to above static_assert.
-            memcpy(&contents[0], &contents[offset], contents_size);
+            memcpy(&contents[0], &contents[offset], bytes_in_buffer);
         }
     }
-    VERIFY(contents_size <= LINE_LENGTH_BYTES - 1);
-    if (contents_size > 0)
-        print_line(&contents[0], contents_size);
+    VERIFY(bytes_in_buffer <= LINE_LENGTH_BYTES - 1);
+    if (bytes_in_buffer > 0)
+        print_line(&contents[0], bytes_in_buffer);
 
     return 0;
 }