1 /*
2  * Copyright (C) 2022-2024 Huawei Device Co., Ltd.
3  * Licensed under the Apache License, Version 2.0 (the "License");
4  * you may not use this file except in compliance with the License.
5  * You may obtain a copy of the License at
6  *
7  *     http://www.apache.org/licenses/LICENSE-2.0
8  *
9  * Unless required by applicable law or agreed to in writing, software
10  * distributed under the License is distributed on an "AS IS" BASIS,
11  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12  * See the License for the specific language governing permissions and
13  * limitations under the License.
14  */
15 
16 #include <cinttypes>
17 #include <csignal>
18 #include <cstdio>
19 #include <cstdlib>
20 #include <cstring>
21 #include <ctime>
22 #include <fcntl.h>
23 #include <fstream>
24 #include <getopt.h>
25 #include <map>
26 #include <regex>
27 #include <sstream>
28 #include <string>
29 #include <sys/stat.h>
30 #include <sys/types.h>
31 #include <thread>
32 #include <unistd.h>
33 #include <vector>
34 #include <memory>
35 #include <iostream>
36 #include <zlib.h>
37 
38 #include "hitrace_meter.h"
39 #include "common_utils.h"
40 #include "trace_collector.h"
41 #include "hitrace_osal.h"
42 #include "securec.h"
43 
44 using namespace std;
45 using namespace OHOS::HiviewDFX::HitraceOsal;
46 
47 namespace {
48 
49 struct TraceArgs {
50     std::string tags;
51     std::string tagGroups;
52     std::string clockType;
53     int bufferSize = 0;
54     int fileSize = 0;
55     bool overwrite = true;
56     std::string output;
57 
58     int duration = 0;
59     bool isCompress = false;
60 };
61 
62 enum RunningState {
63     /* Initial value */
64     STATE_NULL = 0,
65 
66     /* Record a short trace */
67     RECORDING_SHORT_TEXT = 1,  // --text
68     RECORDING_SHORT_RAW = 2,  // --raw
69 
70     /* Record a long trace */
71     RECORDING_LONG_BEGIN = 10,  // --trace_begin
72     RECORDING_LONG_DUMP = 11,  // --trace_dump
73     RECORDING_LONG_FINISH = 12,  // --trace_finish
74     RECORDING_LONG_FINISH_NODUMP = 13,  // --trace_finish_nodump
75     RECORDING_LONG_BEGIN_RECORD = 14,  // --trace_begin --record
76     RECORDING_LONG_FINISH_RECORD = 15,  // --trace_finish --record
77 
78     /* Manipulating trace services in snapshot mode */
79     SNAPSHOT_START = 20,  // --start_bgsrv
80     SNAPSHOT_DUMP = 21,  // --dump_bgsrv
81     SNAPSHOT_STOP = 22,  // --stop_bgsrv
82 
83     /* Help Info */
84     SHOW_HELP = 31,  // -h, --help
85     SHOW_LIST_CATEGORY = 32,  // -l, --list_categories
86 };
87 
88 const std::map<RunningState, std::string> STATE_INFO = {
89     { STATE_NULL, "STATE_NULL" },
90     { RECORDING_SHORT_TEXT, "RECORDING_SHORT_TEXT" },
91     { RECORDING_SHORT_RAW, "RECORDING_SHORT_RAW" },
92     { RECORDING_LONG_BEGIN, "RECORDING_LONG_BEGIN" },
93     { RECORDING_LONG_DUMP, "RECORDING_LONG_DUMP" },
94     { RECORDING_LONG_FINISH_NODUMP, "RECORDING_LONG_FINISH_NODUMP" },
95     { RECORDING_LONG_BEGIN_RECORD, "RECORDING_LONG_BEGIN_RECORD" },
96     { RECORDING_LONG_FINISH_RECORD, "RECORDING_LONG_FINISH_RECORD" },
97     { SNAPSHOT_START, "SNAPSHOT_START" },
98     { SNAPSHOT_DUMP, "SNAPSHOT_DUMP" },
99     { SNAPSHOT_STOP, "SNAPSHOT_STOP" },
100     { SHOW_HELP, "SHOW_HELP" },
101     { SHOW_LIST_CATEGORY, "SHOW_LIST_CATEGORY" },
102 };
103 
104 constexpr struct option LONG_OPTIONS[] = {
105     { "buffer_size",       required_argument, nullptr, 0 },
106     { "trace_clock",       required_argument, nullptr, 0 },
107     { "help",              no_argument,       nullptr, 0 },
108     { "output",            required_argument, nullptr, 0 },
109     { "time",              required_argument, nullptr, 0 },
110     { "text",              no_argument,       nullptr, 0 },
111     { "raw",               no_argument,       nullptr, 0 },
112     { "trace_begin",       no_argument,       nullptr, 0 },
113     { "trace_finish",      no_argument,       nullptr, 0 },
114     { "trace_finish_nodump",      no_argument,       nullptr, 0 },
115     { "record",            no_argument,       nullptr, 0 },
116     { "trace_dump",        no_argument,       nullptr, 0 },
117     { "list_categories",   no_argument,       nullptr, 0 },
118     { "overwrite",         no_argument,       nullptr, 0 },
119     { "start_bgsrv",       no_argument,       nullptr, 0 },
120     { "dump_bgsrv",        no_argument,       nullptr, 0 },
121     { "stop_bgsrv",        no_argument,       nullptr, 0 },
122     { "file_size",         required_argument, nullptr, 0 },
123     { nullptr,             0,                 nullptr, 0 },
124 };
125 const unsigned int CHUNK_SIZE = 65536;
126 
127 constexpr const char *TRACE_TAG_PROPERTY = "debug.hitrace.tags.enableflags";
128 
129 // various operating paths of ftrace
130 constexpr const char *TRACING_ON_PATH = "tracing_on";
131 constexpr const char *TRACE_PATH = "trace";
132 constexpr const char *TRACE_MARKER_PATH = "trace_marker";
133 
134 // support customization of some parameters
135 const int KB_PER_MB = 1024;
136 const int MIN_BUFFER_SIZE = 256;
137 const int MAX_BUFFER_SIZE = 307200; // 300 MB
138 const int HM_MAX_BUFFER_SIZE = 1024 * KB_PER_MB; // 1024 MB
139 const int DEFAULT_BUFFER_SIZE = 18432; // 18 MB
140 constexpr unsigned int MAX_OUTPUT_LEN = 255;
141 const int PAGE_SIZE_KB = 4; // 4 KB
142 const int MIN_FILE_SIZE = 51200; // 50 MB
143 const int MAX_FILE_SIZE = 512000; // 500 MB
144 
145 string g_traceRootPath;
146 
147 std::shared_ptr<OHOS::HiviewDFX::UCollectClient::TraceCollector> g_traceCollector;
148 
149 TraceArgs g_traceArgs;
150 std::map<std::string, OHOS::HiviewDFX::Hitrace::TagCategory> g_allTags;
151 std::map<std::string, std::vector<std::string>> g_allTagGroups;
152 RunningState g_runningState = STATE_NULL;
153 }
154 
ConsoleLog(const std::string & logInfo)155 static void ConsoleLog(const std::string& logInfo)
156 {
157     // get localtime
158     time_t currentTime;
159     time(&currentTime);
160     struct tm timeInfo = {};
161     const int bufferSize = 20;
162     char timeStr[bufferSize] = {0};
163     localtime_r(&currentTime, &timeInfo);
164     strftime(timeStr, bufferSize, "%Y/%m/%d %H:%M:%S", &timeInfo);
165     std::cout << timeStr << " " << logInfo << std::endl;
166 }
167 
GetStateInfo(const RunningState state)168 static std::string GetStateInfo(const RunningState state)
169 {
170     if (STATE_INFO.find(state) == STATE_INFO.end()) {
171         ConsoleLog("error: running_state is invalid.");
172         return "";
173     }
174     return STATE_INFO.at(state);
175 }
176 
IsTraceMounted()177 static bool IsTraceMounted()
178 {
179     const string debugfsPath = "/sys/kernel/debug/tracing/";
180     const string tracefsPath = "/sys/kernel/tracing/";
181 
182     if (access((debugfsPath + TRACE_MARKER_PATH).c_str(), F_OK) != -1) {
183         g_traceRootPath = debugfsPath;
184         return true;
185     }
186     if (access((tracefsPath + TRACE_MARKER_PATH).c_str(), F_OK) != -1) {
187         g_traceRootPath = tracefsPath;
188         return true;
189     }
190     return false;
191 }
192 
WriteStrToFile(const string & filename,const std::string & str)193 static bool WriteStrToFile(const string& filename, const std::string& str)
194 {
195     ofstream out;
196     std::string inSpecPath =
197         OHOS::HiviewDFX::Hitrace::CanonicalizeSpecPath((g_traceRootPath + filename).c_str());
198     out.open(inSpecPath, ios::out);
199     if (out.fail()) {
200         ConsoleLog("error: open " + inSpecPath + " failed.");
201         return false;
202     }
203     out << str;
204     if (out.bad()) {
205         ConsoleLog("error: can not write " + inSpecPath);
206         out.close();
207         return false;
208     }
209     out.flush();
210     out.close();
211     return true;
212 }
213 
SetFtraceEnabled(const string & path,bool enabled)214 static bool SetFtraceEnabled(const string& path, bool enabled)
215 {
216     return WriteStrToFile(path, enabled ? "1" : "0");
217 }
218 
SetProperty(const string & property,const string & value)219 static bool SetProperty(const string& property, const string& value)
220 {
221     return SetPropertyInner(property, value);
222 }
223 
SetTraceTagsEnabled(uint64_t tags)224 static bool SetTraceTagsEnabled(uint64_t tags)
225 {
226     string value = std::to_string(tags);
227     return SetProperty(TRACE_TAG_PROPERTY, value);
228 }
229 
ShowListCategory()230 static void ShowListCategory()
231 {
232     printf("  %18s   description:\n", "tagName:");
233     for (auto it = g_allTags.begin(); it != g_allTags.end(); ++it) {
234         printf("  %18s - %s\n", it->first.c_str(), it->second.description.c_str());
235     }
236 }
237 
ShowHelp(const string & cmd)238 static void ShowHelp(const string& cmd)
239 {
240     printf("usage: %s [options] [categories...]\n", cmd.c_str());
241     printf("options include:\n"
242            "  -b N               Sets the size of the buffer (KB) for storing and reading traces. The default \n"
243            "                     buffer size is 2048 KB.\n"
244            "  --buffer_size N    Like \"-b N\".\n"
245            "  -l                 Lists available hitrace categories.\n"
246            "  --list_categories  Like \"-l\".\n"
247            "  -t N               Sets the hitrace running duration in seconds (5s by default), which depends on \n"
248            "                     the time required for analysis.\n"
249            "  --time N           Like \"-t N\".\n"
250            "  --trace_clock clock\n"
251            "                     Sets the type of the clock for adding a timestamp to a trace, which can be\n"
252            "                     boot (default), global, mono, uptime, or perf.\n"
253            "  --trace_begin      Starts capturing traces.\n"
254            "  --trace_dump       Dumps traces to a specified path (stdout by default).\n"
255            "  --trace_finish     Stops capturing traces and dumps traces to a specified path (stdout by default).\n"
256            "  --trace_finish_nodump\n"
257            "                     Stops capturing traces and not dumps traces.\n"
258            "  --record           Enable or disable long-term trace collection tasks in conjunction with\n"
259            "                    \"--trace_begin\" and \"--trace_finish\".\n"
260            "  --overwrite        Sets the action to take when the buffer is full. If this option is used,\n"
261            "                     the latest traces are discarded; if this option is not used (default setting),\n"
262            "                     the earliest traces are discarded.\n"
263            "  -o filename        Specifies the name of the target file (stdout by default).\n"
264            "  --output filename\n"
265            "                     Like \"-o filename\".\n"
266            "  -z                 Compresses a captured trace.\n"
267            "  --text             Specify the output format of trace as text.\n"
268            "  --raw              Specify the output format of trace as raw trace, the default format is text.\n"
269            "  --start_bgsrv      Enable trace_service in snapshot mode.\n"
270            "  --dump_bgsrv       Trigger the dump trace task of the trace_service.\n"
271            "  --stop_bgsrv       Disable trace_service in snapshot mode.\n"
272            "  --file_size        Sets the size of the raw trace (KB). The default file size is 102400 KB.\n"
273            "                     Only effective in raw trace mode\n"
274     );
275 }
276 
277 template <typename T>
StrToNum(const std::string & sString,T & tX)278 inline bool StrToNum(const std::string& sString, T &tX)
279 {
280     std::istringstream iStream(sString);
281     return (iStream >> tX) ? true : false;
282 }
283 
SetRunningState(const RunningState & setValue)284 static bool SetRunningState(const RunningState& setValue)
285 {
286     if (g_runningState != STATE_NULL) {
287         ConsoleLog("error: the parameter is set incorrectly, " + GetStateInfo(g_runningState) +
288                    " and " + GetStateInfo(setValue) + " cannot coexist.");
289         return false;
290     }
291     g_runningState = setValue;
292     return true;
293 }
294 
CheckOutputFile(const char * path)295 static bool CheckOutputFile(const char* path)
296 {
297     struct stat buf;
298     size_t len = strnlen(path, MAX_OUTPUT_LEN);
299     if (len == MAX_OUTPUT_LEN || len < 1 || (stat(path, &buf) == 0 && (buf.st_mode & S_IFDIR))) {
300         ConsoleLog("error: output file is illegal");
301         return false;
302     }
303     g_traceArgs.output = path;
304     return true;
305 }
306 
ParseLongOpt(const string & cmd,int optionIndex)307 static bool ParseLongOpt(const string& cmd, int optionIndex)
308 {
309     bool isTrue = true;
310     if (!strcmp(LONG_OPTIONS[optionIndex].name, "buffer_size")) {
311         int bufferSizeKB = 0;
312         int maxBufferSizeKB = MAX_BUFFER_SIZE;
313         if (IsHmKernel()) {
314             maxBufferSizeKB = HM_MAX_BUFFER_SIZE;
315         }
316         if (!StrToNum(optarg, bufferSizeKB)) {
317             ConsoleLog("error: buffer size is illegal input. eg: \"--buffer_size 1024\".");
318             isTrue = false;
319         } else if (bufferSizeKB < MIN_BUFFER_SIZE || bufferSizeKB > maxBufferSizeKB) {
320             ConsoleLog("error: buffer size must be from 256 KB to " + std::to_string(maxBufferSizeKB / KB_PER_MB) +
321                 " MB. eg: \"--buffer_size 1024\".");
322             isTrue = false;
323         }
324         g_traceArgs.bufferSize = bufferSizeKB / PAGE_SIZE_KB * PAGE_SIZE_KB;
325     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "trace_clock")) {
326         regex re("[a-zA-Z]{4,6}");
327         if (regex_match(optarg, re)) {
328             g_traceArgs.clockType = optarg;
329         } else {
330             ConsoleLog("error: \"--trace_clock\" is illegal input. eg: \"--trace_clock boot\".");
331             isTrue = false;
332         }
333     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "help")) {
334         isTrue = SetRunningState(SHOW_HELP);
335     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "time")) {
336         if (!StrToNum(optarg, g_traceArgs.duration)) {
337             ConsoleLog("error: the time is illegal input. eg: \"--time 5\".");
338             isTrue = false;
339         } else if (g_traceArgs.duration < 1) {
340             ConsoleLog("error: \"-t " + std::string(optarg) + "\" to be greater than zero. eg: \"--time 5\".");
341             isTrue = false;
342         }
343     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "list_categories")) {
344         isTrue = SetRunningState(SHOW_LIST_CATEGORY);
345     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "output")) {
346         isTrue = CheckOutputFile(optarg);
347     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "overwrite")) {
348         g_traceArgs.overwrite = false;
349     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "trace_begin")) {
350         isTrue = SetRunningState(RECORDING_LONG_BEGIN);
351     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "trace_finish")) {
352         isTrue = SetRunningState(RECORDING_LONG_FINISH);
353     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "trace_finish_nodump")) {
354         isTrue = SetRunningState(RECORDING_LONG_FINISH_NODUMP);
355     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "trace_dump")) {
356         isTrue = SetRunningState(RECORDING_LONG_DUMP);
357     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "record")) {
358         if (g_runningState == RECORDING_LONG_BEGIN) {
359             g_runningState = RECORDING_LONG_BEGIN_RECORD;
360         } else if (g_runningState == RECORDING_LONG_FINISH) {
361             g_runningState = RECORDING_LONG_FINISH_RECORD;
362         } else {
363             ConsoleLog("error: \"--record\" is set incorrectly. eg: \"--trace_begin --record\","
364                        " \"--trace_finish --record\".");
365             isTrue = false;
366         }
367     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "start_bgsrv")) {
368         isTrue = SetRunningState(SNAPSHOT_START);
369     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "dump_bgsrv")) {
370         isTrue = SetRunningState(SNAPSHOT_DUMP);
371     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "stop_bgsrv")) {
372         isTrue = SetRunningState(SNAPSHOT_STOP);
373     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "text")) {
374         isTrue = SetRunningState(RECORDING_SHORT_TEXT);
375     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "raw")) {
376         isTrue = SetRunningState(RECORDING_SHORT_RAW);
377     } else if (!strcmp(LONG_OPTIONS[optionIndex].name, "file_size")) {
378         int fileSizeKB = 0;
379         if (!StrToNum(optarg, fileSizeKB)) {
380             ConsoleLog("error: file size is illegal input. eg: \"--file_size 1024\".");
381             isTrue = false;
382         } else if (fileSizeKB < MIN_FILE_SIZE || fileSizeKB > MAX_FILE_SIZE) {
383             ConsoleLog("error: file size must be from 50 MB to 500 MB. eg: \"--file_size 102400\".");
384             isTrue = false;
385         }
386         g_traceArgs.fileSize = fileSizeKB;
387     }
388 
389     return isTrue;
390 }
391 
ParseOpt(int opt,char ** argv,int optIndex)392 static bool ParseOpt(int opt, char** argv, int optIndex)
393 {
394     bool isTrue = true;
395     switch (opt) {
396         case 'b': {
397             int bufferSizeKB = 0;
398             int maxBufferSizeKB = MAX_BUFFER_SIZE;
399             if (IsHmKernel()) {
400                 maxBufferSizeKB = HM_MAX_BUFFER_SIZE;
401             }
402             if (!StrToNum(optarg, bufferSizeKB)) {
403                 ConsoleLog("error: buffer size is illegal input. eg: \"--buffer_size 1024\".");
404                 isTrue = false;
405             } else if (bufferSizeKB < MIN_BUFFER_SIZE || bufferSizeKB > maxBufferSizeKB) {
406                 ConsoleLog("error: buffer size must be from 256 KB to " + std::to_string(maxBufferSizeKB / KB_PER_MB) +
407                 " MB. eg: \"--buffer_size 1024\".");
408                 isTrue = false;
409             }
410             g_traceArgs.bufferSize = bufferSizeKB / PAGE_SIZE_KB * PAGE_SIZE_KB;
411             break;
412         }
413         case 'h':
414             isTrue = SetRunningState(SHOW_HELP);
415             break;
416         case 'l':
417             isTrue = SetRunningState(SHOW_LIST_CATEGORY);
418             break;
419         case 't': {
420             if (!StrToNum(optarg, g_traceArgs.duration)) {
421                 ConsoleLog("error: the time is illegal input. eg: \"--time 5\".");
422                 isTrue = false;
423             } else if (g_traceArgs.duration < 1) {
424                 ConsoleLog("error: \"-t " + std::string(optarg) + "\" to be greater than zero. eg: \"--time 5\".");
425                 isTrue = false;
426             }
427             break;
428         }
429         case 'o': {
430             isTrue = CheckOutputFile(optarg);
431             break;
432         }
433         case 'z':
434             g_traceArgs.isCompress = true;
435             break;
436         case 0: // long options
437             isTrue = ParseLongOpt(argv[0], optIndex);
438             break;
439         case '?':
440             isTrue = false;
441             break;
442         default:
443             break;
444     }
445     return isTrue;
446 }
447 
AddTagItems(int argc,char ** argv)448 static bool AddTagItems(int argc, char** argv)
449 {
450     for (int i = optind; i < argc; i++) {
451         std::string tag = std::string(argv[i]);
452         if (g_allTags.find(tag) == g_allTags.end()) {
453             std::string errorInfo = "error: " + tag + " is not support category on this device.";
454             ConsoleLog(errorInfo);
455             return false;
456         }
457 
458         if (i == optind) {
459             g_traceArgs.tags = tag;
460         } else {
461             g_traceArgs.tags += ("," + tag);
462         }
463     }
464     return true;
465 }
466 
HandleOpt(int argc,char ** argv)467 static bool HandleOpt(int argc, char** argv)
468 {
469     bool isTrue = true;
470     int opt = 0;
471     int optionIndex = 0;
472     string shortOption = "b:c:hlo:t:z";
473     int argcSize = argc;
474     while (isTrue && argcSize-- > 0) {
475         opt = getopt_long(argc, argv, shortOption.c_str(), LONG_OPTIONS, &optionIndex);
476         if (opt < 0 && (!AddTagItems(argc, argv))) {
477             isTrue = false;
478             break;
479         }
480         isTrue = ParseOpt(opt, argv, optionIndex);
481     }
482 
483     return isTrue;
484 }
485 
StopTrace()486 static void StopTrace()
487 {
488     const int napTime = 10000;
489     usleep(napTime);
490     SetTraceTagsEnabled(0);
491     SetFtraceEnabled(TRACING_ON_PATH, false);
492 }
493 
DumpCompressedTrace(int traceFd,int outFd)494 static void DumpCompressedTrace(int traceFd, int outFd)
495 {
496     z_stream zs { nullptr };
497     int flush = Z_NO_FLUSH;
498     ssize_t bytesWritten;
499     ssize_t bytesRead;
500     if (memset_s(&zs, sizeof(zs), 0, sizeof(zs)) != 0) {
501         ConsoleLog("error: zip stream buffer init failed.");
502         return;
503     }
504     int ret = deflateInit(&zs, Z_DEFAULT_COMPRESSION);
505     if (ret != Z_OK) {
506         ConsoleLog("error: initializing zlib failed ret " + std::to_string(ret));
507         return;
508     }
509     std::unique_ptr<uint8_t[]>  in = std::make_unique<uint8_t[]>(CHUNK_SIZE);
510     std::unique_ptr<uint8_t[]>  out = std::make_unique<uint8_t[]>(CHUNK_SIZE);
511     if (!in || !out) {
512         ConsoleLog("error: couldn't allocate buffers.");
513         return;
514     }
515     zs.next_out = reinterpret_cast<Bytef*>(out.get());
516     zs.avail_out = CHUNK_SIZE;
517 
518     do {
519         if (zs.avail_in == 0 && flush == Z_NO_FLUSH) {
520             bytesRead = TEMP_FAILURE_RETRY(read(traceFd, in.get(), CHUNK_SIZE));
521             if (bytesRead == 0) {
522                 flush = Z_FINISH;
523             } else if (bytesRead == -1) {
524                 ConsoleLog("error: reading trace, errno " + std::to_string(errno));
525                 break;
526             } else {
527                 zs.next_in = reinterpret_cast<Bytef*>(in.get());
528                 zs.avail_in = bytesRead;
529             }
530         }
531         if (zs.avail_out == 0) {
532             bytesWritten = TEMP_FAILURE_RETRY(write(outFd, out.get(), CHUNK_SIZE));
533             if (bytesWritten < static_cast<ssize_t>(CHUNK_SIZE)) {
534                 ConsoleLog("error: writing deflated trace, errno " + std::to_string(errno));
535                 break;
536             }
537             zs.next_out = reinterpret_cast<Bytef*>(out.get());
538             zs.avail_out = CHUNK_SIZE;
539         }
540         ret = deflate(&zs, flush);
541         if (flush == Z_FINISH && ret == Z_STREAM_END) {
542             size_t have = CHUNK_SIZE - zs.avail_out;
543             bytesWritten = TEMP_FAILURE_RETRY(write(outFd, out.get(), have));
544             if (static_cast<size_t>(bytesWritten) < have) {
545                 ConsoleLog("error: writing deflated trace, errno " + std::to_string(errno));
546             }
547             break;
548         } else if (ret != Z_OK) {
549             if (ret == Z_ERRNO) {
550                 ConsoleLog("error: deflate failed with errno " + std::to_string(errno));
551             } else {
552                 ConsoleLog("error: deflate failed return " + std::to_string(ret));
553             }
554             break;
555         }
556     } while (ret == Z_OK);
557 
558     ret = deflateEnd(&zs);
559     if (ret != Z_OK) {
560         ConsoleLog("error: cleaning up zlib return " + std::to_string(ret));
561     }
562 }
563 
DumpTrace()564 static void DumpTrace()
565 {
566     std::string tracePath = g_traceRootPath + TRACE_PATH;
567     string traceSpecPath = OHOS::HiviewDFX::Hitrace::CanonicalizeSpecPath(tracePath.c_str());
568     int traceFd = open(traceSpecPath.c_str(), O_RDONLY);
569     if (traceFd == -1) {
570         ConsoleLog("error: opening " + tracePath + ", errno: " + std::to_string(errno));
571         return;
572     }
573 
574     int outFd = STDOUT_FILENO;
575     if (g_traceArgs.output.size() > 0) {
576         string outSpecPath = OHOS::HiviewDFX::Hitrace::CanonicalizeSpecPath(g_traceArgs.output.c_str());
577         outFd = open(outSpecPath.c_str(), O_WRONLY | O_CREAT | O_TRUNC, S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH);
578     }
579 
580     if (outFd == -1) {
581         ConsoleLog("error: opening " + g_traceArgs.output + ", errno: " + std::to_string(errno));
582         close(traceFd);
583         return;
584     }
585 
586     ssize_t bytesWritten;
587     ssize_t bytesRead;
588     if (g_traceArgs.isCompress) {
589         DumpCompressedTrace(traceFd, outFd);
590     } else {
591         const int blockSize = 4096;
592         char buffer[blockSize];
593         do {
594             bytesRead = TEMP_FAILURE_RETRY(read(traceFd, buffer, blockSize));
595             if ((bytesRead == 0) || (bytesRead == -1)) {
596                 break;
597             }
598             bytesWritten = TEMP_FAILURE_RETRY(write(outFd, buffer, bytesRead));
599         } while (bytesWritten > 0);
600     }
601 
602     if (outFd != STDOUT_FILENO) {
603         ConsoleLog("trace read done, output: " + g_traceArgs.output);
604         close(outFd);
605     }
606     close(traceFd);
607 }
608 
InitAllSupportTags()609 static bool InitAllSupportTags()
610 {
611     if (!OHOS::HiviewDFX::Hitrace::ParseTagInfo(g_allTags, g_allTagGroups)) {
612         ConsoleLog("error: hitrace_utils.json file is damaged.");
613         return false;
614     }
615     return true;
616 }
617 
ReloadTraceArgs()618 static std::string ReloadTraceArgs()
619 {
620     if (g_traceArgs.tags.size() == 0) {
621         ConsoleLog("error: tag is empty, please add.");
622         return "";
623     }
624     std::string args = "tags:" + g_traceArgs.tags;
625 
626     if (g_traceArgs.bufferSize > 0) {
627         args += (" bufferSize:" + std::to_string(g_traceArgs.bufferSize));
628     } else {
629         args += (" bufferSize:" + std::to_string(DEFAULT_BUFFER_SIZE));
630     }
631 
632     if (g_traceArgs.clockType.size() > 0) {
633         args += (" clockType:" + g_traceArgs.clockType);
634     }
635 
636     if (g_traceArgs.overwrite) {
637         args += " overwrite:";
638         args += "1";
639     } else {
640         args += " overwrite:";
641         args += "0";
642     }
643 
644     if (g_traceArgs.fileSize > 0) {
645         if (g_runningState == RECORDING_SHORT_RAW || g_runningState == RECORDING_LONG_BEGIN_RECORD) {
646             args += (" fileSize:" + std::to_string(g_traceArgs.fileSize));
647         } else {
648             ConsoleLog("warning: The current state does not support specifying the file size, file size: " +
649                 std::to_string(g_traceArgs.fileSize) + " is invalid.");
650         }
651     }
652 
653     if (g_runningState != RECORDING_SHORT_TEXT) {
654         ConsoleLog("args: " + args);
655     }
656     return args;
657 }
658 
HandleRecordingShortRaw()659 static bool HandleRecordingShortRaw()
660 {
661     std::string args = ReloadTraceArgs();
662     if (g_traceArgs.output.size() > 0) {
663         ConsoleLog("warning: The current state does not support specifying the output file path, " +
664                    g_traceArgs.output + " is invalid.");
665     }
666     auto openRet = g_traceCollector->OpenRecording(args);
667     if (openRet.retCode != OHOS::HiviewDFX::UCollect::UcError::SUCCESS) {
668         ConsoleLog("error: OpenRecording failed, errorCode(" + std::to_string(openRet.retCode) +")");
669         if (openRet.retCode == OHOS::HiviewDFX::UCollect::UcError::TRACE_IS_OCCUPIED ||
670             openRet.retCode == OHOS::HiviewDFX::UCollect::UcError::TRACE_CALL_ERROR) {
671             return false;
672         } else {
673             g_traceCollector->Recover();
674             return false;
675         }
676     }
677 
678     auto recOnRet = g_traceCollector->RecordingOn();
679     if (recOnRet.retCode != OHOS::HiviewDFX::UCollect::UcError::SUCCESS) {
680         ConsoleLog("error: RecordingOn failed, errorCode(" + std::to_string(recOnRet.retCode) +")");
681         g_traceCollector->Recover();
682         return false;
683     }
684     ConsoleLog("start capture, please wait " + std::to_string(g_traceArgs.duration) +"s ...");
685     sleep(g_traceArgs.duration);
686 
687     auto recOffRet = g_traceCollector->RecordingOff();
688     if (recOffRet.retCode != OHOS::HiviewDFX::UCollect::UcError::SUCCESS) {
689         ConsoleLog("error: RecordingOff failed, errorCode(" + std::to_string(recOffRet.retCode) +")");
690         g_traceCollector->Recover();
691         return false;
692     }
693     ConsoleLog("capture done, output files:");
694     for (std::string item : recOffRet.data) {
695         std::cout << "    " << item << std::endl;
696     }
697     g_traceCollector->Recover();
698     return true;
699 }
700 
HandleRecordingShortText()701 static bool HandleRecordingShortText()
702 {
703     std::string args = ReloadTraceArgs();
704     auto openRet = g_traceCollector->OpenRecording(args);
705     if (openRet.retCode != OHOS::HiviewDFX::UCollect::UcError::SUCCESS) {
706         ConsoleLog("error: OpenRecording failed, errorCode(" + std::to_string(openRet.retCode) +")");
707         if (openRet.retCode == OHOS::HiviewDFX::UCollect::UcError::TRACE_IS_OCCUPIED ||
708             openRet.retCode == OHOS::HiviewDFX::UCollect::UcError::TRACE_CALL_ERROR) {
709             return false;
710         } else {
711             g_traceCollector->Recover();
712             return false;
713         }
714     }
715     ConsoleLog("start capture, please wait " + std::to_string(g_traceArgs.duration) +"s ...");
716     sleep(g_traceArgs.duration);
717 
718     OHOS::HiviewDFX::Hitrace::MarkClockSync(g_traceRootPath);
719     StopTrace();
720 
721     if (g_traceArgs.output.size() > 0) {
722         ConsoleLog("capture done, start to read trace.");
723     }
724     DumpTrace();
725     g_traceCollector->Recover();
726     return true;
727 }
728 
HandleRecordingLongBegin()729 static bool HandleRecordingLongBegin()
730 {
731     std::string args = ReloadTraceArgs();
732     if (g_traceArgs.output.size() > 0) {
733         ConsoleLog("warning: The current state does not support specifying the output file path, " +
734                    g_traceArgs.output + " is invalid.");
735     }
736     auto openRet = g_traceCollector->OpenRecording(args);
737     if (openRet.retCode != OHOS::HiviewDFX::UCollect::UcError::SUCCESS) {
738         ConsoleLog("error: OpenRecording failed, errorCode(" + std::to_string(openRet.retCode) +")");
739         if (openRet.retCode == OHOS::HiviewDFX::UCollect::UcError::TRACE_IS_OCCUPIED ||
740             openRet.retCode == OHOS::HiviewDFX::UCollect::UcError::TRACE_CALL_ERROR) {
741             return false;
742         } else {
743             g_traceCollector->Recover();
744             return false;
745         }
746     }
747     ConsoleLog("OpenRecording done.");
748     return true;
749 }
750 
HandleRecordingLongDump()751 static bool HandleRecordingLongDump()
752 {
753     OHOS::HiviewDFX::Hitrace::MarkClockSync(g_traceRootPath);
754     ConsoleLog("start to read trace.");
755     DumpTrace();
756     return true;
757 }
758 
HandleRecordingLongFinish()759 static bool HandleRecordingLongFinish()
760 {
761     OHOS::HiviewDFX::Hitrace::MarkClockSync(g_traceRootPath);
762     StopTrace();
763     ConsoleLog("start to read trace.");
764     DumpTrace();
765     g_traceCollector->Recover();
766     return true;
767 }
768 
HandleRecordingLongFinishNodump()769 static bool HandleRecordingLongFinishNodump()
770 {
771     g_traceCollector->Recover();
772     ConsoleLog("end capture trace.");
773     return true;
774 }
775 
HandleRecordingLongBeginRecord()776 static bool HandleRecordingLongBeginRecord()
777 {
778     std::string args = ReloadTraceArgs();
779     if (g_traceArgs.output.size() > 0) {
780         ConsoleLog("warning: The current state does not support specifying the output file path, " +
781                    g_traceArgs.output + " is invalid.");
782     }
783     auto openRet = g_traceCollector->OpenRecording(args);
784     if (openRet.retCode != OHOS::HiviewDFX::UCollect::UcError::SUCCESS) {
785         ConsoleLog("error: OpenRecording failed, errorCode(" + std::to_string(openRet.retCode) +")");
786         g_traceCollector->Recover();
787         return false;
788     }
789 
790     auto recOnRet = g_traceCollector->RecordingOn();
791     if (recOnRet.retCode != OHOS::HiviewDFX::UCollect::UcError::SUCCESS) {
792         ConsoleLog("error: RecordingOn failed, errorCode(" + std::to_string(recOnRet.retCode) +")");
793         g_traceCollector->Recover();
794         return false;
795     }
796     ConsoleLog("trace capturing. ");
797     return true;
798 }
799 
HandleRecordingLongFinishRecord()800 static bool HandleRecordingLongFinishRecord()
801 {
802     auto recOffRet = g_traceCollector->RecordingOff();
803     if (recOffRet.retCode != OHOS::HiviewDFX::UCollect::UcError::SUCCESS) {
804         ConsoleLog("error: RecordingOff failed, errorCode(" + std::to_string(recOffRet.retCode) +")");
805         g_traceCollector->Recover();
806         return false;
807     }
808     ConsoleLog("capture done, output files:");
809     for (std::string item : recOffRet.data) {
810         std::cout << "    " << item << std::endl;
811     }
812     g_traceCollector->Recover();
813     return true;
814 }
815 
HandleOpenSnapshot()816 static bool HandleOpenSnapshot()
817 {
818     std::vector<std::string> tagGroups = { "scene_performance" };
819     auto openRet = g_traceCollector->OpenSnapshot(tagGroups);
820     if (openRet.retCode != OHOS::HiviewDFX::UCollect::UcError::SUCCESS) {
821         ConsoleLog("error: OpenSnapshot failed, errorCode(" + std::to_string(openRet.retCode) +")");
822         if (openRet.retCode == OHOS::HiviewDFX::UCollect::UcError::TRACE_IS_OCCUPIED ||
823             openRet.retCode == OHOS::HiviewDFX::UCollect::UcError::TRACE_CALL_ERROR) {
824             return false;
825         } else {
826             g_traceCollector->Recover();
827             return false;
828         }
829     }
830     ConsoleLog("OpenSnapshot done.");
831     return true;
832 }
833 
HandleDumpSnapshot()834 static bool HandleDumpSnapshot()
835 {
836     bool isSuccess = true;
837     auto dumpRet = g_traceCollector->DumpSnapshot(OHOS::HiviewDFX::UCollect::TraceCaller::DEVELOP);
838     if (dumpRet.retCode != OHOS::HiviewDFX::UCollect::UcError::SUCCESS) {
839         ConsoleLog("error: DumpSnapshot failed, errorCode(" + std::to_string(dumpRet.retCode) +")");
840         isSuccess = false;
841     } else {
842         ConsoleLog("DumpSnapshot done, output:");
843         for (std::string item : dumpRet.data) {
844             std::cout << "    " << item << std::endl;
845         }
846     }
847     return isSuccess;
848 }
849 
HandleCloseSnapshot()850 static bool HandleCloseSnapshot()
851 {
852     bool isSuccess = true;
853     auto closeRet = g_traceCollector->Close();
854     if (closeRet.retCode != OHOS::HiviewDFX::UCollect::UcError::SUCCESS) {
855         ConsoleLog("error: CloseSnapshot failed, errorCode(" + std::to_string(closeRet.retCode) +")");
856         isSuccess = false;
857     } else {
858         ConsoleLog("CloseSnapshot done.");
859     }
860     return isSuccess;
861 }
862 
InterruptExit(int signo)863 static void InterruptExit(int signo)
864 {
865     /**
866      * trace reset.
867     */
868     g_traceCollector->Recover();
869     _exit(-1);
870 }
871 
main(int argc,char ** argv)872 int main(int argc, char **argv)
873 {
874     if (!IsDeveloperMode()) {
875         ConsoleLog("error: not in developermode, exit");
876         return -1;
877     }
878 
879     if (argc < 0 || argc > 256) { // 256 : max input argument counts
880         ConsoleLog("error: the number of input arguments exceeds the upper limit.");
881         return -1;
882     }
883     bool isSuccess = true;
884     g_traceCollector = OHOS::HiviewDFX::UCollectClient::TraceCollector::Create();
885     if (g_traceCollector == nullptr) {
886         ConsoleLog("error: traceCollector create failed, exit.");
887         return -1;
888     }
889     (void)signal(SIGKILL, InterruptExit);
890     (void)signal(SIGINT, InterruptExit);
891 
892     if (!IsTraceMounted()) {
893         ConsoleLog("error: trace isn't mounted, exit.");
894         return -1;
895     }
896 
897     if (!InitAllSupportTags()) {
898         return -1;
899     }
900 
901     if (!HandleOpt(argc, argv)) {
902         ConsoleLog("error: parsing args failed, exit.");
903         return -1;
904     }
905 
906     if (g_runningState == STATE_NULL) {
907         g_runningState = RECORDING_SHORT_TEXT;
908     }
909 
910     if (g_runningState != RECORDING_SHORT_TEXT && g_runningState != RECORDING_LONG_DUMP &&
911         g_runningState != RECORDING_LONG_FINISH) {
912         ConsoleLog(std::string(argv[0]) + " enter, running_state is " + GetStateInfo(g_runningState));
913     }
914 
915     switch (g_runningState) {
916         case RECORDING_SHORT_RAW:
917             isSuccess = HandleRecordingShortRaw();
918             break;
919         case RECORDING_SHORT_TEXT:
920             isSuccess = HandleRecordingShortText();
921             break;
922         case RECORDING_LONG_BEGIN:
923             isSuccess = HandleRecordingLongBegin();
924             break;
925         case RECORDING_LONG_DUMP:
926             isSuccess = HandleRecordingLongDump();
927             break;
928         case RECORDING_LONG_FINISH:
929             isSuccess = HandleRecordingLongFinish();
930             break;
931         case RECORDING_LONG_FINISH_NODUMP:
932             isSuccess = HandleRecordingLongFinishNodump();
933             break;
934         case RECORDING_LONG_BEGIN_RECORD:
935             isSuccess = HandleRecordingLongBeginRecord();
936             break;
937         case RECORDING_LONG_FINISH_RECORD:
938             isSuccess = HandleRecordingLongFinishRecord();
939             break;
940         case SNAPSHOT_START:
941             isSuccess = HandleOpenSnapshot();
942             break;
943         case SNAPSHOT_DUMP:
944             isSuccess = HandleDumpSnapshot();
945             break;
946         case SNAPSHOT_STOP:
947             isSuccess = HandleCloseSnapshot();
948             break;
949         case SHOW_HELP:
950             ShowHelp(argv[0]);
951             break;
952         case SHOW_LIST_CATEGORY:
953             ShowListCategory();
954             break;
955         default:
956             ShowHelp(argv[0]);
957             isSuccess = false;
958             break;
959     }
960     return isSuccess ? 0 : -1;
961 }
962