1// Copyright (c) 2011 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 <string>
6#include <vector>
7
8#include "base/at_exit.h"
9#include "base/basictypes.h"
10#include "base/command_line.h"
11#include "base/files/file_path.h"
12#include "base/files/file_util.h"
13#include "base/logging.h"
14#include "base/strings/string_number_conversions.h"
15#include "base/strings/string_util.h"
16#include "base/strings/utf_string_conversions.h"
17#include "courgette/courgette.h"
18#include "courgette/streams.h"
19#include "courgette/third_party/bsdiff.h"
20
21
22void PrintHelp() {
23  fprintf(stderr,
24    "Usage:\n"
25    "  courgette -supported <executable_file>\n"
26    "  courgette -dis <executable_file> <binary_assembly_file>\n"
27    "  courgette -asm <binary_assembly_file> <executable_file>\n"
28    "  courgette -disadj <executable_file> <reference> <binary_assembly_file>\n"
29    "  courgette -gen <v1> <v2> <patch>\n"
30    "  courgette -apply <v1> <patch> <v2>\n"
31    "\n");
32}
33
34void UsageProblem(const char* message) {
35  fprintf(stderr, "%s", message);
36  fprintf(stderr, "\n");
37  PrintHelp();
38  exit(1);
39}
40
41void Problem(const char* format, ...) {
42  va_list args;
43  va_start(args, format);
44  vfprintf(stderr, format, args);
45  fprintf(stderr, "\n");
46  va_end(args);
47  exit(1);
48}
49
50std::string ReadOrFail(const base::FilePath& file_name, const char* kind) {
51  int64 file_size = 0;
52  if (!base::GetFileSize(file_name, &file_size))
53    Problem("Can't read %s file.", kind);
54  std::string buffer;
55  buffer.reserve(static_cast<size_t>(file_size));
56  if (!base::ReadFileToString(file_name, &buffer))
57    Problem("Can't read %s file.", kind);
58  return buffer;
59}
60
61void WriteSinkToFile(const courgette::SinkStream *sink,
62                     const base::FilePath& output_file) {
63  int count =
64      base::WriteFile(output_file,
65                           reinterpret_cast<const char*>(sink->Buffer()),
66                           static_cast<int>(sink->Length()));
67  if (count == -1)
68    Problem("Can't write output.");
69  if (static_cast<size_t>(count) != sink->Length())
70    Problem("Incomplete write.");
71}
72
73void Disassemble(const base::FilePath& input_file,
74                 const base::FilePath& output_file) {
75  std::string buffer = ReadOrFail(input_file, "input");
76
77  courgette::AssemblyProgram* program = NULL;
78  const courgette::Status parse_status =
79      courgette::ParseDetectedExecutable(buffer.c_str(), buffer.length(),
80                                         &program);
81
82  if (parse_status != courgette::C_OK)
83    Problem("Can't parse input.");
84
85  courgette::EncodedProgram* encoded = NULL;
86  const courgette::Status encode_status = Encode(program, &encoded);
87
88  courgette::DeleteAssemblyProgram(program);
89
90  if (encode_status != courgette::C_OK)
91    Problem("Can't encode program.");
92
93  courgette::SinkStreamSet sinks;
94
95  const courgette::Status write_status =
96    courgette::WriteEncodedProgram(encoded, &sinks);
97  if (write_status != courgette::C_OK)
98    Problem("Can't serialize encoded program.");
99
100  courgette::DeleteEncodedProgram(encoded);
101
102  courgette::SinkStream sink;
103  if (!sinks.CopyTo(&sink))
104    Problem("Can't combine serialized encoded program streams.");
105
106  WriteSinkToFile(&sink, output_file);
107}
108
109bool Supported(const base::FilePath& input_file) {
110  bool result = false;
111
112  std::string buffer = ReadOrFail(input_file, "input");
113
114  courgette::ExecutableType type;
115  size_t detected_length;
116
117  DetectExecutableType(buffer.c_str(), buffer.length(),
118                       &type,
119                       &detected_length);
120
121  // If the detection fails, we just fall back on UNKNOWN
122  std::string format = "Unsupported";
123
124  switch (type)
125  {
126    case courgette::EXE_UNKNOWN:
127      break;
128
129    case courgette::EXE_WIN_32_X86:
130      format = "Windows 32 PE";
131      result = true;
132      break;
133
134    case courgette::EXE_ELF_32_X86:
135      format = "ELF 32 X86";
136      result = true;
137      break;
138
139    case courgette::EXE_ELF_32_ARM:
140      format = "ELF 32 ARM";
141      result = true;
142      break;
143
144    case courgette::EXE_WIN_32_X64:
145      format = "Windows 64 PE";
146      result = true;
147      break;
148  }
149
150  printf("%s Executable\n", format.c_str());
151  return result;
152}
153
154void DisassembleAndAdjust(const base::FilePath& program_file,
155                          const base::FilePath& model_file,
156                          const base::FilePath& output_file) {
157  std::string program_buffer = ReadOrFail(program_file, "program");
158  std::string model_buffer = ReadOrFail(model_file, "reference");
159
160  courgette::AssemblyProgram* program = NULL;
161  const courgette::Status parse_program_status =
162      courgette::ParseDetectedExecutable(program_buffer.c_str(),
163                                         program_buffer.length(),
164                                         &program);
165  if (parse_program_status != courgette::C_OK)
166    Problem("Can't parse program input.");
167
168  courgette::AssemblyProgram* model = NULL;
169  const courgette::Status parse_model_status =
170      courgette::ParseDetectedExecutable(model_buffer.c_str(),
171                                         model_buffer.length(),
172                                         &model);
173  if (parse_model_status != courgette::C_OK)
174    Problem("Can't parse model input.");
175
176  const courgette::Status adjust_status = Adjust(*model, program);
177  if (adjust_status != courgette::C_OK)
178    Problem("Can't adjust program.");
179
180  courgette::EncodedProgram* encoded = NULL;
181  const courgette::Status encode_status = Encode(program, &encoded);
182
183  courgette::DeleteAssemblyProgram(program);
184
185  if (encode_status != courgette::C_OK)
186    Problem("Can't encode program.");
187
188  courgette::SinkStreamSet sinks;
189
190  const courgette::Status write_status =
191    courgette::WriteEncodedProgram(encoded, &sinks);
192  if (write_status != courgette::C_OK)
193    Problem("Can't serialize encoded program.");
194
195  courgette::DeleteEncodedProgram(encoded);
196
197  courgette::SinkStream sink;
198  if (!sinks.CopyTo(&sink))
199    Problem("Can't combine serialized encoded program streams.");
200
201  WriteSinkToFile(&sink, output_file);
202}
203
204// Diffs two executable files, write a set of files for the diff, one file per
205// stream of the EncodedProgram format.  Each file is the bsdiff between the
206// original file's stream and the new file's stream.  This is completely
207// uninteresting to users, but it is handy for seeing how much each which
208// streams are contributing to the final file size.  Adjustment is optional.
209void DisassembleAdjustDiff(const base::FilePath& model_file,
210                           const base::FilePath& program_file,
211                           const base::FilePath& output_file_root,
212                           bool adjust) {
213  std::string model_buffer = ReadOrFail(model_file, "'old'");
214  std::string program_buffer = ReadOrFail(program_file, "'new'");
215
216  courgette::AssemblyProgram* model = NULL;
217  const courgette::Status parse_model_status =
218      courgette::ParseDetectedExecutable(model_buffer.c_str(),
219                                         model_buffer.length(),
220                                         &model);
221  if (parse_model_status != courgette::C_OK)
222    Problem("Can't parse model input.");
223
224  courgette::AssemblyProgram* program = NULL;
225  const courgette::Status parse_program_status =
226      courgette::ParseDetectedExecutable(program_buffer.c_str(),
227                                         program_buffer.length(),
228                                         &program);
229  if (parse_program_status != courgette::C_OK)
230    Problem("Can't parse program input.");
231
232  if (adjust) {
233    const courgette::Status adjust_status = Adjust(*model, program);
234    if (adjust_status != courgette::C_OK)
235      Problem("Can't adjust program.");
236  }
237
238  courgette::EncodedProgram* encoded_program = NULL;
239  const courgette::Status encode_program_status =
240      Encode(program, &encoded_program);
241  courgette::DeleteAssemblyProgram(program);
242  if (encode_program_status != courgette::C_OK)
243    Problem("Can't encode program.");
244
245  courgette::EncodedProgram* encoded_model = NULL;
246  const courgette::Status encode_model_status = Encode(model, &encoded_model);
247  courgette::DeleteAssemblyProgram(model);
248  if (encode_model_status != courgette::C_OK)
249    Problem("Can't encode model.");
250
251  courgette::SinkStreamSet program_sinks;
252  const courgette::Status write_program_status =
253    courgette::WriteEncodedProgram(encoded_program, &program_sinks);
254  if (write_program_status != courgette::C_OK)
255    Problem("Can't serialize encoded program.");
256  courgette::DeleteEncodedProgram(encoded_program);
257
258  courgette::SinkStreamSet model_sinks;
259  const courgette::Status write_model_status =
260    courgette::WriteEncodedProgram(encoded_model, &model_sinks);
261  if (write_model_status != courgette::C_OK)
262    Problem("Can't serialize encoded model.");
263  courgette::DeleteEncodedProgram(encoded_model);
264
265  courgette::SinkStream empty_sink;
266  for (int i = 0;  ; ++i) {
267    courgette::SinkStream* old_stream = model_sinks.stream(i);
268    courgette::SinkStream* new_stream = program_sinks.stream(i);
269    if (old_stream == NULL && new_stream == NULL)
270      break;
271
272    courgette::SourceStream old_source;
273    courgette::SourceStream new_source;
274    old_source.Init(old_stream ? *old_stream : empty_sink);
275    new_source.Init(new_stream ? *new_stream : empty_sink);
276    courgette::SinkStream patch_stream;
277    courgette::BSDiffStatus status =
278        courgette::CreateBinaryPatch(&old_source, &new_source, &patch_stream);
279    if (status != courgette::OK) Problem("-xxx failed.");
280
281    std::string append = std::string("-") + base::IntToString(i);
282
283    WriteSinkToFile(&patch_stream,
284                    output_file_root.InsertBeforeExtensionASCII(append));
285  }
286}
287
288void Assemble(const base::FilePath& input_file,
289              const base::FilePath& output_file) {
290  std::string buffer = ReadOrFail(input_file, "input");
291
292  courgette::SourceStreamSet sources;
293  if (!sources.Init(buffer.c_str(), buffer.length()))
294    Problem("Bad input file.");
295
296  courgette::EncodedProgram* encoded = NULL;
297  const courgette::Status read_status = ReadEncodedProgram(&sources, &encoded);
298  if (read_status != courgette::C_OK)
299    Problem("Bad encoded program.");
300
301  courgette::SinkStream sink;
302
303  const courgette::Status assemble_status = courgette::Assemble(encoded, &sink);
304  if (assemble_status != courgette::C_OK)
305    Problem("Can't assemble.");
306
307  WriteSinkToFile(&sink, output_file);
308}
309
310void GenerateEnsemblePatch(const base::FilePath& old_file,
311                           const base::FilePath& new_file,
312                           const base::FilePath& patch_file) {
313  std::string old_buffer = ReadOrFail(old_file, "'old' input");
314  std::string new_buffer = ReadOrFail(new_file, "'new' input");
315
316  courgette::SourceStream old_stream;
317  courgette::SourceStream new_stream;
318  old_stream.Init(old_buffer);
319  new_stream.Init(new_buffer);
320
321  courgette::SinkStream patch_stream;
322  courgette::Status status =
323      courgette::GenerateEnsemblePatch(&old_stream, &new_stream, &patch_stream);
324
325  if (status != courgette::C_OK) Problem("-gen failed.");
326
327  WriteSinkToFile(&patch_stream, patch_file);
328}
329
330void ApplyEnsemblePatch(const base::FilePath& old_file,
331                        const base::FilePath& patch_file,
332                        const base::FilePath& new_file) {
333  // We do things a little differently here in order to call the same Courgette
334  // entry point as the installer.  That entry point point takes file names and
335  // returns an status code but does not output any diagnostics.
336
337  courgette::Status status =
338      courgette::ApplyEnsemblePatch(old_file.value().c_str(),
339                                    patch_file.value().c_str(),
340                                    new_file.value().c_str());
341
342  if (status == courgette::C_OK)
343    return;
344
345  // Diagnose the error.
346  switch (status) {
347    case courgette::C_BAD_ENSEMBLE_MAGIC:
348      Problem("Not a courgette patch");
349      break;
350
351    case courgette::C_BAD_ENSEMBLE_VERSION:
352      Problem("Wrong version patch");
353      break;
354
355    case courgette::C_BAD_ENSEMBLE_HEADER:
356      Problem("Corrupt patch");
357      break;
358
359    case courgette::C_DISASSEMBLY_FAILED:
360      Problem("Disassembly failed (could be because of memory issues)");
361      break;
362
363    case courgette::C_STREAM_ERROR:
364      Problem("Stream error (likely out of memory or disk space)");
365      break;
366
367    default:
368      break;
369  }
370
371  //  If we failed due to a missing input file, this will
372  // print the message.
373  std::string old_buffer = ReadOrFail(old_file, "'old' input");
374  old_buffer.clear();
375  std::string patch_buffer = ReadOrFail(patch_file, "'patch' input");
376  patch_buffer.clear();
377
378  // Non-input related errors:
379  if (status == courgette::C_WRITE_OPEN_ERROR)
380    Problem("Can't open output");
381  if (status == courgette::C_WRITE_ERROR)
382    Problem("Can't write output");
383
384  Problem("-apply failed.");
385}
386
387void GenerateBSDiffPatch(const base::FilePath& old_file,
388                         const base::FilePath& new_file,
389                         const base::FilePath& patch_file) {
390  std::string old_buffer = ReadOrFail(old_file, "'old' input");
391  std::string new_buffer = ReadOrFail(new_file, "'new' input");
392
393  courgette::SourceStream old_stream;
394  courgette::SourceStream new_stream;
395  old_stream.Init(old_buffer);
396  new_stream.Init(new_buffer);
397
398  courgette::SinkStream patch_stream;
399  courgette::BSDiffStatus status =
400      courgette::CreateBinaryPatch(&old_stream, &new_stream, &patch_stream);
401
402  if (status != courgette::OK) Problem("-genbsdiff failed.");
403
404  WriteSinkToFile(&patch_stream, patch_file);
405}
406
407void ApplyBSDiffPatch(const base::FilePath& old_file,
408                      const base::FilePath& patch_file,
409                      const base::FilePath& new_file) {
410  std::string old_buffer = ReadOrFail(old_file, "'old' input");
411  std::string patch_buffer = ReadOrFail(patch_file, "'patch' input");
412
413  courgette::SourceStream old_stream;
414  courgette::SourceStream patch_stream;
415  old_stream.Init(old_buffer);
416  patch_stream.Init(patch_buffer);
417
418  courgette::SinkStream new_stream;
419  courgette::BSDiffStatus status =
420      courgette::ApplyBinaryPatch(&old_stream, &patch_stream, &new_stream);
421
422  if (status != courgette::OK) Problem("-applybsdiff failed.");
423
424  WriteSinkToFile(&new_stream, new_file);
425}
426
427int main(int argc, const char* argv[]) {
428  base::AtExitManager at_exit_manager;
429  CommandLine::Init(argc, argv);
430  const CommandLine& command_line = *CommandLine::ForCurrentProcess();
431
432  logging::LoggingSettings settings;
433  settings.logging_dest = logging::LOG_TO_ALL;
434  settings.log_file = FILE_PATH_LITERAL("courgette.log");
435  (void)logging::InitLogging(settings);
436  logging::SetMinLogLevel(logging::LOG_VERBOSE);
437
438  bool cmd_sup = command_line.HasSwitch("supported");
439  bool cmd_dis = command_line.HasSwitch("dis");
440  bool cmd_asm = command_line.HasSwitch("asm");
441  bool cmd_disadj = command_line.HasSwitch("disadj");
442  bool cmd_make_patch = command_line.HasSwitch("gen");
443  bool cmd_apply_patch = command_line.HasSwitch("apply");
444  bool cmd_make_bsdiff_patch = command_line.HasSwitch("genbsdiff");
445  bool cmd_apply_bsdiff_patch = command_line.HasSwitch("applybsdiff");
446  bool cmd_spread_1_adjusted = command_line.HasSwitch("gen1a");
447  bool cmd_spread_1_unadjusted = command_line.HasSwitch("gen1u");
448
449  std::vector<base::FilePath> values;
450  const CommandLine::StringVector& args = command_line.GetArgs();
451  for (size_t i = 0; i < args.size(); ++i) {
452    values.push_back(base::FilePath(args[i]));
453  }
454
455  // '-repeat=N' is for debugging.  Running many iterations can reveal leaks and
456  // bugs in cleanup.
457  int repeat_count = 1;
458  std::string repeat_switch = command_line.GetSwitchValueASCII("repeat");
459  if (!repeat_switch.empty())
460    if (!base::StringToInt(repeat_switch, &repeat_count))
461      repeat_count = 1;
462
463  if (cmd_sup + cmd_dis + cmd_asm + cmd_disadj + cmd_make_patch +
464      cmd_apply_patch + cmd_make_bsdiff_patch + cmd_apply_bsdiff_patch +
465      cmd_spread_1_adjusted + cmd_spread_1_unadjusted
466      != 1)
467    UsageProblem(
468        "Must have exactly one of:\n"
469        "  -supported -asm, -dis, -disadj, -gen or -apply, -genbsdiff"
470        " or -applybsdiff.");
471
472  while (repeat_count-- > 0) {
473    if (cmd_sup) {
474      if (values.size() != 1)
475        UsageProblem("-supported <executable_file>");
476      return !Supported(values[0]);
477    } else if (cmd_dis) {
478        if (values.size() != 2)
479          UsageProblem("-dis <executable_file> <courgette_file>");
480        Disassemble(values[0], values[1]);
481    } else if (cmd_asm) {
482      if (values.size() != 2)
483        UsageProblem("-asm <courgette_file_input> <executable_file_output>");
484      Assemble(values[0], values[1]);
485    } else if (cmd_disadj) {
486      if (values.size() != 3)
487        UsageProblem("-disadj <executable_file> <model> <courgette_file>");
488      DisassembleAndAdjust(values[0], values[1], values[2]);
489    } else if (cmd_make_patch) {
490      if (values.size() != 3)
491        UsageProblem("-gen <old_file> <new_file> <patch_file>");
492      GenerateEnsemblePatch(values[0], values[1], values[2]);
493    } else if (cmd_apply_patch) {
494      if (values.size() != 3)
495        UsageProblem("-apply <old_file> <patch_file> <new_file>");
496      ApplyEnsemblePatch(values[0], values[1], values[2]);
497    } else if (cmd_make_bsdiff_patch) {
498      if (values.size() != 3)
499        UsageProblem("-genbsdiff <old_file> <new_file> <patch_file>");
500      GenerateBSDiffPatch(values[0], values[1], values[2]);
501    } else if (cmd_apply_bsdiff_patch) {
502      if (values.size() != 3)
503        UsageProblem("-applybsdiff <old_file> <patch_file> <new_file>");
504      ApplyBSDiffPatch(values[0], values[1], values[2]);
505    } else if (cmd_spread_1_adjusted || cmd_spread_1_unadjusted) {
506      if (values.size() != 3)
507        UsageProblem("-gen1[au] <old_file> <new_file> <patch_files_root>");
508      DisassembleAdjustDiff(values[0], values[1], values[2],
509                            cmd_spread_1_adjusted);
510    } else {
511      UsageProblem("No operation specified");
512    }
513  }
514
515  return 0;
516}
517