jcs.org
Serenity Operating System
1/*
2 * Copyright (c) 2018-2020, Andreas Kling <kling@serenityos.org>
3 *
4 * SPDX-License-Identifier: BSD-2-Clause
5 */
6
7#include <AK/StdLibExtras.h>
8#include <LibCore/ArgsParser.h>
9#include <LibCore/System.h>
10#include <LibMain/Main.h>
11#include <errno.h>
12#include <fcntl.h>
13#include <stdio.h>
14#include <string.h>
15#include <unistd.h>
16
17int head(DeprecatedString const& filename, bool print_filename, ssize_t line_count, ssize_t byte_count);
18
19ErrorOr<int> serenity_main(Main::Arguments args)
20{
21 TRY(Core::System::pledge("stdio rpath"));
22
23 int line_count = -1;
24 int byte_count = -1;
25 bool never_print_filenames = false;
26 bool always_print_filenames = false;
27 Vector<DeprecatedString> files;
28
29 Core::ArgsParser args_parser;
30 args_parser.set_general_help("Print the beginning ('head') of a file.");
31 args_parser.add_option(line_count, "Number of lines to print (default 10)", "lines", 'n', "number");
32 args_parser.add_option(byte_count, "Number of bytes to print", "bytes", 'c', "number");
33 args_parser.add_option(never_print_filenames, "Never print filenames", "quiet", 'q');
34 args_parser.add_option(always_print_filenames, "Always print filenames", "verbose", 'v');
35 args_parser.add_positional_argument(files, "File to process", "file", Core::ArgsParser::Required::No);
36 args_parser.parse(args);
37
38 if (line_count == -1 && byte_count == -1) {
39 line_count = 10;
40 }
41
42 bool print_filenames = files.size() > 1;
43 if (always_print_filenames)
44 print_filenames = true;
45 else if (never_print_filenames)
46 print_filenames = false;
47
48 if (files.is_empty()) {
49 return head("", print_filenames, line_count, byte_count);
50 }
51
52 int rc = 0;
53
54 for (auto& file : files) {
55 if (head(file, print_filenames, line_count, byte_count) != 0) {
56 rc = 1;
57 }
58 }
59
60 return rc;
61}
62
63int head(DeprecatedString const& filename, bool print_filename, ssize_t line_count, ssize_t byte_count)
64{
65 bool is_stdin = false;
66 int fd = -1;
67
68 ScopeGuard fd_close_guard = [&fd] {
69 if (fd > 0)
70 close(fd);
71 };
72
73 if (filename == "" || filename == "-") {
74 fd = 0;
75 is_stdin = true;
76 } else {
77 fd = open(filename.characters(), O_RDONLY);
78 if (fd < 0) {
79 warnln("Failed to open {}: {}", filename, strerror(errno));
80 return 1;
81 }
82 }
83
84 if (print_filename) {
85 if (is_stdin) {
86 outln("==> standard input <==");
87 } else {
88 outln("==> {} <==", filename);
89 }
90 }
91
92 fflush(stdout);
93
94 size_t buffer_size = line_count != -1 ? BUFSIZ : min((size_t)BUFSIZ, (size_t)byte_count);
95 char buffer[buffer_size];
96
97 while (line_count > 0 || byte_count > 0) {
98 size_t ntoread = line_count != -1 ? buffer_size : min(buffer_size, (size_t)byte_count);
99 ssize_t nread = read(fd, buffer, ntoread);
100 if (nread < 0) {
101 perror("read");
102 return 1;
103 } else if (nread == 0) {
104 break;
105 }
106
107 size_t ntowrite;
108 if (byte_count != -1) {
109 // Write out everything we've read, since we have explicitly ensured
110 // that we wouldn't read more than we want to write.
111 ntowrite = nread;
112 byte_count -= nread;
113 } else {
114 // Count line breaks.
115 ntowrite = 0;
116 while (line_count) {
117 char const* newline = strchr(buffer + ntowrite, '\n');
118 if (newline) {
119 // Found another line break, include this line.
120 ntowrite = newline - buffer + 1;
121 line_count--;
122 } else {
123 // No more line breaks, write the whole thing.
124 ntowrite = nread;
125 break;
126 }
127 }
128 }
129
130 size_t ncomplete = 0;
131 while (ncomplete < ntowrite) {
132 ssize_t nwritten = write(1, buffer + ncomplete, ntowrite - ncomplete);
133 if (nwritten < 0) {
134 perror("write");
135 return 1;
136 }
137 ncomplete += nwritten;
138 }
139 }
140
141 if (print_filename) {
142 puts("");
143 }
144
145 return 0;
146}