xref: /aosp_15_r20/external/zucchini/main_utils.cc (revision a03ca8b91e029cd15055c20c78c2e087c84792e4)
1 // Copyright 2017 The Chromium Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4 
5 #include "components/zucchini/main_utils.h"
6 
7 #include <stddef.h>
8 
9 #include <memory>
10 #include <ostream>
11 #include <string>
12 #include <vector>
13 
14 #include "base/command_line.h"
15 #include "base/files/file_path.h"
16 #include "base/files/file_util.h"
17 #include "base/logging.h"
18 #include "base/process/process_handle.h"
19 #include "base/strings/string_number_conversions.h"
20 #include "base/strings/string_split.h"
21 #include "base/strings/string_util.h"
22 #include "base/time/time.h"
23 #include "build/build_config.h"
24 #include "components/zucchini/io_utils.h"
25 #include "components/zucchini/version_info.h"
26 #include "components/zucchini/zucchini_commands.h"
27 
28 #if defined(OS_WIN)
29 #include <windows.h>  // This include must come first.
30 
31 #include <psapi.h>
32 #endif
33 
34 namespace {
35 
36 /******** Command ********/
37 
38 // Specifications for a Zucchini command.
39 struct Command {
Command__anonf9d0c39c0111::Command40   constexpr Command(const char* name_in,
41                     const char* usage_in,
42                     int num_args_in,
43                     CommandFunction command_function_in)
44       : name(name_in),
45         usage(usage_in),
46         num_args(num_args_in),
47         command_function(command_function_in) {}
48   Command(const Command&) = default;
49   ~Command() = default;
50 
51   // Unique name of command. |-name| is used to select from command-line.
52   const char* const name;
53 
54   // Usage help text of command.
55   const char* const usage;
56 
57   // Number of arguments (assumed to be filenames) used by the command.
58   const int num_args;
59 
60   // Main function to run for the command.
61   const CommandFunction command_function;
62 };
63 
64 /******** List of Zucchini commands ********/
65 
66 constexpr Command kCommands[] = {
67     {"gen",
68      "-gen <old_file> <new_file> <patch_file> [-raw] [-keep]"
69      " [-impose=#+#=#+#,#+#=#+#,...]",
70      3, &MainGen},
71     {"apply", "-apply <old_file> <patch_file> <new_file> [-keep]", 3,
72      &MainApply},
73     {"verify", "-verify <patch_file>", 1, &MainVerify},
74     {"read", "-read <exe> [-dump]", 1, &MainRead},
75     {"detect", "-detect <archive_file>", 1, &MainDetect},
76     {"match", "-match <old_file> <new_file> [-impose=#+#=#+#,#+#=#+#,...]", 2,
77      &MainMatch},
78     {"crc32", "-crc32 <file>", 1, &MainCrc32},
79 };
80 
81 /******** GetPeakMemoryMetrics ********/
82 
83 #if defined(OS_LINUX) || defined(OS_CHROMEOS)
84 // Linux does not have an exact mapping to the values used on Windows so use a
85 // close approximation:
86 // peak_virtual_memory ~= peak_page_file_usage
87 // resident_set_size_hwm (high water mark) ~= peak_working_set_size
88 //
89 // On failure the input values will be set to 0.
GetPeakMemoryMetrics(size_t * peak_virtual_memory,size_t * resident_set_size_hwm)90 void GetPeakMemoryMetrics(size_t* peak_virtual_memory,
91                           size_t* resident_set_size_hwm) {
92   *peak_virtual_memory = 0;
93   *resident_set_size_hwm = 0;
94   auto status_path =
95       base::FilePath("/proc")
96           .Append(base::NumberToString(base::GetCurrentProcessHandle()))
97           .Append("status");
98   std::string contents_string;
99   base::ReadFileToString(status_path, &contents_string);
100   std::vector<base::StringPiece> lines = base::SplitStringPiece(
101       contents_string, "\n", base::TRIM_WHITESPACE, base::SPLIT_WANT_NONEMPTY);
102 
103   for (const auto& line : lines) {
104     // Tokens should generally be of the form "Metric: <val> kB"
105     std::vector<base::StringPiece> tokens = base::SplitStringPiece(
106         line, " ", base::TRIM_WHITESPACE, base::SPLIT_WANT_NONEMPTY);
107     if (tokens.size() < 2)
108       continue;
109 
110     if (tokens[0] == "VmPeak:") {
111       if (base::StringToSizeT(tokens[1], peak_virtual_memory)) {
112         *peak_virtual_memory *= 1024;  // in kiB
113         if (*resident_set_size_hwm)
114           return;
115       }
116     } else if (tokens[0] == "VmHWM:") {
117       if (base::StringToSizeT(tokens[1], resident_set_size_hwm)) {
118         *resident_set_size_hwm *= 1024;  // in kiB
119         if (*peak_virtual_memory)
120           return;
121       }
122     }
123   }
124 }
125 #endif  // defined(OS_LINUX) || defined(OS_CHROMEOS)
126 
127 #if defined(OS_WIN)
128 // On failure the input values will be set to 0.
GetPeakMemoryMetrics(size_t * peak_page_file_usage,size_t * peak_working_set_size)129 void GetPeakMemoryMetrics(size_t* peak_page_file_usage,
130                           size_t* peak_working_set_size) {
131   *peak_page_file_usage = 0;
132   *peak_working_set_size = 0;
133   PROCESS_MEMORY_COUNTERS pmc;
134   if (::GetProcessMemoryInfo(::GetCurrentProcess(), &pmc, sizeof(pmc))) {
135     *peak_page_file_usage = pmc.PeakPagefileUsage;
136     *peak_working_set_size = pmc.PeakWorkingSetSize;
137   }
138 }
139 #endif  // defined(OS_WIN)
140 
141 /******** ScopedResourceUsageTracker ********/
142 
143 // A class to track and log system resource usage.
144 class ScopedResourceUsageTracker {
145  public:
146   // Initializes states for tracking.
ScopedResourceUsageTracker()147   ScopedResourceUsageTracker() {
148     start_time_ = base::TimeTicks::Now();
149 
150 #if defined(OS_LINUX) || defined(OS_CHROMEOS) || defined(OS_WIN)
151     GetPeakMemoryMetrics(&start_peak_page_file_usage_,
152                          &start_peak_working_set_size_);
153 #endif  // defined(OS_LINUX) || defined(OS_CHROMEOS) || defined(OS_WIN)
154   }
155 
156   // Computes and prints usage.
~ScopedResourceUsageTracker()157   ~ScopedResourceUsageTracker() {
158     base::TimeTicks end_time = base::TimeTicks::Now();
159 
160 #if defined(OS_LINUX) || defined(OS_CHROMEOS) || defined(OS_WIN)
161     size_t cur_peak_page_file_usage = 0;
162     size_t cur_peak_working_set_size = 0;
163     GetPeakMemoryMetrics(&cur_peak_page_file_usage, &cur_peak_working_set_size);
164 
165     LOG(INFO) << "Zucchini.PeakPagefileUsage "
166               << cur_peak_page_file_usage / 1024 << " KiB";
167     LOG(INFO) << "Zucchini.PeakPagefileUsageChange "
168               << (cur_peak_page_file_usage - start_peak_page_file_usage_) / 1024
169               << " KiB";
170     LOG(INFO) << "Zucchini.PeakWorkingSetSize "
171               << cur_peak_working_set_size / 1024 << " KiB";
172     LOG(INFO) << "Zucchini.PeakWorkingSetSizeChange "
173               << (cur_peak_working_set_size - start_peak_working_set_size_) /
174                      1024
175               << " KiB";
176 #endif  // defined(OS_LINUX) || defined(OS_CHROMEOS) || defined(OS_WIN)
177 
178     LOG(INFO) << "Zucchini.TotalTime " << (end_time - start_time_).InSecondsF()
179               << " s";
180   }
181 
182  private:
183   base::TimeTicks start_time_;
184 #if defined(OS_LINUX) || defined(OS_CHROMEOS) || defined(OS_WIN)
185   size_t start_peak_page_file_usage_ = 0;
186   size_t start_peak_working_set_size_ = 0;
187 #endif  // defined(OS_LINUX) || defined(OS_CHROMEOS) || defined(OS_WIN)
188 };
189 
190 /******** Helper functions ********/
191 
192 // Translates |command_line| arguments to a vector of base::FilePath (expecting
193 // exactly |expected_count|). On success, writes the results to |paths| and
194 // returns true. Otherwise returns false.
CheckAndGetFilePathParams(const base::CommandLine & command_line,size_t expected_count,std::vector<base::FilePath> * paths)195 bool CheckAndGetFilePathParams(const base::CommandLine& command_line,
196                                size_t expected_count,
197                                std::vector<base::FilePath>* paths) {
198   const base::CommandLine::StringVector& args = command_line.GetArgs();
199   if (args.size() != expected_count)
200     return false;
201 
202   paths->clear();
203   paths->reserve(args.size());
204   for (const auto& arg : args)
205     paths->emplace_back(arg);
206   return true;
207 }
208 
209 // Prints main Zucchini usage text.
PrintUsage(std::ostream & err)210 void PrintUsage(std::ostream& err) {
211   err << "Version: " << zucchini::kMajorVersion << "."
212       << zucchini::kMinorVersion << std::endl;
213   err << "Usage:" << std::endl;
214   for (const Command& command : kCommands)
215     err << "  zucchini " << command.usage << std::endl;
216 }
217 
218 }  // namespace
219 
220 /******** Exported Functions ********/
221 
RunZucchiniCommand(const base::CommandLine & command_line,std::ostream & out,std::ostream & err)222 zucchini::status::Code RunZucchiniCommand(const base::CommandLine& command_line,
223                                           std::ostream& out,
224                                           std::ostream& err) {
225   // Look for a command with name that matches input.
226   const Command* command_use = nullptr;
227   for (const Command& command : kCommands) {
228     if (command_line.HasSwitch(command.name)) {
229       if (command_use) {        // Too many commands found.
230         command_use = nullptr;  // Set to null to flag error.
231         break;
232       }
233       command_use = &command;
234     }
235   }
236 
237   // Expect exactly 1 matching command. If 0 or >= 2, print usage and quit.
238   if (!command_use) {
239     err << "Must have exactly one of:" << std::endl;
240     err << "  [";
241     zucchini::PrefixSep sep(", ");
242     for (const Command& command : kCommands)
243       err << sep << "-" << command.name;
244     err << "]" << std::endl;
245     PrintUsage(err);
246     return zucchini::status::kStatusInvalidParam;
247   }
248 
249   // Try to parse filename arguments. On failure, print usage and quit.
250   std::vector<base::FilePath> paths;
251   if (!CheckAndGetFilePathParams(command_line, command_use->num_args, &paths)) {
252     err << command_use->usage << std::endl;
253     PrintUsage(err);
254     return zucchini::status::kStatusInvalidParam;
255   }
256 
257   ScopedResourceUsageTracker resource_usage_tracker;
258   return command_use->command_function({command_line, paths, out, err});
259 }
260