courgette/encode_decode_unittest.cc

   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 "courgette/base_test_unittest.h"
   6 #include "courgette/courgette.h"
   7 #include "courgette/streams.h"
   8
   9 class EncodeDecodeTest : public BaseTest {
  10  public:
  11   void TestAssembleToStreamDisassemble(std::string file,
  12                                        size_t expected_encoded_lenth) const;
  13 };
  14
  15 void EncodeDecodeTest::TestAssembleToStreamDisassemble(
  16     std::string file,
  17     size_t expected_encoded_lenth) const {
  18   const void* original_buffer = file.c_str();
  19   size_t original_length = file.length();
  20
  21   courgette::AssemblyProgram* program = NULL;
  22   const courgette::Status parse_status =
  23       courgette::ParseDetectedExecutable(original_buffer,
  24                                          original_length,
  25                                          &program);
  26   EXPECT_EQ(courgette::C_OK, parse_status);
  27
  28   courgette::EncodedProgram* encoded = NULL;
  29
  30   const courgette::Status encode_status = Encode(program, &encoded);
  31   EXPECT_EQ(courgette::C_OK, encode_status);
  32
  33   DeleteAssemblyProgram(program);
  34
  35   courgette::SinkStreamSet sinks;
  36   const courgette::Status write_status = WriteEncodedProgram(encoded, &sinks);
  37   EXPECT_EQ(courgette::C_OK, write_status);
  38
  39   DeleteEncodedProgram(encoded);
  40
  41   courgette::SinkStream sink;
  42   bool can_collect = sinks.CopyTo(&sink);
  43   EXPECT_TRUE(can_collect);
  44
  45   const void* buffer = sink.Buffer();
  46   size_t length = sink.Length();
  47
  48   EXPECT_EQ(expected_encoded_lenth, length);
  49
  50   courgette::SourceStreamSet sources;
  51   bool can_get_source_streams = sources.Init(buffer, length);
  52   EXPECT_TRUE(can_get_source_streams);
  53
  54   courgette::EncodedProgram *encoded2 = NULL;
  55   const courgette::Status read_status = ReadEncodedProgram(&sources, &encoded2);
  56   EXPECT_EQ(courgette::C_OK, read_status);
  57
  58   courgette::SinkStream assembled;
  59   const courgette::Status assemble_status = Assemble(encoded2, &assembled);
  60   EXPECT_EQ(courgette::C_OK, assemble_status);
  61
  62   const void* assembled_buffer = assembled.Buffer();
  63   size_t assembled_length = assembled.Length();
  64
  65   EXPECT_EQ(original_length, assembled_length);
  66   EXPECT_EQ(0, memcmp(original_buffer, assembled_buffer, original_length));
  67
  68   DeleteEncodedProgram(encoded2);
  69 }
  70
  71 TEST_F(EncodeDecodeTest, PE) {
  72   std::string file = FileContents("setup1.exe");
  73   TestAssembleToStreamDisassemble(file, 971850);
  74 }
  75
  76 TEST_F(EncodeDecodeTest, PE64) {
  77   std::string file = FileContents("chrome64_1.exe");
  78   TestAssembleToStreamDisassemble(file, 814709);
  79 }
  80
  81 TEST_F(EncodeDecodeTest, Elf_Small) {
  82   std::string file = FileContents("elf-32-1");
  83   TestAssembleToStreamDisassemble(file, 135988);
  84 }